Repository navigation
Conversation
…llbacks, speculate with clipping - Save the worker process's Adam state in checkpoints, in the ZeRO-3 CPU offload layout, and give a loaded state to the worker. - Step a sub-group that has a param without a gradient (in step()), give such a param a zero gradient, keep the gradients of earlier micro-steps, and add up a second gradient of a param in the same step. - Do not step during backward on inf/nan gradients; clamp exp_avg_sq at 0 when a step is rolled back. - Pass adam_w_mode, bias_correction and fp32_optimizer_states to the worker, allocate its gradient buffer in the master weights dtype, and step during backward only in the settings adam_rollback supports. - Step during backward also with gradient clipping, and roll back and redo the steps when clipping triggers. - Unscale fp16 gradients in the worker for steps taken during backward. - Partial offload (ZeRO-3 and SuperOffload): the backup optimizer follows the LR schedule, and its state is checkpointed. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
… commit decoupled saves - Send lr, betas, eps and weight_decay with every CPU step, so that momentum and weight decay schedules and the values of a resumed checkpoint reach the worker; a rollback uses those of the step it undoes. - Undo and redo the steps taken during backward when dynamic loss scaling raises the scale in step(). - With the decoupled checkpoint engine, commit a SuperOffload save before save_checkpoint returns: the next backward updates the fp32 weights and Adam state that the writer reads from shared memory. - Warn at initialization when /dev/shm is too small for the tensors shared with the worker. - Keep a nan exp_avg_sq nan in the rollback clamp. - Move the partial offload backup optimizer changes to their own PR. - Tests: two param groups, bf16_optimizer_states, hyperparameter changes, fp16 loss scale increase, CPU steps submitted during backward; checkpoint tests with one and two param groups. 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: 983cba3c21
ℹ️ 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".
| self.speculative_cpu_step = (adamw_mode and fp32_optimizer_states | ||
| and self.master_weights_and_grads_dtype == torch.float32 | ||
| and all(config["bias_correction"] for config in optimizer_configs)) |
There was a problem hiding this comment.
Disable speculative rollback for zero Adam betas
With SuperOffload configured with a valid Adam beta of 0 (for example betas=[0, 0.999]) and a speculative step later being undone because clipping triggers, another gradient arrives, or another group overflows, this gate still enables speculation. adamw_rollback_inplace reconstructs exp_avg/exp_avg_sq by dividing by beta1/beta2, so rollback writes NaNs into the worker state and the next update corrupts the weights; include nonzero betas in the eligibility check (and re-evaluate it if schedules change them).
Useful? React with 👍 / 👎.
| cpu_state = {} | ||
| for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): | ||
| if self.subgroup_to_device[sub_group_id] == 'cpu' and fp32_param in self.optimizer.state: | ||
| cpu_state[sub_group_id] = self.optimizer.state.pop(fp32_param) |
There was a problem hiding this comment.
Preserve FP32 moments when loading BF16-master checkpoints
When bf16_master_weights_and_grads=true and bf16_optimizer_states=false, the worker checkpoint contains FP32 moments, but super().load_state_dict() first loads them through PyTorch against the BF16 flat parameter and casts floating optimizer state to that parameter dtype. Passing the popped state through here therefore leaves a worker configured for FP32 optimizer states holding BF16 moments, silently changing resumed updates; cast exp_avg and exp_avg_sq back to FP32 before set_state or restore the worker state without the generic parameter-dtype cast.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
The rollback recovers the old moments by dividing by beta1 and beta2, so with a beta of 0 (betas=(0, 0.99), for example) undoing a CPU step after clipping, an overflow or a loss scale change wrote nan into the moments and then into the weights. A sub-group now steps during backward only when both betas of its param group are nonzero at that step; otherwise step() runs it after the checks, as for bf16 master weights. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
Summary
SuperOffload has several bugs that change the training result without an error, mostly in the path that runs the CPU
Adam step during backward and undoes it in
step(). This PR fixes them.Fixes
Before → after, against ZeRO-3 CPU offload on GH200 (base
c6291b7e1):exp_avg_sqnegative (nan in the next step).
adam_w_mode,bias_correctionandfp32_optimizer_states, and failed with bf16 masterweights. Up to 1.4e-1 off, or
RuntimeErrorbetas,epsandweight_decay(onlylrwas sent).gradient_clipping> 0 (the default is 1.0), no CPU step overlapped backward. Now it does;step()undoes andredoes those steps when clipping triggers or the fp16 loss scale changes.
save_checkpointnow commits the save first.