feat(trainer): step CPUOffloadOptimizer in bounded chunks of optimizer state

src/prime_rl/trainer/optim/__init__.py

@@ -93,7 +93,10 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        return CPUOffloadOptimizer(optimizer), None
+        # Chunked stepping calls optimizer.step once per chunk, which is only safe for element-wise
+        # optimizers; Muon issues collectives inside step and keeps a single pass.
+        step_chunk_bytes = 4 << 30 if config.type in ("adamw", "sign_sgd") else None
+        return CPUOffloadOptimizer(optimizer, step_chunk_bytes=step_chunk_bytes), None
 
     return optimizer, None
 

src/prime_rl/trainer/optim/state_offload.py

@@ -1,4 +1,5 @@
 import copy
+import math
 
 import torch
 from torch.distributed.tensor import DTensor
@@ -15,9 +16,14 @@ class CPUOffloadOptimizer(OffloadOptimizer):
     at the same time: peak memory becomes max(activations, opt_states) instead of sum.
     """
 
-    def __init__(self, optimizer: Optimizer, pin_memory: bool = True):
+    def __init__(self, optimizer: Optimizer, pin_memory: bool = True, step_chunk_bytes: int | None = 4 << 30):
         self.optimizer = optimizer
         self.pin_memory = pin_memory
+        # Upper bound on the optimizer state (estimated at two fp32 moments per element) brought
+        # to GPU per step pass. Moving all states at once needs the full sharded state on GPU next
+        # to weights and gradients, which can OOM for very large models. None steps everything in
+        # one pass.
+        self.step_chunk_bytes = step_chunk_bytes
         self._initialized = False
         self._cpu_buffers: dict[tuple[torch.Tensor, str], torch.Tensor] = {}
 
@@ -32,9 +38,12 @@ def _offload_to_cpu(self, param: torch.Tensor, key: str, src: torch.Tensor) -> t
             buffer.copy_(src, non_blocking=True)
         return buffer
 
-    def _move_states(self, device: str):
-        """Move optimizer states to CPU or back to GPU (matching each parameter's device)."""
-        for param in self.optimizer.state:
+    def _move_states(self, device: str, params=None):
+        """Move optimizer states to CPU or back to GPU (matching each parameter's device), for
+        `params` or every parameter with state."""
+        for param in self.optimizer.state if params is None else params:
+            if param not in self.optimizer.state:
+                continue
             state = self.optimizer.state[param]
             for key, value in state.items():
                 if isinstance(value, DTensor):
@@ -57,20 +66,72 @@ def _move_states(self, device: str):
     def step(self, closure=None):
         # First step initializes states on GPU - offload after
         if not self._initialized:
-            result = self.optimizer.step(closure)
-            self._move_states("cpu")
+            if self.step_chunk_bytes is not None:
+                result = self._chunked_step(closure)
+            else:
+                result = self.optimizer.step(closure)
+                self._move_states("cpu")
             self._initialized = True
             return result
 
-        # Move states to GPU
-        self._move_states("cuda")
-
-        # Run optimizer step
-        result = self.optimizer.step(closure)
-
-        # Move states back to CPU
-        self._move_states("cpu")
-
+        if self.step_chunk_bytes is None:
+            self._move_states("cuda")
+            result = self.optimizer.step(closure)
+            self._move_states("cpu")
+            return result
+        return self._chunked_step(closure)
+
+    @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
+
+    def _chunked_step(self, closure=None):
+        """Step the parameters chunk by chunk, with only one chunk's states on GPU. Each pass sees
+        only its chunk's gradients; the optimizer skips parameters without one, so every parameter
+        is updated exactly once with the same element-wise math as a single step."""
+        # Plan over every parameter (not only those with a gradient) so the plan is identical on all
+        # ranks; parameters without a gradient are skipped by the optimizer within their pass.
+        params = [p for group in self.optimizer.param_groups for p in group["params"]]
+        grads = {id(p): p.grad for p in params}
+        result = None
+        try:
+            for chunk in self._step_chunks(params):
+                ids = {id(p) for p in chunk}
+                for param in params:
+                    param.grad = grads[id(param)] if id(param) in ids else None
+                self._move_states("cuda", chunk)
+                result = self.optimizer.step(closure)
+                self._move_states("cpu", chunk)
+        finally:
+            for param in params:
+                param.grad = grads[id(param)]
         return result
 
     def zero_grad(self, set_to_none: bool = True):

tests/unit/train/test_state_offload.py

@@ -112,3 +112,60 @@ def test_offload_reuses_pinned_buffers_across_steps() -> None:
     assert pointers_before.keys() == pointers_after.keys()
     for key, pointer in pointers_before.items():
         assert pointers_after[key] == pointer, f"state buffer reallocated for {key}"
+
+
+@pytest.mark.parametrize("step_chunk_bytes", [1, 64 * 128 * 8, None])
+def test_chunked_steps_match_plain_adamw(step_chunk_bytes: int | None) -> None:
+    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)
+
+    chunked_model = make_model(0)
+    chunked_optimizer = CPUOffloadOptimizer(
+        torch.optim.AdamW(chunked_model.parameters(), lr=1e-2), step_chunk_bytes=step_chunk_bytes
+    )
+    run_steps(chunked_model, chunked_optimizer, num_steps=3, seed=1)
+
+    assert_params_equal(reference_model, chunked_model)
+    assert_states_offloaded(chunked_optimizer)
+    assert all(param.grad is None for param in chunked_model.parameters())
+
+
+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()