Warn when DeepSpeed ZeRO silently changes custom optimizer semantics

Open
  • Is something wrong?
#4,392 2 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

  • #4414 by @Eryx-Z — closed without merging

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
48/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
backend, machine-learning

Research direction

Start by tracing the DeepSpeed integration path that sets zero_allow_untested_optimizer, then inspect existing optimizer integration tests. Reproduce the reported ZeRO-1 behavior with the linked deterministic example and compare it with DDP. Done means the risk is made explicit through an opt-in or warning, with tests covering the behavior; the issue also proposes documentation and a ZeroRedundancyOptimizer recipe.

Written by the indexing model from the issue text.

Description

System Info
- Accelerate version: 1.15.0
- Platform: Linux 5.4.0 x86_64, glibc 2.31
- Python version: 3.12.13
- NumPy version: 2.5.3
- PyTorch version: 2.14.1+cu132
- PyTorch accelerator: CUDA
- DeepSpeed version: 0.19.7
- ProdigyPlusScheduleFree version: 2.0.1
- GPU: 2 x NVIDIA RTX A6000
- Accelerate default config: not used; all settings explicit
Information
  • The official example scripts
  • My own modified scripts
Tasks
  • One of the scripts in the examples/ folder
  • A deterministic optimizer-semantics test using a custom optimizer with DeepSpeed ZeRO-1
Reproduction

This is a follow-up to #4159, whose wrapper traversal is fixed in 1.15.0 and passed our test.
The remaining problem is independent: Accelerate accepts an arbitrary torch.optim.Optimizer
with DeepSpeed, automatically enables zero_allow_untested_optimizer, and the job can exit 0 even
when ZeRO partitioning changes the optimizer's mathematics.

With ProdigyPlusScheduleFree, DeepSpeed replaces the optimizer group with a local flattened FP32
master partition. Prodigy's group-level d estimate becomes rank-local, and original tensor
boundaries used by factorization, estimator slicing, and RMS scaling disappear.

Standalone two-rank reproduction: https://gist.github.com/XIONGPEILIN/c2d12d59fecd7ef7e347667e551e9fdd

Observed after 64 deterministic updates:

execution final rank d max model abs diff vs DDP
replicated DDP 0.016384, 0.016384 reference
DeepSpeed ZeRO-1 0.00172216, 0.000215269 0.175511
ZeRO-1 + globally summed d statistics 0.00180069, 0.00180069 0.161856

All runs completed without an exception. The 1.15.0 optimizer.eval() / optimizer.train()
forwarding fix worked on both ranks.

We also tested a solution that keeps the stock optimizer recipe but uses tensor-wise PyTorch
ZeroRedundancyOptimizer: max model difference from DDP was 7.152557e-07 at a predeclared
1e-6 tolerance, with identical d on both ranks.

Expected behavior

I do not expect Accelerate to prove that every custom optimizer is partition-safe. I do expect the
integration to make this risk explicit rather than silently opting arbitrary optimizers into
zero_allow_untested_optimizer.

Possible improvements:

  1. require an explicit opt-in for untested custom optimizers under ZeRO, or emit a prominent
    Accelerate warning explaining that ZeRO may replace original tensors with local flat masters;
  2. document that optimizers using group-global statistics or tensor shapes/boundaries need a
    partition-aware implementation;
  3. add a documented Accelerate recipe for PyTorch ZeroRedundancyOptimizer that avoids optimizer
    device-placement state_dict() calls and supports Schedule-Free train() / eval() forwarding.

This would complement, rather than change, the fix already released for #4159.

Dominant language
Python
Stars
9.9k
Forks
1.6k
Avg merge
20h 23m
Merged PRs (30d)
12

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from huggingface/accelerate

All issues in huggingface/accelerate

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.