feat: chunked per-layer optimizer step with stream-overlapped CPU offload

[samsja]

Summary

The CPUOffloadOptimizer now optionally performs optimizer.step() per-transformer-layer ("chunked") instead of all-at-once, with optional stream-overlapped H2D/D2H transfers. Each layer's optimizer states are moved to GPU, the step runs for that layer only, and states move back to CPU before the next layer.

Problem

With the old all-at-once offload, the entire model's optimizer states land on GPU simultaneously during step(). When weight + grad + all_opt_states > VRAM — even without activations — this causes OOM on large models.

Config

Two new config options (both default false — opt-in):

[trainer.model]
optim_cpu_offload_chunked = true   # per-layer step vs all-at-once
optim_cpu_offload_stream = true    # stream overlap vs sequential move->step->move

• optim_cpu_offload_chunked = false: all-at-once offloading (loads full model states to GPU in one shot)
• optim_cpu_offload_stream = false: chunked step without CUDA stream overlap — simple sequential loop, no stream objects in the code path

Implementation

• Per-layer chunking: Parameters grouped by transformer layer (regex layers.{N}.). Non-layer params (embeddings, lm_head, norm) go into a trailing "misc" chunk.
• Stream-overlapped (_step_chunked_streamed): H2D/D2H on dedicated CUDA streams; next layer prefetched while current computes; previous layer evicted while next computes. Ends with torch.cuda.synchronize() to block CPU.
• No-stream (_step_chunked): Simple sequential move_to_gpu -> step -> move_to_cpu per chunk. No CUDA stream objects.
• First-step chunked: The init step also uses the chunked path, avoiding materializing all optimizer states on GPU at once during initialization.
• Pinned memory: D2H pre-allocates pinned CPU destination and async-copies into it. state_dict() and ckpt paths sync after D2H to prevent reading incomplete copies.
• Muon step counter sync: Per-group step integer synced after all chunks since param_groups are temporarily swapped.

VRAM impact (synthetic benchmark, float32, AdamW)

Model (layers, hidden)	Opt states	All-at-once overhead	Chunked overhead	Savings
Small (14, 768)	1,383 MiB	2,076 MiB (3.4 GiB)	564 MiB (1.9 GiB)	1.5 GiB
Medium (28, 1024)	4,084 MiB	6,132 MiB (10.2 GiB)	756 MiB (4.8 GiB)	5.4 GiB
Large (40, 2560)	33,250 MiB	50,350 MiB (83.9 GiB)	2,462 MiB (36.0 GiB)	47.9 GiB

Validation

• Full reverse-text RL training (20 steps, Qwen3-0.6B, 2 GPUs) completed successfully
• Unit tests for chunked/stream, chunked/no-stream, non-chunked, first-step chunked, state_dict sync, checkpoint compatibility all pass
• base_optimizer, state, param_groups, zero_grad, state_dict, load_state_dict APIs unchanged

Bugbot issues addressed

1. Async D2H checkpoint corruption (HIGH): Added torch.cuda.synchronize() after _move_states("cpu") in state_dict() and ckpt.py
2. Incomplete stream dependency (MEDIUM): Added torch.cuda.synchronize() at end of streamed step; h2d_stream.wait_stream(d2h_stream) before prefetch to prevent pin buffer races
3. First step OOM (HIGH): Init step now uses chunked path instead of full optimizer.step()
4. Defaults disagree (MEDIUM): Aligned setup_optimizer defaults with config (false)

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 053beda. Configure here.}