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 15a7f1eca2c..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 @@ -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): diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py index 68ba4015623..f47f40e47e3 100644 --- a/tests/unit/train/test_state_offload.py +++ b/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()