[rasdani]

Summary

With optimizer-state CPU offload (CPUOffloadOptimizer), step() moved every parameter's optimizer state to GPU at once. Peak GPU memory during the step therefore still included the full sharded optimizer state next to weights and gradients. On a very large model this OOMed at the first optimizer step even though forward and backward fit.

This PR steps the parameters in chunks bounded by step_chunk_bytes of optimizer state, so peak GPU memory during the step stays flat and the results are identical.

Changes

• CPUOffloadOptimizer.__init__ gets a step_chunk_bytes: int | None = 4 << 30 parameter. State size is estimated as two fp32 moments per local element. None keeps the previous single-pass behavior.
• New _chunked_step: groups parameters that have gradients into chunks, and for each chunk exposes only that chunk's gradients, moves its state to GPU, calls optimizer.step(), then moves the state back to CPU. The optimizer skips parameters without a gradient, so every parameter is updated exactly once with the same element-wise math. All gradients are restored afterwards, including on error.
• _move_states takes an optional parameter subset.
• The first (state-initializing) step also goes through the chunked path.
• Test: test_chunked_steps_match_plain_adamw is parametrized over a tiny chunk size (one parameter per chunk), a mid-size chunk and None. It checks that the parameters match plain torch.optim.AdamW exactly after 3 steps, that states end up offloaded and that gradients are cleared.

FullCPUOffloadOptimizer (full offload) is unchanged.

Testing

• ruff check and ruff format --check on the changed files pass.
• tests/unit/train/test_state_offload.py needs CUDA and could not run on the CPU-only machine I used. Please run it on a GPU runner.
• As a CPU-only sanity check, I ran the new step logic with _move_states patched to stay on CPU. It was bit-identical to plain AdamW for chunk sizes 1, 641288 and None over 3 steps.
• The same change has been used in training runs of a large MoE model with optimizer-state offload, where it removed the optimizer-step OOM.

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes distributed training optimizer stepping and memory behavior on large models; incorrect chunk planning or grad masking could desync ranks or alter updates, though tests target parity and rank agreement.

Overview
CPU-offloaded optimizer steps no longer pull the full optimizer state onto GPU at once, which was causing OOM on very large models during step() even when forward/backward fit.

CPUOffloadOptimizer gains step_chunk_bytes (default 4 GiB of estimated state per pass). It chunks parameters, moves only that chunk’s state to GPU, masks gradients so each optimizer.step() updates one chunk, then offloads again—numerically matching a single full step for element-wise optimizers. _move_states can target a parameter subset; chunk sizing uses _planned_numel on global DTensor shapes so all ranks issue the same number of step calls (important for uneven shards and for optimizers like Muon that run collectives inside step).

The trainer wires step_chunk_bytes=4<<30 for adamw and sign_sgd, and None (legacy single-pass) for other types such as Muon. Tests assert AdamW parity across chunk sizes, stable chunk plans when some params lack grads, and correct largest-shard sizing for uneven DTensor shards.

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

[cursor]

Stale Bugbot comment from a previous run.

[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 11c8a0c. Configure here.}

[garrett361]

Can certainly see a use case for this (what was the model and setup where you found you needed this? would be good to have some measurements), but the implementation needs various fixes/improvements.

