From 664c7b0d93396274b4b16739952289aa41a63d64 Mon Sep 17 00:00:00 2001 From: rasdani <73563550+rasdani@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:52:56 +0000 Subject: [PATCH 1/3] feat(trainer): step CPUOffloadOptimizer in bounded chunks of optimizer state With optimizer-state CPU offload, step() moved every parameter's state to GPU at once, so peak memory still needed the full sharded optimizer state next to weights and gradients. For very large models this OOMs at the first optimizer step even when forward and backward fit. Add a step_chunk_bytes parameter (default 4 GiB, estimated as two fp32 moments per element) and step the parameters chunk by chunk with only that chunk's state on GPU. Each pass only exposes its chunk's gradients and the optimizer skips parameters without one, so every parameter is updated exactly once with the same element-wise math. step_chunk_bytes=None keeps the previous single-pass behavior. Co-Authored-By: Claude Opus 5.5 --- src/prime_rl/trainer/optim/state_offload.py | 70 ++++++++++++++++----- tests/unit/train/test_state_offload.py | 17 +++++ 2 files changed, 72 insertions(+), 15 deletions(-) diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py index 15a7f1eca2c..3e738fb412a 100644 --- a/src/prime_rl/trainer/optim/state_offload.py +++ b/src/prime_rl/trainer/optim/state_offload.py @@ -15,9 +15,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 +37,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 +65,52 @@ 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) + + def _step_chunks(self, params: list[torch.Tensor]) -> list[list[torch.Tensor]]: + chunks: list[list[torch.Tensor]] = [[]] + chunk_bytes = 0 + for param in params: + local = param._local_tensor if isinstance(param, DTensor) else param + state_bytes = 8 * local.numel() + 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.""" + params = [p for group in self.optimizer.param_groups for p in group["params"] if p.grad is not None] + 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): diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py index 68ba4015623..8d0248a2a42 100644 --- a/tests/unit/train/test_state_offload.py +++ b/tests/unit/train/test_state_offload.py @@ -112,3 +112,20 @@ 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()) From 11c8a0c3a426f8b757e4549e5a44315866c2a9a4 Mon Sep 17 00:00:00 2001 From: rasdani <73563550+rasdani@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:34:33 +0000 Subject: [PATCH 2/3] fix(trainer): plan CPUOffloadOptimizer step chunks identically on every rank Chunk sizes came from each rank's local shard and only parameters with a gradient, so uneven FSDP shards could give ranks different numbers of optimizer.step calls. Plan over all parameters with a per-rank estimate derived from the global shape and placements instead, and enable chunked stepping only for element-wise optimizers (AdamW, SignSGD); Muon runs collectives inside step and keeps a single pass. Co-Authored-By: Claude Opus 5.5 --- src/prime_rl/trainer/optim/__init__.py | 5 ++++- src/prime_rl/trainer/optim/state_offload.py | 23 ++++++++++++++++++--- tests/unit/train/test_state_offload.py | 20 ++++++++++++++++++ 3 files changed, 44 insertions(+), 4 deletions(-) diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 5988996e969..431c6f2d8ce 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/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 diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py index 3e738fb412a..1ad46fd33bb 100644 --- a/src/prime_rl/trainer/optim/state_offload.py +++ b/src/prime_rl/trainer/optim/state_offload.py @@ -80,12 +80,27 @@ def step(self, closure=None): return result return self._chunked_step(closure) + @staticmethod + def _planned_numel(param: torch.Tensor) -> int: + """Per-rank element count estimated from the global shape: identical on every rank.""" + if not isinstance(param, DTensor): + return param.numel() + shards = 1 + for dim, placement in enumerate(param.placements): + if placement.is_shard(): + shards *= param.device_mesh.size(dim) + return -(-param.numel() // shards) + 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: - local = param._local_tensor if isinstance(param, DTensor) else param - state_bytes = 8 * local.numel() + state_bytes = 8 * self._planned_numel(param) if chunks[-1] and chunk_bytes + state_bytes > self.step_chunk_bytes: chunks.append([]) chunk_bytes = 0 @@ -97,7 +112,9 @@ 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.""" - params = [p for group in self.optimizer.param_groups for p in group["params"] if p.grad is not None] + # 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: diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py index 8d0248a2a42..5ae1bad93fa 100644 --- a/tests/unit/train/test_state_offload.py +++ b/tests/unit/train/test_state_offload.py @@ -129,3 +129,23 @@ def test_chunked_steps_match_plain_adamw(step_chunk_bytes: int | None) -> None: 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)) From ab15c010bc61aba29927146af7bdef2fd547631f Mon Sep 17 00:00:00 2001 From: rasdani <73563550+rasdani@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:01:53 +0000 Subject: [PATCH 3/3] fix(trainer): size offload step chunks by the largest uneven shard _planned_numel used ceil(numel / shards), which undercounts a tensor sharded on a dim that does not divide evenly: (3, 1000) on Shard(0) over 2 ranks holds 2000 elements on rank 0, not 1500. Round up per sharded dim instead, still from the global shape so every rank derives the same plan. Co-Authored-By: Claude Opus 5.5 --- src/prime_rl/trainer/optim/state_offload.py | 14 +++++++++----- tests/unit/train/test_state_offload.py | 20 ++++++++++++++++++++ 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py index 1ad46fd33bb..ac2d833f4b4 100644 --- a/src/prime_rl/trainer/optim/state_offload.py +++ b/src/prime_rl/trainer/optim/state_offload.py @@ -1,4 +1,5 @@ import copy +import math import torch from torch.distributed.tensor import DTensor @@ -82,14 +83,17 @@ def step(self, closure=None): @staticmethod def _planned_numel(param: torch.Tensor) -> int: - """Per-rank element count estimated from the global shape: identical on every rank.""" + """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() - shards = 1 - for dim, placement in enumerate(param.placements): + shape = list(param.shape) + for mesh_dim, placement in enumerate(param.placements): if placement.is_shard(): - shards *= param.device_mesh.size(dim) - return -(-param.numel() // shards) + 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. diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py index 5ae1bad93fa..f47f40e47e3 100644 --- a/tests/unit/train/test_state_offload.py +++ b/tests/unit/train/test_state_offload.py @@ -149,3 +149,23 @@ def test_chunk_plan_covers_params_without_grad(monkeypatch) -> None: # 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()