feat: chunked per-layer optimizer step with stream-overlapped CPU offload

docs/scaling.md

@@ -146,6 +146,16 @@ optim_cpu_offload = true   # already the default
 
 Mutually exclusive with `fsdp_cpu_offload`. Also incompatible with `trainer.max_concurrent_runs > 1` (multi-tenant training) — set `optim_cpu_offload = false` for multi-run. Muon doesn't support `fsdp_cpu_offload` but does support `optim_cpu_offload`.
 
+The optimizer step is performed per-transformer-layer ("chunked"): each layer's optimizer states are moved to GPU, the step runs for that layer only, and the states move back to CPU before the next layer. This bounds peak GPU optimizer-state memory to a single layer's worth, preventing OOM when `weight + grad + all_opt_states` exceeds available VRAM even without activations. H2D and D2H transfers are pipelined on dedicated CUDA streams so the next layer's states are prefetched while the current layer computes:
+
+```toml
+[trainer.model]
+optim_cpu_offload_chunked = true
+optim_cpu_offload_stream = true
+```
+
+Set `optim_cpu_offload_chunked = false` to fall back to all-at-once offloading (loads the full model's optimizer states to GPU in one shot — simpler but requires enough VRAM to hold them all simultaneously). Set `optim_cpu_offload_stream = false` to disable stream overlap within the chunked step, falling back to a simple sequential move→step→move loop.
+
 ### LM Head Chunking
 
 The vanilla LM head materializes a `[batch * seq, vocab]` logits tensor on every step — a major memory tax when the vocabulary is large (often >100K). `fused_lm_head_token_chunk_size` swaps in a custom fused linear + logprob/entropy kernel that streams through `chunk_size` tokens at a time, avoiding the materialization. It defaults to `1024` for RL training:

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

@@ -171,6 +171,12 @@ class ModelConfig(BaseModelConfig):
     optim_cpu_offload: bool = True
     """Offload only optimizer states (momentum, variance) to CPU, keeping weights on GPU. Avoids the H2D all-gather overhead of FSDP CPU offload while still saving GPU memory."""
 
+    optim_cpu_offload_chunked: bool = False
+    """When optimizer CPU offload is enabled, perform the optimizer step per-transformer-layer instead of all-at-once. Each layer's states are moved to GPU, the step runs for that layer only, and states move back to CPU before the next layer. Bounds peak GPU optimizer-state memory to one layer's worth, preventing OOM when weight + grad + all_opt_states exceeds available VRAM."""
+
+    optim_cpu_offload_stream: bool = False
+    """When chunked optimizer CPU offload is enabled, overlap H2D (state→GPU) and D2H (state→CPU) transfers with optimizer compute using dedicated CUDA streams. Disabling simplifies the transfer logic to a sequential move→step→move loop at the cost of serializing transfers with compute."""
+
     reshard_after_forward: bool = True
     """Reshard the model after each forward pass."""
 

src/prime_rl/trainer/optim.py

@@ -1,4 +1,5 @@
 import copy
+import re
 from typing import Callable
 
 import torch
@@ -20,58 +21,263 @@ class CPUOffloadOptimizer:
     Unlike FSDP's CPUOffload which offloads weights too, this keeps weights on GPU.
     With activation checkpointing, activations and optimizer states are never on GPU
     at the same time: peak memory becomes max(activations, opt_states) instead of sum.
+
+    When ``named_params`` is provided, the step is performed per-transformer-layer
+    ("chunked") instead of all-at-once.  Each chunk's optimizer states are moved to
+    GPU, the optimizer step is run for just that chunk, and the states are moved back
+    to CPU before the next chunk.  This bounds peak GPU optimizer-state memory to a
+    single layer's worth rather than the full model's, preventing OOM when
+    weight + grad + all_opt_states exceeds available VRAM.
+
+    H2D (state→GPU) and D2H (state→CPU) transfers are issued on dedicated CUDA
+    streams so they overlap with the optimizer compute of the current chunk: the next
+    chunk's states are prefetched while the current chunk computes, and the previous
+    chunk's states are evicted while the next chunk computes. The cost is one extra
+    layer's worth of optimizer state on GPU during the overlap window, which is
+    negligible compared to the savings from not loading the full model's states.
     """
 
-    def __init__(self, optimizer: Optimizer, pin_memory: bool = True):
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        named_params: list[tuple[str, nn.Parameter]] | None = None,
+        pin_memory: bool = True,
+        stream: bool = True,
+    ):
         self.optimizer = optimizer
         self.pin_memory = pin_memory
+        self.stream = stream
         self._initialized = False
 
+        # Build per-layer chunks from the optimizer's actual param_groups so
+        # we only include params the optimizer tracks (excludes frozen params
+        # that Muon / _create_optimizer filtered out).
+        self._chunks: list[list[nn.Parameter]] | None = None
+        if named_params is not None:
+            self._chunks = self._build_chunks(named_params)
+
+    # ------------------------------------------------------------------ #
+    # Chunk construction
+    # ------------------------------------------------------------------ #
+
+    @staticmethod
+    def _extract_layer_idx(name: str) -> int | None:
+        """Return the transformer-layer index encoded in a parameter name.
+
+        Matches patterns like ``model.layers.0.self_attn.q_proj.weight`` or
+        ``language_model.layers.12.mlp.experts.3.weight`` and returns the
+        integer layer index.  Returns ``None`` for non-layer params
+        (embed_tokens, lm_head, norm, …).
+        """
+        m = re.search(r"layers\.(\d+)\.", name)
+        return int(m.group(1)) if m else None
+
+    def _build_chunks(self, named_params: list[tuple[str, nn.Parameter]]) -> list[list[nn.Parameter]]:
+        """Partition optimizer parameters into per-layer chunks.
+
+        Parameters from the same transformer layer form one chunk; non-layer
+        parameters (embeddings, lm_head, final norm) are collected into a
+        trailing "misc" chunk.  Only parameters actually present in the
+        optimizer's ``param_groups`` are included.
+        """
+        # Map param identity → layer index
+        param_layer: dict[int, int | None] = {}
+        for name, param in named_params:
+            param_layer[id(param)] = self._extract_layer_idx(name)
+
+        # Group optimizer params by layer
+        by_layer: dict[int, list[nn.Parameter]] = {}
+        misc: list[nn.Parameter] = []
+        for group in self.optimizer.param_groups:
+            for param in group["params"]:
+                layer_idx = param_layer.get(id(param))
+                if layer_idx is not None:
+                    by_layer.setdefault(layer_idx, []).append(param)
+                else:
+                    misc.append(param)
+
+        chunks = [by_layer[k] for k in sorted(by_layer)]
+        if misc:
+            chunks.append(misc)
+        return chunks or [misc]  # edge-case: no layer params at all
+
+    def _chunk_param_ids(self, chunk_idx: int) -> set[int]:
+        return {id(p) for p in self._chunks[chunk_idx]}
+
+    # ------------------------------------------------------------------ #
+    # State movement
+    # ------------------------------------------------------------------ #
+
+    def _move_tensor(self, v: torch.Tensor, device: str) -> torch.Tensor:
+        if device == "cpu":
+            if self.pin_memory:
+                # Allocate pinned CPU destination first, then async-copy into it.
+                # This avoids the race where pin_memory() would read a tensor
+                # whose async D2H copy hasn't completed yet.
+                dst = torch.empty_like(v, device="cpu").pin_memory()
+                dst.copy_(v, non_blocking=True)
+                return dst
+            return v.to("cpu")
+        return v.to(device, non_blocking=True)
+
+    def _move_dtensor(self, v: DTensor, device: str) -> DTensor:
+        new_local = self._move_tensor(v._local_tensor, device)
+        new_dtensor = copy.copy(v)
+        new_dtensor._local_tensor = new_local
+        return new_dtensor
+
     def _move_states(self, device: str):
-        """Move optimizer states to CPU or back to GPU (matching each parameter's device)."""
+        """Move **all** optimizer states to *device* (used by state_dict / first step)."""
         for p in self.optimizer.state:
             state = self.optimizer.state[p]
             for k, v in state.items():
                 if isinstance(v, DTensor):
-                    local_tensor = v._local_tensor
-                    if device == "cpu":
-                        non_blocking = not self.pin_memory
-                        new_local = local_tensor.to("cpu", non_blocking=non_blocking)
-                        if self.pin_memory and not new_local.is_pinned():
-                            new_local = new_local.pin_memory()
-                    else:
-                        new_local = local_tensor.to(device, non_blocking=True)
-                    new_dtensor = copy.copy(v)
-                    new_dtensor._local_tensor = new_local
-                    state[k] = new_dtensor
+                    state[k] = self._move_dtensor(v, device)
                 elif isinstance(v, torch.Tensor):
-                    if device == "cpu":
-                        non_blocking = not self.pin_memory
-                        cpu_tensor = v.to("cpu", non_blocking=non_blocking)
-                        if self.pin_memory and not cpu_tensor.is_pinned():
-                            cpu_tensor = cpu_tensor.pin_memory()
-                        state[k] = cpu_tensor
-                    else:
-                        state[k] = v.to(device, non_blocking=True)
+                    state[k] = self._move_tensor(v, device)
+
+    def _move_chunk_states(self, chunk_idx: int, device: str, stream: torch.cuda.Stream | None = None):
+        """Move optimizer states for a single chunk to/from *device*.
+
+        When *stream* is provided the transfers are issued on that stream so they
+        can overlap with work on the default stream.
+        """
+        chunk_ids = self._chunk_param_ids(chunk_idx)
+
+        def _do_move():
+            for p in self.optimizer.state:
+                if id(p) not in chunk_ids:
+                    continue
+                state = self.optimizer.state[p]
+                for k, v in state.items():
+                    if isinstance(v, DTensor):
+                        state[k] = self._move_dtensor(v, device)
+                    elif isinstance(v, torch.Tensor):
+                        state[k] = self._move_tensor(v, device)
+
+        if stream is not None:
+            with torch.cuda.stream(stream):
+                _do_move()
+        else:
+            _do_move()
+
+    # ------------------------------------------------------------------ #
+    # Per-chunk optimizer step
+    # ------------------------------------------------------------------ #
+
+    def _build_chunk_param_groups(self, chunk_idx: int) -> list[dict]:
+        """Create param_groups containing only the current chunk's params.
+
+        Hyperparameters (lr, weight_decay, betas, algorithm, step, …) are
+        copied from the original groups so the optimizer behaves identically.
+        """
+        chunk_ids = self._chunk_param_ids(chunk_idx)
+        chunk_groups: list[dict] = []
+        for group in self._original_param_groups:
+            filtered = [p for p in group["params"] if id(p) in chunk_ids]
+            if not filtered:
+                continue
+            new_group = {k: v for k, v in group.items() if k != "params"}
+            new_group["params"] = filtered
+            chunk_groups.append(new_group)
+        return chunk_groups
+
+    def _step_chunk(self, chunk_idx: int, closure=None):
+        """Run ``optimizer.step()`` for a single chunk by temporarily swapping param_groups."""
+        chunk_groups = self._build_chunk_param_groups(chunk_idx)
+        self.optimizer.param_groups = chunk_groups
+        result = self.optimizer.step(closure)
+        self.optimizer.param_groups = self._original_param_groups
+        return result
+
+    def _sync_step_counters(self, original_steps: list):
+        """Increment the original param_groups' step counters by one.
+
+        Muon stores a per-group ``step`` integer that ``step()`` increments at the
+        top of the call.  Because we swap in per-chunk copies, the *original*
+        groups never see the increment.  Standard optimizers (AdamW, SGD,
+        SignSGD) keep the step counter in ``state[p]['step']`` which is
+        per-parameter and thus already correct (each param is in exactly one
+        chunk and its state is touched exactly once).
+        """
+        for group, orig_step in zip(self._original_param_groups, original_steps):
+            if "step" in group:
+                group["step"] = orig_step + 1
+
+    # ------------------------------------------------------------------ #
+    # Public API
+    # ------------------------------------------------------------------ #
 
     def step(self, closure=None):
-        # First step initializes states on GPU - offload after
+        # First step initializes states on GPU — offload after.
         if not self._initialized:
             result = self.optimizer.step(closure)
             self._move_states("cpu")
             self._initialized = True
             return result
 
-        # Move states to GPU
-        self._move_states("cuda")
+        # Non-chunked fallback (no named_params provided).
+        if self._chunks is None or len(self._chunks) <= 1:
+            self._move_states("cuda")
+            result = self.optimizer.step(closure)
+            self._move_states("cpu")
+            return result
 
-        # Run optimizer step
-        result = self.optimizer.step(closure)
+        # Chunked step.
+        self._original_param_groups = self.optimizer.param_groups
+        original_steps = [g.get("step", 0) for g in self._original_param_groups]
 
-        # Move states back to CPU
-        self._move_states("cpu")
+        if self.stream:
+            self._step_chunked_streamed(closure)
+        else:
+            self._step_chunked(closure)
+
+        self._sync_step_counters(original_steps)
+        self._original_param_groups = None
+        return None
+
+    def _step_chunked(self, closure=None):
+        """Per-layer optimizer step without stream overlap — simple sequential loop."""
+        for i in range(len(self._chunks)):
+            self._move_chunk_states(i, "cuda")
+            self._step_chunk(i, closure)
+            self._move_chunk_states(i, "cpu")
+
+    def _step_chunked_streamed(self, closure=None):
+        """Per-layer optimizer step with stream-overlapped H2D / D2H.
+
+        H2D (state→GPU) for chunk i+1 is issued on a dedicated stream while
+        the optimizer compute for chunk i runs on the compute stream.  D2H
+        (state→CPU) for chunk i runs on a second dedicated stream, overlapping
+        with the compute for chunk i+1.
+        """
+        h2d_stream = torch.cuda.Stream()
+        d2h_stream = torch.cuda.Stream()
+        compute_stream = torch.cuda.current_stream()
+        n = len(self._chunks)
 
-        return result
+        # Issue H2D for chunk 0 on the h2d stream.
+        self._move_chunk_states(0, "cuda", h2d_stream)
+
+        for i in range(n):
+            # Wait for H2D of chunk i to land on GPU.
+            compute_stream.wait_stream(h2d_stream)
+
+            # Prefetch H2D for chunk i+1 so it overlaps with step i on the
+            # compute stream.
+            if i + 1 < n:
+                self._move_chunk_states(i + 1, "cuda", h2d_stream)
+
+            # Run optimizer step for chunk i.
+            self._step_chunk(i, closure)
+
+            # Evict chunk i's states back to CPU (overlaps with step i+1).
+            d2h_stream.wait_stream(compute_stream)
+            self._move_chunk_states(i, "cpu", d2h_stream)
+
+        # Ensure all transfers are complete before returning.
+        compute_stream.wait_stream(h2d_stream)
+        compute_stream.wait_stream(d2h_stream)
 
     def zero_grad(self, set_to_none: bool = True):
         self.optimizer.zero_grad(set_to_none=set_to_none)
@@ -114,6 +320,8 @@ def setup_optimizer(
     parallel_dims: ParallelDims,
     lora: bool = False,
     cpu_offload: bool = False,
+    cpu_offload_chunked: bool = True,
+    cpu_offload_stream: bool = True,
 ) -> Optimizer | CPUOffloadOptimizer:
     if lora:
         # Wait for run 0 to be created in the multi run manager
@@ -126,7 +334,11 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        return CPUOffloadOptimizer(optimizer)
+        return CPUOffloadOptimizer(
+            optimizer,
+            named_params=named_params if cpu_offload_chunked else None,
+            stream=cpu_offload_stream,
+        )
 
     return optimizer
 

src/prime_rl/trainer/rl/train.py

@@ -166,6 +166,8 @@ def train(config: TrainerConfig):
             parallel_dims,
             lora=config.model.lora is not None,
             cpu_offload=config.model.optim_cpu_offload,
+            cpu_offload_chunked=config.model.optim_cpu_offload_chunked,
+            cpu_offload_stream=config.model.optim_cpu_offload_stream,
         )
         scheduler = setup_scheduler(optimizer, config.scheduler, config.max_steps, config.optim.lr)
     else:

src/prime_rl/trainer/sft/train.py

@@ -178,7 +178,12 @@ def train(config: SFTConfig):
     # Set up the optimizer
     logger.info(f"Initializing optimizer ({config.optim})")
     optimizer = setup_optimizer(
-        config.optim, list(model.named_parameters()), parallel_dims, cpu_offload=config.model.optim_cpu_offload
+        config.optim,
+        list(model.named_parameters()),
+        parallel_dims,
+        cpu_offload=config.model.optim_cpu_offload,
+        cpu_offload_chunked=config.model.optim_cpu_offload_chunked,
+        cpu_offload_stream=config.model.optim_cpu_offload_stream,
     )
 
     # Set up the learning rate scheduler