perf(optim): overlap the offloaded AdamW step with the next forward (builds on #4008)

docs/scaling.md

@@ -165,7 +165,7 @@ Every chunk runs the MoE's EP all-to-alls, so all ranks of an EP group must run
 
 ### Optimizer Offloading
 
-State-only optimizer offload is enabled by default (`model.offload = "optimizer"`). For full offload, set `model.offload = "full"`; this keeps BF16 compute weights on GPU and runs CPU optimizer chunks as gradients become ready during backward. Full offload only supports AdamW and SignSGD (`optim.type = "sign_sgd"`) and disables gradient clipping. SignSGD is stateless, so it halves the host RAM footprint versus AdamW (8 instead of 16 bytes per parameter: FP32 master + FP32 accumulated gradient, no moments).
+State-only optimizer offload is enabled by default (`model.offload = "optimizer"`). With AdamW, the step streams the states through the GPU in chunks; set `model.offload_overlap_forward = true` to run those copies and updates alongside the next forward, each FSDP module waiting only for its own parameters (off by default). For full offload, set `model.offload = "full"`; this keeps BF16 compute weights on GPU and runs CPU optimizer chunks as gradients become ready during backward. Full offload only supports AdamW and SignSGD (`optim.type = "sign_sgd"`) and disables gradient clipping. SignSGD is stateless, so it halves the host RAM footprint versus AdamW (8 instead of 16 bytes per parameter: FP32 master + FP32 accumulated gradient, no moments).
 
 ### LM Head Chunking
 

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -304,6 +304,9 @@ class ModelConfig(BaseModelConfig):
     offload: OffloadMode = "optimizer"
     """CPU offloading. ``optimizer`` keeps optimizer states (momentum, variance) in CPU RAM and weights on GPU. ``full`` keeps FP32 masters, optimizer states, and accumulated gradients in CPU RAM and runs each optimizer chunk on CPU as soon as its last gradient arrives, overlapped with backward; it supports AdamW and SignSGD only, disables gradient clipping, and rounds each gradient to BF16 once before the FP32 CPU update. ``none`` keeps everything on GPU."""
 
+    offload_overlap_forward: bool = False
+    """With ``offload = "optimizer"`` and AdamW, run the optimizer step's state copies and updates alongside the next forward instead of waiting for them: each FSDP module waits only for the chunks holding its own parameters before its all-gather, and weight broadcasts wait for the whole step first. The math is unchanged; the step's gradients stay allocated until their chunk has updated. Off by default until it has been exercised in RL runs."""
+
     reshard_after_forward: bool = True
     """Reshard the model after each forward pass."""
 

src/prime_rl/trainer/optim/__init__.py

@@ -53,6 +53,7 @@ def setup_optimizer(
     named_params: list[tuple[str, nn.Parameter]],
     parallel_dims: ParallelDims,
     offload: OffloadMode = "none",
+    offload_overlap_forward: bool = False,
     model: nn.Module | None = None,
     full_offload_dtype_policy: dict[int, tuple[torch.dtype, torch.dtype]] | None = None,
 ) -> tuple[OptimizerLike, GradientOffloadManager | None]:
@@ -94,11 +95,18 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        return CPUOffloadOptimizer(optimizer), None
+        overlap_model = model if offload_overlap_forward and config.type == "adamw" else None
+        return CPUOffloadOptimizer(optimizer, overlap_model=overlap_model), None
 
     return optimizer, None
 
 
+def wait_for_optimizer_step(optimizer: OptimizerLike) -> None:
+    """Make the current stream wait for a step that would otherwise overlap the next forward."""
+    if isinstance(optimizer, CPUOffloadOptimizer):
+        optimizer.wait()
+
+
 def _create_optimizer(
     config: OptimizerConfig,
     named_params: list[tuple[str, nn.Parameter]],

src/prime_rl/trainer/optim/state_offload.py

@@ -1,10 +1,140 @@
 import copy
 
 import torch
+from torch import nn
+from torch.distributed.fsdp import FSDPModule
+from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup
 from torch.distributed.tensor import DTensor
-from torch.optim import Optimizer
+from torch.optim import AdamW, Optimizer
 
 from prime_rl.trainer.optim.base import OffloadOptimizer
+from prime_rl.utils.logger import get_logger
+
+# Byte budget of one staging chunk. A single state larger than this gets a chunk of its own.
+CHUNK_BYTES = 1 << 30
+# Alignment of each state view inside a staging buffer.
+ALIGN_BYTES = 512
+# An overlapped step enqueues this many chunks up front, so the copies run while the trainer
+# finishes its step bookkeeping, and later stays this many chunks ahead of the forward. Enqueuing
+# every chunk at once would fill the CUDA launch queue and block the CPU until the copies drain.
+INITIAL_CHUNKS = 8
+LOOKAHEAD_CHUNKS = 4
+
+
+def _local(value: torch.Tensor) -> torch.Tensor:
+    return value._local_tensor if isinstance(value, DTensor) else value
+
+
+def _with_local(value: torch.Tensor, local: torch.Tensor) -> torch.Tensor:
+    if not isinstance(value, DTensor):
+        return local
+    new_dtensor = copy.copy(value)
+    new_dtensor._local_tensor = local
+    return new_dtensor
+
+
+class _ChunkedStep:
+    """One chunked AdamW step. Chunks stream through two GPU staging buffers: while chunk i updates,
+    chunk i+1 loads on the h2d stream and chunk i-1 writes back to pinned memory on the d2h stream.
+
+    Chunks are enqueued on demand by ``advance``. The step's hyperparameters and gradients are
+    captured up front, so a chunk enqueued after the scheduler step or ``zero_grad`` updates
+    exactly as it would have at step time."""
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        chunks: list[tuple[set[torch.Tensor], list[tuple[torch.Tensor, str, int]]]],
+        buffer_bytes: int,
+        h2d: torch.cuda.Stream,
+        d2h: torch.cuda.Stream,
+        update: torch.cuda.Stream,
+    ):
+        self.optimizer = optimizer
+        self.chunks = chunks
+        self.h2d, self.d2h, self.update = h2d, d2h, update
+        compute = torch.cuda.current_stream()
+        self.overlapped = update != compute
+        # Allocated on the compute stream, so the caching allocator reuses the same blocks every step.
+        self.buffers = [torch.empty(buffer_bytes, dtype=torch.uint8, device="cuda") for _ in range(2)]
+        h2d.wait_stream(compute)
+        # The previous step's writebacks must land before the loads read the CPU states again.
+        h2d.wait_stream(d2h)
+        update.wait_stream(compute)
+        self.hyperparams = [{k: v for k, v in group.items() if k != "params"} for group in optimizer.param_groups]
+        self.grads = {param: param.grad for params, _ in chunks for param in params}
+        self.written_back: list[torch.cuda.Event | None] = [None, None]
+        self.cpu_values: list[dict[tuple[torch.Tensor, str], torch.Tensor]] = [{}, {}]
+        self.updated: list[torch.cuda.Event] = []
+        self.loaded = self._load(0)
+
+    @property
+    def done(self) -> bool:
+        return len(self.updated) == len(self.chunks)
+
+    def _load(self, i: int) -> torch.cuda.Event:
+        state = self.optimizer.state
+        buffer, slot = self.buffers[i % 2], self.cpu_values[i % 2]
+        with torch.cuda.stream(self.h2d):
+            if self.written_back[i % 2] is not None:
+                self.h2d.wait_event(self.written_back[i % 2])
+            slot.clear()
+            for param, key, offset in self.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 self.h2d.record_event()
+
+    def _step_params(self, params: set[torch.Tensor]):
+        groups = self.optimizer.param_groups
+        current = [dict(group) for group in groups]
+        current_grads = {param: param.grad for param in params}
+        for group, hyperparams in zip(groups, self.hyperparams):
+            group.update(hyperparams)
+            group["params"] = [param for param in group["params"] if param in params]
+        for param in params:
+            param.grad = self.grads[param]
+        try:
+            self.optimizer.step()
+        finally:
+            for param in params:
+                param.grad = current_grads[param]
+            for group, values in zip(groups, current):
+                group.update(values)
+
+    def advance(self, until: int) -> None:
+        """Enqueue every chunk up to and including chunk ``until``."""
+        state = self.optimizer.state
+        while not self.done and len(self.updated) <= until:
+            i = len(self.updated)
+            params = self.chunks[i][0]
+            next_loaded = self._load(i + 1) if i + 1 < len(self.chunks) else None
+            self.update.wait_event(self.loaded)
+            with torch.cuda.stream(self.update):
+                self._step_params(params)
+                self.updated.append(self.update.record_event())
+            self.d2h.wait_stream(self.update)
+            with torch.cuda.stream(self.d2h):
+                for (param, key), cpu_value in self.cpu_values[i % 2].items():
+                    _local(cpu_value).copy_(_local(state[param][key]), non_blocking=True)
+                    state[param][key] = cpu_value
+                self.written_back[i % 2] = self.d2h.record_event()
+            self.loaded = next_loaded
+            for param in params:
+                grad = self.grads.pop(param)
+                if self.overlapped and grad is not None:
+                    # zero_grad may already have let go of it: keep the compute stream off its
+                    # memory until the update has read it.
+                    _local(grad).record_stream(self.update)
+        if self.done and self.overlapped and self.buffers:
+            for buffer in self.buffers:
+                for stream in (self.h2d, self.update, self.d2h):
+                    buffer.record_stream(stream)
+            self.buffers = []
 
 
 class CPUOffloadOptimizer(OffloadOptimizer):
@@ -15,15 +145,83 @@ 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, overlap_model: nn.Module | None = None):
+        """With ``overlap_model`` (FSDP-sharded), a chunked AdamW step returns before its chunks
+        run: they update on a side stream alongside the next forward, enqueued as the forward reaches
+        them, and each FSDP parameter group waits only for the chunks holding its parameters before
+        its all-gather. Anything else that reads the parameters must call ``wait()`` first."""
         self.optimizer = optimizer
