From 9f8dcee6f3b0ac25af4ea2cab3c08158eff2e3fa Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Wed, 22 Jul 2026 18:55:44 +0000
Subject: [PATCH 1/7] feat: chunked per-layer optimizer step with
 stream-overlapped CPU offload

The CPUOffloadOptimizer now performs the optimizer.step() per-transformer-layer
instead of all-at-once. Each layer's optimizer states are moved to GPU, the step
runs for that layer only, and states are moved back to CPU before the next layer.

This bounds peak GPU optimizer-state memory to ~one layer's worth (plus one
prefetched layer for overlap) instead of the full model's, preventing OOM when
weight + grad + all_opt_states exceeds available VRAM.

H2D and D2H transfers are pipelined on dedicated CUDA streams so the next layer's
states are prefetched while the current layer computes, and the previous layer's
states are evicted while the next layer computes.

Pinned-memory D2H is fixed to pre-allocate the pinned destination and async-copy
into it, avoiding a race where pin_memory() reads a tensor whose async copy hasn't
completed.

Muon's per-group step counter is synced after all chunks complete, since the
param_groups are temporarily swapped during chunked stepping.
---
 docs/scaling.md                   |   2 +
 src/prime_rl/trainer/optim.py     | 247 ++++++++++++++++++++++++++----
 src/prime_rl/trainer/sft/train.py |   5 +-
 3 files changed, 223 insertions(+), 31 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index d10ef6f444..645e9135ff 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -146,6 +146,8 @@ 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.
+
 ### 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:
diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py
index cb0b3d2fb0..2b0bff027c 100644
--- a/src/prime_rl/trainer/optim.py
+++ b/src/prime_rl/trainer/optim.py
@@ -1,4 +1,5 @@
 import copy
+import re
 from typing import Callable
 
 import torch
@@ -20,58 +21,244 @@ 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,
+    ):
         self.optimizer = optimizer
         self.pin_memory = pin_memory
         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 with stream-overlapped H2D / D2H ---- #
 
-        # Move states back to CPU
-        self._move_states("cpu")
+        self._original_param_groups = self.optimizer.param_groups
+        original_steps = [g.get("step", 0) for g in self._original_param_groups]
 
-        return result
+        h2d_stream = torch.cuda.Stream()
+        d2h_stream = torch.cuda.Stream()
+        compute_stream = torch.cuda.current_stream()
+        n = len(self._chunks)
+
+        # 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)
+
+        # Fix step counters on the original param_groups (Muon only).
+        self._sync_step_counters(original_steps)
+        self._original_param_groups = None
+
+        return None
 
     def zero_grad(self, set_to_none: bool = True):
         self.optimizer.zero_grad(set_to_none=set_to_none)
@@ -126,7 +313,7 @@ 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)
 
     return optimizer
 
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index 1f2784bccf..caef6b2e37 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -178,7 +178,10 @@ 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,
     )
 
     # Set up the learning rate scheduler

From eac59e3bbbc944d3e87f77d3343353a41ad53661 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Wed, 22 Jul 2026 19:35:21 +0000
Subject: [PATCH 2/7] feat: add on/off config for chunked offload and stream
 overlap
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add  and  config options
to independently toggle per-layer chunking and stream-overlapped H2D/D2H.

When stream is disabled, the chunked step uses a simple sequential
move→step→move loop with no CUDA stream logic, making the code path clearer
and easier to debug.
---
 docs/scaling.md                               | 10 ++++-
 .../src/prime_rl/configs/trainer.py           |  6 +++
 src/prime_rl/trainer/optim.py                 | 43 +++++++++++++++----
 src/prime_rl/trainer/rl/train.py              |  2 +
 src/prime_rl/trainer/sft/train.py             |  2 +
 5 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 645e9135ff..eae6c81bc2 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -146,7 +146,15 @@ 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.
