[aniravich]

Summary

With CPU optimizer offload (model.offload = "optimizer", the default, or "full"), a resume never restored the optimizer's param groups. Those hold the lr, betas, weight decay and, for Muon, the step counter. SFT and RL both resume through CheckpointManager.load → AppState.load_state_dict, so both were affected.

Problem. In the offload branch, AppState.load_state_dict applies the model weights and writes back the packed optimizer state, but never touches the param groups. The optimizer state is loaded in place: the template that state_dict() hands to dcp_load shares its state tensors with the live optimizer. The param groups are the exception. Optimizer.state_dict() builds new group dicts, and FullCPUOffloadOptimizer.checkpoint_optimizer() copies its groups too. So dcp_load wrote the saved values into those copies, and nothing copied them back. The non-offload branch goes through set_state_dict → Optimizer.load_state_dict, which does restore the groups.

What this broke:

• LinearLR and CosineAnnealingLR compute the next lr from the current group["lr"], and LRScheduler.load_state_dict does not touch the groups. So a resumed run kept the lr its scheduler set when it was built. With a linear or cosine warmup and min_lr = 0, that is lr * 1e-8: 5e-13 instead of 5e-5 until the schedule's next phase starts. For a warmup-only linear schedule, or a cosine schedule resumed after its warmup, that is the rest of the run. Resuming during a decay without warmup restarted the decay near the peak lr instead (cosine: 5e-5 instead of 1e-5). With the default constant schedule the lr symptom does not show as long as the config's lr equals the checkpoint's.
• With a warmup-only linear schedule, SFT's ckpt.skip_scheduler rebuilt its warmup on top of that lr, so it topped out at 5e-13.
• dion's Muon keeps its step counter in the param groups. It restarted at 1, so the bias correction of its AdamW sub-optimizer (embeddings, norms, lm_head, ...) made the first updates after a resume too large: 1.66x, 1.39x, 1.27x, ... of the uninterrupted run in a small run resumed at step 60. This happens with the constant schedule too.

Fix. After dcp_load, copy every loaded param group entry into the matching live group, except "params" and Muon's "distributed_mesh_name". Groups are matched by position, the same way DCP fills the template. The live groups come from self.optimizers, not _get_checkpoint_optimizers(), because full offload's checkpoint optimizer only holds copies. The groups are updated in place, so FullCPUOffloadOptimizer's per-chunk references to them stay valid. The non-offload branch is unchanged.

distributed_mesh_name names the mesh Muon uses for the expert group: "dp_shard_mod_ep" when expert parallelism is on, None when it is off. It describes the current run's layout, not training progress, and with model.ep = "auto" it follows the GPU count. So the resumed run keeps its own value. Restoring the saved one would make the first step after resuming an EP checkpoint without EP raise a KeyError, where the offload path ran before. offload = "none" restores whole groups through Optimizer.load_state_dict, so it still fails on such a resume, as it does on main.

Behaviour change. With offload, a resume now keeps the checkpoint's lr, betas, weight decay and other group hyperparameters, as offload = "none" already does. Before, the config's values applied because of this bug. For example, resuming or forking a run (--resume.dir) with a lower optim.lr used the new lr on main with the default offload, and now uses the checkpoint's. ckpt.skip_scheduler does not change that, because the rebuilt schedule starts from the restored lr and initial_lr. Only ckpt.skip_optimizer, which skips the whole optimizer state, keeps the config's values. Whether the config should be able to override these on resume is a separate question that affects both paths.

Checkpoints written after a resume on an affected version already hold the wrong lr in their param groups (5e-13 after a warmup). This change restores whatever the checkpoint holds.

Added test_checkpoint_resume_restores_param_groups to tests/unit/train/test_state_offload.py, which is GPU-marked like the rest of that file. It trains a CPUOffloadOptimizer run 3 steps into a 4-step LinearLR warmup, then saves and loads it through AppState and DCP. It checks that the lr and the next 2 steps match the uninterrupted run. Full offload is not covered by the test; it is covered by the resume script below and still needs a GPU check.

Validation