-        self.pin_memory = pin_memory
         self._initialized = False
         self._cpu_buffers: dict[tuple[torch.Tensor, str], torch.Tensor] = {}
+        self._streams: tuple[torch.cuda.Stream, torch.cuda.Stream, torch.cuda.Stream] | None = None
+        self._fsdp_groups: dict[torch.Tensor, FSDPParamGroup] = {}
+        self._in_flight: _ChunkedStep | None = None
+        # Last chunk of the in-flight step holding each FSDP group's parameters.
+        self._group_chunks: dict[FSDPParamGroup, int] = {}
+        if overlap_model is not None:
+            self._install_unshard_waits(overlap_model)
+
+    def _install_unshard_waits(self, model: nn.Module) -> None:
+        fsdp_groups = {}
+        for module in model.modules():
+            if isinstance(module, FSDPModule):
+                # A module holds one group per mesh (e.g. experts on the EP mesh next to dense weights).
+                for group in module._get_fsdp_state()._fsdp_param_groups:
+                    for fsdp_param in group.fsdp_params:
+                        fsdp_groups[fsdp_param.sharded_param] = group
+        params = [p for group in self.optimizer.param_groups for p in group["params"]]
+        missing = sum(p not in fsdp_groups for p in params)
+        if missing:
+            get_logger().warning(
+                f"{missing} of {len(params)} optimizer parameters are not managed by FSDP; "
+                "the offloaded optimizer step will not overlap the next forward"
+            )
+            return
+        self._fsdp_groups = fsdp_groups
+        for group in set(fsdp_groups.values()):
+            group.unshard = self._unshard_after_step(group, group.unshard)
+            group.wait_for_unshard = self._wait_for_unshard_after_step(group, group.wait_for_unshard)
+
+    def _updated(self, chunk: int) -> torch.cuda.Event:
+        self._in_flight.advance(chunk + LOOKAHEAD_CHUNKS)
+        return self._in_flight.updated[chunk]
+
+    def _unshard_after_step(self, group: FSDPParamGroup, unshard):
+        def wrapper(async_op: bool = False):
+            chunk = self._group_chunks.get(group)
+            if chunk is not None:
+                event = self._updated(chunk)
+                # The all-gather copy-in reads the sharded parameters on these streams.
+                for stream in set(group.comm_ctx.get_all_gather_streams(async_op, group._training_state)):
+                    stream.wait_event(event)
+            return unshard(async_op)
+
+        return wrapper
+
+    def _wait_for_unshard_after_step(self, group: FSDPParamGroup, wait_for_unshard):
+        def wrapper():
+            chunk = self._group_chunks.pop(group, None)
+            if chunk is not None:
+                torch.cuda.current_stream().wait_event(self._updated(chunk))
+            return wait_for_unshard()
+
+        return wrapper
+
+    def _finish_in_flight(self) -> None:
+        if self._in_flight is not None:
+            self._in_flight.advance(len(self._in_flight.chunks))
+            self._in_flight = None
+            self._group_chunks.clear()
+
+    def wait(self) -> None:
+        """Make the current stream wait for an overlapped step's updates and state writebacks."""
+        self._finish_in_flight()
+        if self._streams is not None:
+            current = torch.cuda.current_stream()
+            for stream in self._streams:
+                current.wait_stream(stream)
 
     def _offload_to_cpu(self, param: torch.Tensor, key: str, src: torch.Tensor) -> torch.Tensor:
-        if not self.pin_memory:
-            return src.to("cpu", non_blocking=True)
         buffer = self._cpu_buffers.get((param, key))
         if buffer is None or buffer.shape != src.shape or buffer.dtype != src.dtype:
             buffer = torch.empty(src.shape, dtype=src.dtype, device="cpu", pin_memory=True)
@@ -34,23 +232,18 @@ def _offload_to_cpu(self, param: torch.Tensor, key: str, src: torch.Tensor) -> t
 
     def _move_states(self, device: str):
         """Move optimizer states to CPU or back to GPU (matching each parameter's device)."""
+        self.wait()
         for param in self.optimizer.state:
             state = self.optimizer.state[param]
             for key, value in state.items():
-                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)
+                if not isinstance(value, torch.Tensor):
+                    continue
+                local = _local(value)
+                if device == "cpu":
+                    new_local = self._offload_to_cpu(param, key, local)
+                else:
+                    new_local = local.to(device, non_blocking=True)
+                state[key] = _with_local(value, new_local)
         if device == "cpu" and torch.cuda.is_initialized():
             torch.cuda.synchronize()
 
@@ -62,6 +255,12 @@ def step(self, closure=None):
             self._initialized = True
             return result
 
+        # A parameter without state yet (no grad on any earlier step) is initialized by the full path below.
+        if isinstance(self.optimizer, AdamW) and all(
+            p in self.state for group in self.param_groups for p in group["params"] if p.grad is not None
+        ):
+            return self._chunked_step()
+
         # Move states to GPU
         self._move_states("cuda")
 