+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   # already the default
+optim_cpu_offload_stream = true    # already the default
+```
+
+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
 
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 11ecb92064..d2225c3a9a 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/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 = True
+    """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 = True
+    """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."""
 
diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py
index 2b0bff027c..623abfa142 100644
--- a/src/prime_rl/trainer/optim.py
+++ b/src/prime_rl/trainer/optim.py
@@ -42,9 +42,11 @@ def __init__(
         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
@@ -221,11 +223,34 @@ def step(self, closure=None):
             self._move_states("cpu")
             return result
 
-        # ---- Chunked step with stream-overlapped H2D / D2H ---- #
-
+        # Chunked step.
         self._original_param_groups = self.optimizer.param_groups
         original_steps = [g.get("step", 0) for g in self._original_param_groups]
 
+        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()
@@ -254,12 +279,6 @@ def step(self, closure=None):
         compute_stream.wait_stream(h2d_stream)
         compute_stream.wait_stream(d2h_stream)
 
-        # Fix step counters on the original param_groups (Muon only).
-        self._sync_step_counters(original_steps)
-        self._original_param_groups = None
-
-        return None
-
     def zero_grad(self, set_to_none: bool = True):
         self.optimizer.zero_grad(set_to_none=set_to_none)
 
@@ -301,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
@@ -313,7 +334,11 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        return CPUOffloadOptimizer(optimizer, named_params=named_params)
+        return CPUOffloadOptimizer(
+            optimizer,
+            named_params=named_params if cpu_offload_chunked else None,
+            stream=cpu_offload_stream,
+        )
 
     return optimizer
 
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 29a178a76f..39cca469a7 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/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:
diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py
index caef6b2e37..0d23275209 100644
--- a/src/prime_rl/trainer/sft/train.py
+++ b/src/prime_rl/trainer/sft/train.py
@@ -182,6 +182,8 @@ def train(config: SFTConfig):
         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

From 2f6cb2aaac81b99340879d971508e4e65fd4317a Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Wed, 22 Jul 2026 20:53:46 +0000
Subject: [PATCH 3/7] chore: default optim_cpu_offload_chunked and
 optim_cpu_offload_stream to false

---
 docs/scaling.md                                           | 4 ++--
 packages/prime-rl-configs/src/prime_rl/configs/trainer.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index eae6c81bc2..a04758ef6b 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -150,8 +150,8 @@ The optimizer step is performed per-transformer-layer ("chunked"): each layer's
 
 ```toml
 [trainer.model]
-optim_cpu_offload_chunked = true   # already the default
-optim_cpu_offload_stream = true    # already the default
+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.
diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index d2225c3a9a..4ea29cbfb6 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -171,10 +171,10 @@ 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 = True
+    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 = True
+    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

From 66e85d5c83e32eb00137f6b5399805525b9c9fa7 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Thu, 23 Jul 2026 04:22:27 +0000
Subject: [PATCH 4/7] fix: address bugbot review issues

1. Async D2H checkpoint corruption (HIGH): Add torch.cuda.synchronize()
   after _move_states('cpu') in state_dict() and in ckpt.py's
   AppState.state_dict() to ensure async pinned D2H copies complete
   before CPU tensors are read.

2. Incomplete stream dependency chain (MEDIUM): Replace wait_stream-only
   finish in _step_chunked_streamed with torch.cuda.synchronize() to
   block CPU. Add h2d_stream.wait_stream(d2h_stream) before prefetching
   chunk i+1 to prevent racing with the previous chunk's D2H into the
   same pinned CPU buffers.

3. First step OOM (HIGH): The init step now uses the same chunked path
   instead of a full optimizer.step() that materializes all states on
   GPU at once. States are created per-chunk and immediately evicted.

4. Defaults disagree (MEDIUM): Align setup_optimizer defaults for
   cpu_offload_chunked and cpu_offload_stream with config defaults
   (both False).

Also add torch.cuda.synchronize() at the end of _step_chunked (no-stream
path) to ensure async D2H completes before returning.
---
 docs/scaling.md               |  2 +-
 src/prime_rl/trainer/ckpt.py  |  1 +
 src/prime_rl/trainer/optim.py | 37 +++++++++++++++++++++--------------
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index a04758ef6b..2c7620f780 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -154,7 +154,7 @@ 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.
+Both default to `false`. Set `optim_cpu_offload_chunked = true` to enable per-layer chunking. Set `optim_cpu_offload_stream = true` to additionally overlap H2D/D2H transfers with optimizer compute on dedicated CUDA streams. When chunked is enabled but stream is not, the step uses a simple sequential move→step→move loop.
 
 ### LM Head Chunking
 
diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py
index 83728c8d9a..c4c92da7ed 100644
--- a/src/prime_rl/trainer/ckpt.py
+++ b/src/prime_rl/trainer/ckpt.py
@@ -109,6 +109,7 @@ def state_dict(self) -> dict[str, Any]:
             if isinstance(opt, CPUOffloadOptimizer):
                 opt._move_states("cpu")
         if has_cpu_offload:
+            torch.cuda.synchronize()
             gc.collect()
             torch.cuda.empty_cache()
 
diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py
index 623abfa142..218d80ba5a 100644
--- a/src/prime_rl/trainer/optim.py
+++ b/src/prime_rl/trainer/optim.py
@@ -209,21 +209,21 @@ def _sync_step_counters(self, original_steps: list):
     # ------------------------------------------------------------------ #
 
     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")
-            self._initialized = True
-            return result
-
-        # Non-chunked fallback (no named_params provided).
+        # Non-chunked path (no named_params provided).
         if self._chunks is None or len(self._chunks) <= 1:
+            if not self._initialized:
+                result = self.optimizer.step(closure)
+                self._move_states("cpu")
+                torch.cuda.synchronize()
+                self._initialized = True
+                return result
             self._move_states("cuda")
             result = self.optimizer.step(closure)
             self._move_states("cpu")
             return result
 
-        # Chunked step.
+        # Chunked path (also used for the first step to avoid materializing
+        # all optimizer states on GPU at once during initialization).
         self._original_param_groups = self.optimizer.param_groups
         original_steps = [g.get("step", 0) for g in self._original_param_groups]
 
@@ -234,6 +234,10 @@ def step(self, closure=None):
 
         self._sync_step_counters(original_steps)
         self._original_param_groups = None
+
+        if not self._initialized:
+            self._initialized = True
+
         return None
 
     def _step_chunked(self, closure=None):
@@ -242,6 +246,7 @@ def _step_chunked(self, closure=None):
             self._move_chunk_states(i, "cuda")
             self._step_chunk(i, closure)
             self._move_chunk_states(i, "cpu")