Rebased on main (7a630865c, includes #4053); the results below are on the rebased tree, compared with main at the same commit.

No GPU was available. Everything below ran on CPU with torch 2.14.1 (prime-rl pins 2.13.0). dion and torchdata were not installed: dion was imported from its source at the pinned commit (6f9242d), with its Triton module stubbed out (Muon runs with use_triton=False). torchdata.stateful_dataloader.StatefulDataLoader was replaced by a stub, which ckpt.py only uses in type annotations. Other packages that are not installed here (triton, fla, wandb, loguru, ...) were replaced by mock modules. The checkpoint and optimizer code checked below does not use any of them.

• New test, with CUDA emulated on CPU. .cuda() was a no-op, and CPUOffloadOptimizer's H2D copy was a clone(). Its pinned buffers were unpinned, and torch.cuda.synchronize() was a no-op. AdamW, LinearLR, AppState and DCP were the real code.

  • The test passes on this branch. On main it fails with assert 1e-10 == 0.007500000025.
  • The other tests in the file give the same results on both: 2 pass, and the 2 that assert pinned memory fail under the emulation.

• Resume script. It used the real AppState, setup_scheduler and DCP save/load, on a single-rank gloo model with DTensor parameters. It ran three wrappers:

  • CPUOffloadOptimizer with AdamW;
  • CPUOffloadOptimizer with dion Muon;
  • FullCPUOffloadOptimizer. Its CUDA-only setup and H2D copy were skipped, and steps ran torch's AdamW on the CPU masters.

  Each run was saved at step 60 and resumed, and the lr for steps 61-65 was compared with the uninterrupted run. Three schedules were tested: linear warmup 50; linear warmup 50 + decay 50; and cosine warmup 50.

  • main: offload = "none" matches. Both "optimizer" wrappers and "full" resume at 5e-13 instead of 5e-5 (cosine: instead of 4.95e-5). Muon's group step reads 6 instead of 65 five steps after the resume.
  • This branch: all 12 cases match, including the 9 with offload. The param groups equal the saved ones right after the load, and the parameters are bitwise equal 5 steps after the resume.
  • Applying the loaded groups to _get_checkpoint_optimizers() instead fixes "optimizer" but not "full".

• Longer runs (CPUOffloadOptimizer, 200 steps). With linear warmup (with and without decay), cosine warmup resumed at step 30 and 60, and cosine or linear decay without warmup resumed at step 140, the lr on this branch matches the uninterrupted run up to step 199. On main the warmup cases recover only where a next phase starts (step 51 for cosine resumed at 30, step 151 for linear warmup + decay) and never for the other two, and the decay-only schedules resume at 5e-5 instead of 1.03e-5 (cosine) and 2.98e-5 (linear).

• ckpt.skip_scheduler flow, with the schedule rebuilt after the load: on main with offload, a warmup-only linear schedule tops out at 5e-13; warmup + decay and cosine reach 5e-5 because SequentialLR resets the lr to initial_lr. On this branch every case reaches 5e-5, with and without offload.

• Multi-rank. A 2-rank FSDP2 run (gloo) with CPUOffloadOptimizer and a linear warmup, saved after 4 steps and resumed at world size 1, 2 and 4: on this branch the groups and lr (0.004) are restored, and the next 3 steps use the same lr as the uninterrupted run. The parameters after them are equal at world size 2 and within 1.5e-8 at 1 and 4, where the gradient reduction differs. On main the lr after the load is 1e-10.

• Muon expert group across an EP change (single rank, mesh dims named as with EP on or off). With CPUOffloadOptimizer, the resumed expert group keeps its own distributed_mesh_name in both directions, its lr and step are restored, and the next step runs. offload = "none" raises the KeyError after EP on → off, on both main and this branch.

• Changed config on resume (constant schedule; saved with lr 1e-2 and weight decay 0.01, resumed with 1e-3 and 0.1). main: "none" keeps 1e-2 / 0.01, the default offload takes 1e-3 / 0.1, with or without ckpt.skip_scheduler. This branch: both modes keep 1e-2 / 0.01, with or without ckpt.skip_scheduler. With ckpt.skip_optimizer, both modes take 1e-3 / 0.1 on both trees.

• CPU CI. test_state_offload.py is GPU-marked, so pytest tests/unit -m "not gpu" only imports it. Its new top-level from prime_rl.trainer.ckpt import AppState needs torchdata, which CI installs with uv sync --all-extras.

• Lint. ruff check and ruff format --check (ruff 0.15.14, repo config) and py_compile are clean on the changed files.

Still needs a GPU:

• The new test on CUDA, in the GPU unit-test job. That job skips draft PRs, so it runs once this one is marked ready for review.
• A real multi-GPU resume with offload = "full" and a warmup schedule. The native CPU AdamW kernel reads group["lr"] through the per-chunk group references. Check that the logged lr right after the resume continues from the run before it.
• The same check for Muon with offload = "optimizer" under FSDP, including an MoE model resumed with expert parallelism switched on or off.

🤖 Generated with Claude Code