[faresobeid]

Overlaps the optimizer-state CPU offload copies with the AdamW step (model.offload = "optimizer", CPUOffloadOptimizer).

On main, the step moves all AdamW states to the GPU, steps, and moves them back, one direction at a time on the compute stream, and the GPU holds every state at once.

Change

• AdamW steps chunk by chunk (whole parameters, up to 1 GiB per chunk) through two GPU staging buffers. While chunk i updates on the compute stream, chunk i+1 loads on one side stream and chunk i-1 writes back to pinned CPU memory on another, so both PCIe directions run at once.
• Streams are created once. Staging buffers are allocated on the compute stream, so the caching allocator reuses the same blocks every step and they are free for activations between steps.
• AdamW's scalar step counters stay on CPU, as in plain torch AdamW, so there is no per-parameter device sync.
• Unchanged path: Muon, the first step, any step where a parameter has a grad but no state yet, and checkpoint save/load. CPUOffloadOptimizer no longer takes pin_memory (its only caller used True).

Numbers (H200)

	main	this PR
Optimizer step, 32 GB AdamW states, 1 GPU	1.37 s	0.79 s
Same, all 8 GPUs at once (host-bandwidth bound)	2.5 s	2.25 s
Peak allocated during the step	80.0 GiB	34.5 GiB
GLM-5.3 step (1 node, 8 layers, 131k, cp4 ep8, fp8)	11.7 s	11.1 s

DeepSeek-V4-Flash (1 node, 8 layers, 65k, cp2 ep8): 17.0-18.1 s -> 16.6-17.7 s. Entropy and mismatch KL identical per step.

Validation

• Parameters and optimizer states bitwise equal to plain AdamW and to main's offload path after 6 steps (fp32 and bf16, 6 chunks; 1 GPU and FSDP2 on 2 GPUs).
• tests/unit/train/test_state_offload.py (includes a parameter that gets its first grad on step 2). The parity, pinned-buffer and round-trip tests set CHUNK_BYTES = 0, so each parameter is its own chunk and they cross chunk boundaries and reuse both staging slots; all pass on 1 H200.
• pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train, ruff check, ruff format --check, uv lock --check on the branch merged with current main.

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes the hot-path AdamW training step with multi-stream memory management and param-group mutation during chunked steps; correctness is covered by GPU tests but regressions could affect peak memory or numerical parity under FSDP/DTensor.

Overview
AdamW optimizer steps through CPUOffloadOptimizer now use a chunked, stream-overlapped path instead of moving all state to GPU, stepping, then moving everything back on the compute stream.

Whole parameters are grouped into ~1 GiB chunks with aligned views in two GPU staging buffers. While chunk i runs on the default stream, the next chunk H2D loads on one side stream and the previous chunk D2H writes to pinned CPU on another, shrinking peak GPU state during the step and overlapping PCIe traffic with compute. AdamW step scalars stay on CPU; DTensor handling is centralized in _local / _with_local. The legacy full-move path remains for the first step, Muon/non-AdamW, and any step where a parameter has a grad but no optimizer state yet. The wrapper drops the pin_memory constructor flag (pinned CPU buffers are always used).

Unit tests force one-parameter-per-chunk via CHUNK_BYTES=0 and add coverage for parameters that only receive gradients after the first step.

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

[faresobeid]

Measured on GPU: Qwen3-30B-A3B on 8×H200, FSDP 8 / EP 8, seq 16384, full AC, compile, 25 steps, mean tok/s over the last 10. Losses agree across arms (≈1.30–1.31 at step 25).

config	main	this PR	gain
no offload (reference)	57.7k (n=4)	–	–
optimizer offload only	27.3k (n=3), 63.2 GiB	32.8k (n=2)	+20%
both offloads, with #4238	25.9k (n=3)	32.4k (n=2)	+25%

Even with this PR, optimizer offload is still ~44% below no-offload (step ≈ 4.9–5.5 s vs 2.26 s). A follow-up investigation of where the rest goes (copies, CPU step, NUMA/PCIe bandwidth, overlap) is underway. Any improvement will go on a separate branch on top of this PR.

