Repository navigation
Conversation
… sub-groups and checkpoint their Adam state With offload_optimizer.ratio < 1, the sub-groups kept on the device are stepped by a separate torch optimizer. - Give each of its param groups its own dict and copy lr, betas, eps and weight_decay from self.optimizer's group before each device step (ZeRO-3 and SuperOffload). LR schedules ignored the first group, and the shared dicts of the other groups kept every CPU partition after a checkpoint load, so device steps also updated CPU sub-groups. - Save the device sub-groups' Adam state with the other sub-groups' in optimizer_state_dict['state'], and move it to the backup optimizer on load and on universal load, so that checkpoints do not depend on the ratio and ds_to_universal finds every sub-group. - Read the state of a device sub-group from the backup optimizer in safe_get/set_full_optimizer_state. - Use torch Adam (L2) when adam_w_mode is off, and reject bias_correction=False. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
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: afded1772a
ℹ️ 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 len(self.optimizer.param_groups) > 1: | ||
| for i in range(1, len(self.optimizer.param_groups)): | ||
| self.backup_optimizer.add_param_group(self.optimizer.param_groups[i]) | ||
| if not all(group.get("bias_correction", True) for group in self.optimizer.param_groups): |
There was a problem hiding this comment.
Revalidate bias correction after loading optimizer groups
When a checkpoint whose saved optimizer groups have bias_correction=False is resumed into a ratio < 1 engine constructed with the current/default value True, this check passes, but Optimizer.load_state_dict (and the universal-checkpoint loop) later restores False only to self.optimizer; _sync_backup_optimizer_group never propagates or rejects it. CPU subgroups then run without bias correction while device subgroups use torch Adam's correction, silently changing the resumed training result instead of raising the promised ValueError, so the loaded groups need to be revalidated.
Useful? React with 👍 / 👎.
| "bf16": { | ||
| "enabled": True | ||
| }, |
There was a problem hiding this comment.
Guard the BF16-only test configuration
When this test file runs on a supported accelerator without BF16 support (for example a CUDA V100), nearly every _config call reaches DeepSpeedEngine._do_sanity_check, which raises ValueError("Type bf16 is not supported on your device.") before any partial-offload behavior is exercised. The test should select a supported precision or skip explicitly via get_accelerator().is_bf16_supported() so the new regression suite remains runnable across supported hardware.
Useful? React with 👍 / 👎.
…ram dtype
torch Adam keeps the moments of the device sub-groups in the param dtype and fails on fp32 moments for bf16 params
("expected dtype float for `end`"). DeepSpeedCPUAdam keeps fp32 moments for bf16 params (fp32_optimizer_states), so a
load that keeps them in fp32 would pass fp32 moments to the backup optimizer. Cast them to the param dtype when they
move there. Test: resume with bf16 master weights at ratio 0.5.
Signed-off-by: Xinyu Lian <lian7@illinois.edu>
Summary
With
offload_optimizer.ratio< 1, ZeRO-3 keeps some sub-groups on the device and steps them with a separate torchoptimizer (
backup_optimizer). That optimizer has several bugs that change the training result without an error.SuperOffload uses the same optimizer for its device sub-groups; the rest of SuperOffload is fixed in #8972.
Fixes
set_lr(), and changes ofbetas,epsandweight_decay. Inafirst step at lr 0, the weights moved by 3.0e-3. The backup optimizer now copies these values before each step.
with the state of the CPU sub-groups: a resume is bit-identical, and a checkpoint loads with another
ratio.dicts were shared with the main optimizer (3.0e-3 extra per step at lr 3e-3). Each param group of the backup
optimizer now has its own dict.
safe_get_full_optimizer_stateandsafe_set_full_optimizer_stateraisedKeyErrorfor the device sub-groups, andds_to_universalraisedKeyError: 0. Both work now.adam_w_mode: false, the device sub-groups used decoupled weight decay (2.4e-2 off after 4 steps); they now usetorch.optim.Adam(L2).bias_correction: falsewithratio< 1 now raises aValueError.