+        torch.cuda.synchronize()
 
     def _step_chunked_streamed(self, closure=None):
         """Per-layer optimizer step with stream-overlapped H2D / D2H.
@@ -264,8 +269,10 @@ def _step_chunked_streamed(self, closure=None):
             compute_stream.wait_stream(h2d_stream)
 
             # Prefetch H2D for chunk i+1 so it overlaps with step i on the
-            # compute stream.
+            # compute stream.  Wait for the previous chunk's D2H to finish
+            # first so we don't overwrite pinned CPU buffers still in flight.
             if i + 1 < n:
+                h2d_stream.wait_stream(d2h_stream)
                 self._move_chunk_states(i + 1, "cuda", h2d_stream)
 
             # Run optimizer step for chunk i.
@@ -275,9 +282,8 @@ def _step_chunked_streamed(self, closure=None):
             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)
+        # Block the CPU until all async transfers are complete.
+        torch.cuda.synchronize()
 
     def zero_grad(self, set_to_none: bool = True):
         self.optimizer.zero_grad(set_to_none=set_to_none)
@@ -290,6 +296,7 @@ def state_dict(self):
         sd = self.optimizer.state_dict()
         if self._initialized:
             self._move_states("cpu")
+            torch.cuda.synchronize()
         return sd
 
     def load_state_dict(self, state_dict):
@@ -320,8 +327,8 @@ def setup_optimizer(
     parallel_dims: ParallelDims,
     lora: bool = False,
     cpu_offload: bool = False,
-    cpu_offload_chunked: bool = True,
-    cpu_offload_stream: bool = True,
+    cpu_offload_chunked: bool = False,
+    cpu_offload_stream: bool = False,
 ) -> Optimizer | CPUOffloadOptimizer:
     if lora:
         # Wait for run 0 to be created in the multi run manager

From 053beda4d5b5e0e8ac36d7f46b3856759a45c071 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Thu, 23 Jul 2026 04:47:04 +0000
Subject: [PATCH 5/7] fix: always sync step counters in _sync_step_counters

Remove the 'if step in group' guard so _sync_step_counters always writes
group['step'] = orig_step + 1. This handles optimizers that lazily create
the step key on per-chunk copies (e.g. Muon) where the original groups
might not have had the key set yet.
---
 src/prime_rl/trainer/optim.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py
index 218d80ba5a..9625c74dc1 100644
--- a/src/prime_rl/trainer/optim.py
+++ b/src/prime_rl/trainer/optim.py
@@ -201,8 +201,7 @@ def _sync_step_counters(self, original_steps: list):
         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
+            group["step"] = orig_step + 1
 
     # ------------------------------------------------------------------ #
     # Public API

From 8350424d04af99b741d2ac99932082915e7be3a7 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Thu, 23 Jul 2026 04:54:18 +0000
Subject: [PATCH 6/7] fix: evaluate closure only once per chunked step (not per
 chunk)

Pass closure only to the first chunk's optimizer.step() and None to the
rest, so loss/grad recomputation happens once per training step rather than
once per layer.
---
 src/prime_rl/trainer/optim.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py
index 9625c74dc1..5c6ea65567 100644
--- a/src/prime_rl/trainer/optim.py
+++ b/src/prime_rl/trainer/optim.py
@@ -243,7 +243,7 @@ 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._step_chunk(i, closure if i == 0 else None)
             self._move_chunk_states(i, "cpu")
         torch.cuda.synchronize()
 
@@ -275,7 +275,9 @@ def _step_chunked_streamed(self, closure=None):
                 self._move_chunk_states(i + 1, "cuda", h2d_stream)
 
             # Run optimizer step for chunk i.
-            self._step_chunk(i, closure)
+            # Only evaluate the closure once (first chunk) to avoid recomputing
+            # loss/grads per chunk — the optimizer contract expects one evaluation per step.
+            self._step_chunk(i, closure if i == 0 else None)
 
             # Evict chunk i's states back to CPU (overlaps with step i+1).
             d2h_stream.wait_stream(compute_stream)

From 150b5d3e3722bdc6d72005bc49cc808dfd712b2d Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Fri, 24 Jul 2026 08:36:22 +0000
Subject: [PATCH 7/7] refactor: clean up CPUOffloadOptimizer for readability

- Remove redundant section separators and verbose docstrings
- Inline _build_chunk_param_groups into _step_chunk
- Inline _chunk_param_ids into _move_chunk_states
- Simplify _move_states to iterate values() directly
- Shorten comments to be targeted, not narrative
---
 src/prime_rl/trainer/optim.py | 182 ++++++++--------------------------
 1 file changed, 43 insertions(+), 139 deletions(-)

diff --git a/src/prime_rl/trainer/optim.py b/src/prime_rl/trainer/optim.py
index 5c6ea65567..d402a7cd87 100644
--- a/src/prime_rl/trainer/optim.py
+++ b/src/prime_rl/trainer/optim.py
@@ -16,25 +16,16 @@
 
 
 class CPUOffloadOptimizer:
-    """Wraps an optimizer to keep states on CPU, moving to GPU only for step().
+    """Keeps optimizer states on CPU, moving them to GPU only for step().
 
-    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.
+    Weights stay on GPU (unlike FSDP CPUOffload). With activation checkpointing,
+    activations and optimizer states are never on GPU at the same time, so peak
+    memory is 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.
+    ("chunked") instead of all-at-once, bounding peak GPU optimizer-state memory to
+    one layer's worth. H2D/D2H transfers can optionally overlap with compute on
+    dedicated CUDA streams (``stream=True``).
     """
 
     def __init__(
@@ -48,131 +39,74 @@ def __init__(
         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
+        """Group optimizer params by transformer layer; non-layer params go last."""
+        param_layer = {id(p): self._extract_layer_idx(n) for n, p in named_params}
         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)
+            for p in group["params"]:
+                idx = param_layer.get(id(p))
+                if idx is not None:
+                    by_layer.setdefault(idx, []).append(p)
                 else:
-                    misc.append(param)
-
+                    misc.append(p)
         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
-    # ------------------------------------------------------------------ #
+        return chunks or [misc]
 
     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 **all** optimizer states to *device* (used by state_dict / first step)."""
-        for p in self.optimizer.state:
-            state = self.optimizer.state[p]
+        """Move all optimizer states to *device*."""
+        for state in self.optimizer.state.values():
             for k, v in state.items():
                 if isinstance(v, DTensor):
-                    state[k] = self._move_dtensor(v, device)
+                    new_local = self._move_tensor(v._local_tensor, device)
+                    new_dtensor = copy.copy(v)
+                    new_dtensor._local_tensor = new_local
+                    state[k] = new_dtensor
                 elif isinstance(v, torch.Tensor):
                     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:
+        """Move optimizer states for a single chunk to/from *device*."""
+        chunk_ids = {id(p) for p in self._chunks[chunk_idx]}
+        ctx = torch.cuda.stream(stream) if stream is not None else torch.cuda.default_stream()
+        with ctx:
+            for p, state in self.optimizer.state.items():
                 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)
+                        new_local = self._move_tensor(v._local_tensor, device)
+                        new_dtensor = copy.copy(v)
+                        new_dtensor._local_tensor = new_local
+                        state[k] = new_dtensor
                     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] = []
+    def _step_chunk(self, chunk_idx: int, closure=None):
+        """Run optimizer.step() for a single chunk by temporarily swapping param_groups."""
+        chunk_ids = {id(p) for p in self._chunks[chunk_idx]}
+        chunk_groups = []
         for group in self._original_param_groups:
             filtered = [p for p in group["params"] if id(p) in chunk_ids]
             if not filtered:
@@ -180,11 +114,6 @@ def _build_chunk_param_groups(self, chunk_idx: int) -> list[dict]:
             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
@@ -193,22 +122,16 @@ def _step_chunk(self, chunk_idx: int, closure=None):
     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).
+        Muon stores a per-group ``step`` that ``step()`` increments. Because we swap
+        in per-chunk copies, the original groups never see the increment. Standard
+        optimizers (AdamW, SGD, SignSGD) keep step in ``state[p]['step']`` which is
+        per-parameter and already correct.
         """
         for group, orig_step in zip(self._original_param_groups, original_steps):
             group["step"] = orig_step + 1
 
-    # ------------------------------------------------------------------ #
-    # Public API
-    # ------------------------------------------------------------------ #
-
     def step(self, closure=None):
-        # Non-chunked path (no named_params provided).
+        # Non-chunked path.
         if self._chunks is None or len(self._chunks) <= 1:
             if not self._initialized:
                 result = self.optimizer.step(closure)
@@ -233,14 +156,11 @@ def step(self, closure=None):
 
         self._sync_step_counters(original_steps)
         self._original_param_groups = None
-
-        if not self._initialized:
-            self._initialized = True
-
+        self._initialized = True
         return None
 
     def _step_chunked(self, closure=None):
-        """Per-layer optimizer step without stream overlap — simple sequential loop."""
+        """Sequential per-layer step without stream overlap."""
         for i in range(len(self._chunks)):
             self._move_chunk_states(i, "cuda")
             self._step_chunk(i, closure if i == 0 else None)
@@ -248,49 +168,33 @@ def _step_chunked(self, closure=None):
         torch.cuda.synchronize()
 
     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.
-        """
+        """Per-layer step with stream-overlapped H2D / D2H."""
         h2d_stream = torch.cuda.Stream()
         d2h_stream = torch.cuda.Stream()
         compute_stream = torch.cuda.current_stream()
         n = len(self._chunks)
 
-        # 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.  Wait for the previous chunk's D2H to finish
-            # first so we don't overwrite pinned CPU buffers still in flight.
             if i + 1 < n:
+                # Wait for previous D2H before reusing pinned CPU buffers.
                 h2d_stream.wait_stream(d2h_stream)
                 self._move_chunk_states(i + 1, "cuda", h2d_stream)
 
-            # Run optimizer step for chunk i.
-            # Only evaluate the closure once (first chunk) to avoid recomputing
-            # loss/grads per chunk — the optimizer contract expects one evaluation per step.
             self._step_chunk(i, closure if i == 0 else None)
 
-            # 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)
 
-        # Block the CPU until all async transfers are complete.
         torch.cuda.synchronize()
 
     def zero_grad(self, set_to_none: bool = True):
         self.optimizer.zero_grad(set_to_none=set_to_none)
 
     def state_dict(self):
-        # Move to GPU temporarily for consistent state dict
         if self._initialized:
             self._move_states("cuda")
             torch.cuda.synchronize()