[garrett361]

Looks ok to me, but maybe wait on Sami as there's apparently a history here that I don't know about.

Inline comments are mostly agent-written. Generated below, for your agent:

---

Approve with comments. With model.offload = "optimizer" and AdamW, every step after the first now streams states through two GPU staging buffers chunk by chunk (H2D, update, D2H on three streams) and keeps the step counters on CPU, instead of moving every state to the GPU and back on the compute stream.

1. The multi-chunk path has no test coverage: every test model fits in one chunk. CHUNK_BYTES is 1 GiB and the test model is a few hundred KiB, so CI never runs chunk boundaries, staging-buffer offsets, or the reuse of a buffer slot by chunk i+2.

   prime-rl/src/prime_rl/trainer/optim/state_offload.py

   Lines 149 to 176 in 1a8052f

   	def load(i: int) -> torch.cuda.Event:
   	buffer, slot = buffers[i % 2], cpu_values[i % 2]
   	with torch.cuda.stream(h2d):
   	if written_back[i % 2] is not None:
   	h2d.wait_event(written_back[i % 2])
   	slot.clear()
   	for param, key, offset in chunks[i][1]:
   	cpu_value = state[param][key]
   	cpu_local = _local(cpu_value)
   	nbytes = cpu_local.numel() * cpu_local.element_size()
   	gpu_local = buffer[offset : offset + nbytes].view(cpu_local.dtype).view(cpu_local.shape)
   	gpu_local.copy_(cpu_local, non_blocking=True)
   	slot[param, key] = cpu_value
   	state[param][key] = _with_local(cpu_value, gpu_local)
   	return h2d.record_event()
   	loaded = load(0)
   	for i, (params, _) in enumerate(chunks):
   	next_loaded = load(i + 1) if i + 1 < len(chunks) else None
   	compute.wait_event(loaded)
   	self._step_params(params)
   	d2h.wait_stream(compute)
   	with torch.cuda.stream(d2h):
   	for (param, key), cpu_value in cpu_values[i % 2].items():
   	_local(cpu_value).copy_(_local(state[param][key]), non_blocking=True)
   	state[param][key] = cpu_value
   	written_back[i % 2] = d2h.record_event()
   	loaded = next_loaded

   Fix: monkeypatch state_offload.CHUNK_BYTES to 0 in the trained_pair fixture, which puts each of the model's four parameters in its own chunk and runs the existing bitwise parity, pinned-buffer and round-trip tests through that path. No new test is needed.

   prime-rl/tests/unit/train/test_state_offload.py

   Lines 69 to 79 in 1a8052f

   	@pytest.fixture
   	def trained_pair() -> TrainedPair:
   	reference_model = make_model(0)
   	reference_optimizer = torch.optim.AdamW(reference_model.parameters(), lr=1e-2)
   	run_steps(reference_model, reference_optimizer, num_steps=3, seed=1)
   	offloaded_model = make_model(0)
   	offloaded_optimizer = CPUOffloadOptimizer(torch.optim.AdamW(offloaded_model.parameters(), lr=1e-2))
   	run_steps(offloaded_model, offloaded_optimizer, num_steps=3, seed=1)

[garrett361]

Optional: _move_states still spells out this DTensor unwrap and rewrap inline, so using these helpers there too would write it once.

prime-rl/src/prime_rl/trainer/optim/state_offload.py

Lines 55 to 68 in 1a8052f

	if isinstance(value, DTensor):
	local_tensor = value._local_tensor
	if device == "cpu":
	new_local = self._offload_to_cpu(param, key, local_tensor)
	else:
	new_local = local_tensor.to(device, non_blocking=True)
	new_dtensor = copy.copy(value)
	new_dtensor._local_tensor = new_local
	state[key] = new_dtensor
	elif isinstance(value, torch.Tensor):
	if device == "cpu":
	state[key] = self._offload_to_cpu(param, key, value)
	else:
	state[key] = value.to(device, non_blocking=True)

local = _local(value)
new_local = self._offload_to_cpu(param, key, local) if device == "cpu" else local.to(device, non_blocking=True)
state[key] = _with_local(value, new_local)