From the code, it looks like you intend to use this for AdamW and SignSGD, but only the former has any state (so it doesn't make sense to use for the latter). There's also some rank-aware chunking logic, which seems unnecessary since this only applies to AdamW which operates elementwise. The chunk size is also hard-coded and you can't turn it off as it. Also some probable cuda sync issues and other nits.

Claude generated below, for your clanker

---

The chunking is correct for AdamW: torch 2.13 AdamW skips parameters whose grad is None and returns early when a group has no grads, so stepping chunk by chunk gives the same result as one full step. Smaller fixes are in the inline comments; two larger points:

1. The rank-identical chunk plan exists for Muon, which never chunks. _step_chunks sizes chunks from global shapes so every rank makes the same number of step calls, citing Muon's collectives (

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

   Lines 84 to 113 in ab15c01

   	@staticmethod
   	def _planned_numel(param: torch.Tensor) -> int:
   	"""Largest per-rank element count, from the global shape: identical on every rank.
   	Rounds up per sharded tensor dim, so a dim that does not divide evenly counts its largest shard."""
   	if not isinstance(param, DTensor):
   	return param.numel()
   	shape = list(param.shape)
   	for mesh_dim, placement in enumerate(param.placements):
   	if placement.is_shard():
   	dim = placement.dim
   	shape[dim] = -(-shape[dim] // param.device_mesh.size(mesh_dim))
   	return math.prod(shape)
   	def _step_chunks(self, params: list[torch.Tensor]) -> list[list[torch.Tensor]]:
   	"""Group parameters into passes of at most ``step_chunk_bytes`` of estimated state.
   	Sized from each parameter's global shape and placements rather than its local shard, so every
   	rank derives the same plan: uneven shards must not give ranks different numbers of
   	``optimizer.step`` calls (optimizers such as Muon run collectives inside ``step``)."""
   	chunks: list[list[torch.Tensor]] = [[]]
   	chunk_bytes = 0
   	for param in params:
   	state_bytes = 8 * self._planned_numel(param)
   	if chunks[-1] and chunk_bytes + state_bytes > self.step_chunk_bytes:
   	chunks.append([])
   	chunk_bytes = 0
   	chunks[-1].append(param)
   	chunk_bytes += state_bytes
   	return chunks

   ), but setup_optimizer disables chunking for Muon (

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

   Line 98 in ab15c01

   step_chunk_bytes = 4 << 30 if config.type in ("adamw", "sign_sgd") else None

   ) and AdamW's step has no collectives. Drop _planned_numel, the global-shape planning, and the two tests covering them (

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

   Lines 134 to 171 in ab15c01

   	def test_chunk_plan_covers_params_without_grad(monkeypatch) -> None:
   	"""The plan must not depend on which parameters have a gradient on this rank, or ranks can
   	disagree on the number of optimizer.step calls."""
   	params = [torch.nn.Parameter(torch.zeros(n)) for n in (64, 64, 64, 64)]
   	wrapper = CPUOffloadOptimizer(torch.optim.SGD(params, lr=0.1), step_chunk_bytes=8 * 128)
   	monkeypatch.setattr(wrapper, "_move_states", lambda *args, **kwargs: None)
   	plans = []
   	original = wrapper._step_chunks
   	monkeypatch.setattr(wrapper, "_step_chunks", lambda ps: plans.append(original(ps)) or plans[-1])
   	for with_grad in ((True, True, True, True), (True, False, True, False)):
   	for param, has_grad in zip(params, with_grad):
   	param.grad = torch.ones_like(param) if has_grad else None
   	wrapper._chunked_step()
   	assert [len(plan) for plan in plans] == [2, 2]
   	assert [[id(p) for p in chunk] for chunk in plans[0]] == [[id(p) for p in chunk] for chunk in plans[1]]
   	# The second pass leaves parameters without a gradient untouched and their grad None.
   	assert params[1].grad is None and torch.equal(params[1].data, torch.full((64,), -0.1))
   	assert torch.equal(params[0].data, torch.full((64,), -0.2))
   	def test_planned_numel_counts_the_largest_uneven_shard() -> None:
   	"""A dim that does not divide evenly across ranks must count its largest shard, not numel / ranks."""
   	import torch.distributed as dist
   	from torch.distributed.device_mesh import init_device_mesh
   	from torch.distributed.tensor import DTensor, Shard
   	from torch.testing._internal.distributed.fake_pg import FakeStore
   	dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=2)
   	try:
   	mesh = init_device_mesh("cpu", (2,))
   	param = DTensor.from_local(
   	torch.zeros(2, 1000), mesh, [Shard(0)], shape=torch.Size([3, 1000]), stride=(1000, 1)
   	)
   	# Rank 0 holds rows 0-1 (2000 elements); ceil(3000 / 2) would say 1500.
   	assert CPUOffloadOptimizer._planned_numel(param) == 2000
   	assert CPUOffloadOptimizer._planned_numel(torch.zeros(3, 1000)) == 3000
   	finally:
   	dist.destroy_process_group()

   ), or name the chunked optimizer that needs them. The check that a parameter without a gradient is left untouched is worth keeping.
2. Evidence. Please post peak allocated and peak reserved memory, and step time, before and after this change on the run that OOMed.

[garrett361]

Can this default to None, with the size coming from a config field that setup_optimizer passes through? 4 GiB is also hard-coded in optim/__init__.py:98, and with this default any caller that omits the argument gets chunking, Muon included.

[garrett361]

SignSGD keeps no optimizer state, so chunking it moves nothing and only adds step() calls. Chunk adamw only?

[garrett361]

Nit: nothing passes a closure today, but if something did, it would run once per chunk (loss and backward recomputed each time). Suggest assert closure is None on the chunked path.

[garrett361]

_move_states("cpu", ...) ends with torch.cuda.synchronize(), so this is now one host sync per chunk instead of one per step. One sync after the last chunk should be enough, since the caching allocator is stream-ordered.

[garrett361]

Nit: this isn't strictly an upper bound. Chunks are built from whole parameters, so a parameter whose local shard alone exceeds step_chunk_bytes becomes its own chunk, above the limit.

[garrett361]

These tests need CUDA, and the PR says they weren't run. Please run them and confirm they pass.