Repository navigation
Fix bug in Adagrad optimizer initialization - #8976
hwchen2017 wants to merge 2 commits into
Conversation
Signed-off-by: Hongwei Chen <hongweichenmail@gmail.com>
Signed-off-by: Hongwei Chen <hongweichenmail@gmail.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5c86f15a2c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if is_adagrad: | ||
| self.optimizer = torch.optim.Adagrad(self.fp32_partitioned_groups_flat, **self.optimizer.defaults) | ||
| self._set_fp32_optimizer_param_groups() | ||
| self.optimizer = torch.optim.Adagrad(self.optimizer.param_groups, **self.optimizer.defaults) |
There was a problem hiding this comment.
Preserve Adagrad groups in ZenFlow overlap paths
When ZenFlow overlap stepping is enabled, this corrected method is bypassed: deepspeed/runtime/zenflow/engine_stage3.py::initialize_optimizer_states still rebuilds Adagrad from the bare fp32_partitioned_groups_flat list, and ZenFlowZeroOptimizerParallel.initialize_optimizer_states does the same for ZeRO 1/2. Consequently, the same multi-group IndexError and loss of per-group hyperparameters remain for zenflow.overlap_step=true; apply the parameter-group reconstruction to those sibling paths and cover them in the test.
AGENTS.md reference: AGENTS.md:L17-L19
Useful? React with 👍 / 👎.
0z5a
left a comment
There was a problem hiding this comment.
Read at 5c86f15a.
The fix looks right to me. On CPU (gloo, DP=2) your four tests pass on this branch, and the same file run against master fails all four: the two multi-group cases with IndexError, the two single-group cases on values, because the group's lr/weight_decay were replaced by the constructor defaults.
This also fixes a client LR scheduler, which isn't covered yet. Passing self.optimizer.param_groups to the new Adagrad reuses the same group dicts, because add_param_group keeps the dict it is given. So a scheduler built on the client optimizer before deepspeed.initialize(optimizer=..., lr_scheduler=...) now drives the optimizer ZeRO actually steps. I checked this with a client Adagrad(lr=0.1) plus StepLR(step_size=1, gamma=0.1) against a plain torch reference for 3 steps with clipping off:
- master, ZeRO-1/2/3: the engine's lr stays at 0.1 while the reference goes 0.01 → 1e-4, and the weights are 0.063 apart after 3 steps.
- this branch, ZeRO-1/2/3: lr matches and the weights are bitwise equal at every step.
A scheduler case in test_zero_adagrad.py would pin this, since it depends on the dicts being shared.
I agree with Codex on the ZenFlow paths: zenflow/engine_stage3.py:429 and zenflow/zenflow_stage_1_and_2.py:651 still rebuild Adagrad from the flat partition list. I only read those; I didn't run ZenFlow.
| # which do lazy initialization of the state at the first call to step. | ||
| if isinstance(self.optimizer, torch.optim.Adagrad): | ||
| self.optimizer = torch.optim.Adagrad(self.single_partition_of_fp32_groups, **self.optimizer.defaults) | ||
| self.optimizer = torch.optim.Adagrad(self.optimizer.param_groups, **self.optimizer.defaults) |
There was a problem hiding this comment.
Reusing self.optimizer.param_groups here also keeps the client's group dicts, so an LR scheduler built on the client optimizer now takes effect (on master it is silently ignored for Adagrad). Worth a test, since nothing else guarantees the dicts stay shared.
|
Looks like the zenflow path will never get triggered. • |
|
You're right, I read it the same way now: ZenFlow steps with |
ZeRO rebuilt Adagrad from a flat list of FP32 partitions instead of parameter-group dictionaries. It combined every partition into one optimizer group, using the constructor defaults.
This caused two problems:
• Multi-group crashes: ZeRO still referenced the original group IDs, but the rebuilt optimizer only had group 0, causing
IndexErrorduring gradient processing or optimizer stepping.• Lost per-group settings: custom learning rates, weight decay, and other options were replaced by global defaults. Even a single-group configuration could silently produce incorrect updates if it used group-specific overrides.
With this fix, now it use parameter-group dictionaries with the custom learning rates, weight decay and etc. In zero 3, the param group were created temporarily for initialization and then were cleared.