@@ -73,7 +272,56 @@ def step(self, closure=None):
 
         return result
 
+    def _plan_chunks(self) -> tuple[list[tuple[set[torch.Tensor], list[tuple[torch.Tensor, str, int]]]], int]:
+        """Group whole parameters into chunks; each state except the CPU step counter gets an offset into a staging buffer."""
+        chunks, params, entries, used, max_used = [], set(), [], 0, 0
+        for param, state in self.optimizer.state.items():
+            param_entries, param_bytes = [], 0
+            for key, value in state.items():
+                if key == "step":
+                    continue
+                local = _local(value)
+                param_entries.append((param, key, param_bytes))
+                param_bytes += -(-local.numel() * local.element_size() // ALIGN_BYTES) * ALIGN_BYTES
+            if params and used + param_bytes > CHUNK_BYTES:
+                chunks.append((params, entries))
+                params, entries, used = set(), [], 0
+            params.add(param)
+            entries.extend((param, key, used + offset) for param, key, offset in param_entries)
+            used += param_bytes
+            max_used = max(max_used, used)
+        if params:
+            chunks.append((params, entries))
+        return chunks, max_used
+
+    def _chunked_step(self):
+        """Step chunk by chunk (see ``_ChunkedStep``). Without an overlap model the chunks update on
+        the compute stream and the step waits for the GPU; with one they update on a side stream,
+        enqueued as the next forward reaches their parameters.
+
+        The step counters stay on CPU, as in plain torch AdamW, so the update never syncs the device."""
+        self._finish_in_flight()
+        if self._streams is None:
+            self._streams = (torch.cuda.Stream(), torch.cuda.Stream(), torch.cuda.Stream())
+        h2d, d2h, side = self._streams
+        chunks, buffer_bytes = self._plan_chunks()
+        if not chunks:
+            return
+        if not self._fsdp_groups:
+            step = _ChunkedStep(self.optimizer, chunks, buffer_bytes, h2d, d2h, torch.cuda.current_stream())
+            step.advance(len(chunks))
+            torch.cuda.synchronize()
+            return
+        self._in_flight = _ChunkedStep(self.optimizer, chunks, buffer_bytes, h2d, d2h, side)
+        for i, (params, _) in enumerate(chunks):
+            for param in params:
+                self._group_chunks[self._fsdp_groups[param]] = i
+        self._in_flight.advance(INITIAL_CHUNKS - 1)
+
     def zero_grad(self, set_to_none: bool = True):
+        if not set_to_none:
+            # Zeroing in place would race the in-flight step's reads of the gradients.
+            self.wait()
         self.optimizer.zero_grad(set_to_none=set_to_none)
 
     def state_dict(self):