[faresobeid]

Summary

utils.py had five wrappers (prepare_gradient_offload, begin_backward, finish_backward, scale_gradients_, clip_grad_norm_) whose only job was to branch on gradient_manager is None. SFT also had two scaling paths: one with offload and one without.

This PR:

• adds GPUGradientManager (optim/base.py). It has the same begin_step / begin_backward / finish_backward / clip_grad_norm_ / close interface as GradientOffloadManager. It applies the step's gradient scale after the final backward of the step.
• makes setup_optimizer always return a manager. model is now a required argument.
• deletes the five wrappers. RL and SFT call the manager directly.
• makes SFT count the global loss tokens before the step in both modes, as the offload path and RL already do. This leaves one scaling path.
• removes GradientOffloadManager.scale_ and the wait_for_copies argument of finish_backward. Nothing used them: the wrapper always passed False.

Behavior change

None. One side effect: SFT without offload now keeps the grad_accum CPU micro-batches of a step in memory. The offload path already did this.

What I ran

• pytest tests/unit/orchestrator tests/unit/test_configs* -m "not gpu": 291 passed
• pytest tests/unit/train/rl tests/unit/train/sft tests/unit/train/test_state_offload.py tests/unit/train/test_cpu_adam.py tests/unit/train/test_perf.py tests/unit/train/test_world.py -m "not gpu": 93 passed, 1 failed. The failure is test_native_cpu_adamw_matches_fused_torch_and_preserves_gradients, which needs ninja to build the C++ extension. It fails the same way on main in this environment.
• ruff check and ruff format --check on the changed files
• CPU differential check (throwaway script, not in the repo). Tiny model, 4 micro-batches, single-rank gloo. I ran it with plain tensors and with FSDP2 DTensor gradients, using max_norm None / 0.05 / 100, for both the SFT and the RL scaling flow. Old wrappers with None (main) vs GPUGradientManager (branch): the gradients after scale and clip and the returned grad norm are bit-identical in all 12 cases. I also checked that the SFT up-front count (sum(full loss_mask) over dp_cp // cp) equals the old streamed count (sum of CP-shard counts) for cp in {1, 2, 4}.

GPU checks to do

• SFT, no offload, with max_norm (single node, with and without CP): loss and grad_norm match main for the first steps.
• RL, no offload, with max_norm: loss and grad_norm match main.
• SFT and RL with model.full_offload: still run (the manager path is unchanged apart from the removed unused methods).
• SFT with optim_cpu_offload: grad_norm matches main.

Conflicts

git merge-tree merges cleanly with all of these. None of them call the removed wrappers.

• chore: remove Kubernetes deployment path (Helm chart, trainer metrics server) #3767: rl/train.py. Nearby hunks: imports and the end of train().
• feat: exact score centering over the sampling mask (arXiv:2609.20807) #3794: rl/train.py. Nearby hunks: imports and the compute_loss call just before the backward block.
• chore: remove legacy config options #3770: optim/__init__.py (imports) and sft/train.py (different region).
• chore: remove legacy dead code #3769, feat: Total Router Recall (replay sampler MoE routing weights) #3793, feat(trainer): Composer 2-style plausibility filter for router replay #3796, feat: by-handle payload segments: reader, packing, CP-window reads, GC #3819: rl/train.py, different regions.

---

Note

Medium Risk
Touches core training-step gradient scaling, clipping, and FSDP normalization in RL/SFT; intended behavior-preserving but affects every optimizer step and distributed grad-norm logic.

Overview
Introduces GPUGradientManager so GPU-resident training uses the same step/backward/clip/close API as GradientOffloadManager. setup_optimizer now always returns a manager (GPU path gets GPUGradientManager; full CPU offload still returns the offload manager) and model is required.

Removes the trainer/utils.py helpers that only branched on gradient_manager is None. RL and SFT call gradient_manager.begin_step / begin_backward / finish_backward / clip_grad_norm_ directly; post-step FSDP undo scaling for RL moves into begin_step instead of a separate scale_gradients_ pass.

SFT drops the dual no-offload vs offload scaling paths: every step prefetches grad-accum micro-batches and computes the global token denominator up front (aligned with RL/offload). GradientOffloadManager loses unused scale_ and the wait_for_copies flag on finish_backward.

^{Reviewed by Cursor Bugbot for commit 216a8f5. Bugbot is set up for automated code reviews on this repo. Configure here.}