From 4309f95a601d15e9cc466747f09831f8f1bcbb86 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 10 Oct 2026 21:10:53 +0000
Subject: [PATCH 01/12] perf(optim): overlap optimizer state offload copies
 with the AdamW step

---
 src/prime_rl/trainer/optim/state_offload.py | 108 +++++++++++++++++++-
 1 file changed, 107 insertions(+), 1 deletion(-)

diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py
index 15a7f1eca2c..4500fefc6a7 100644
--- a/src/prime_rl/trainer/optim/state_offload.py
+++ b/src/prime_rl/trainer/optim/state_offload.py
@@ -2,10 +2,27 @@
 
 import torch
 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
 
+# 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
+
+
+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 CPUOffloadOptimizer(OffloadOptimizer):
     """Wraps an optimizer to keep states on CPU, moving to GPU only for step().
@@ -20,6 +37,7 @@ def __init__(self, optimizer: Optimizer, pin_memory: bool = True):
         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] | None = None
 
     def _offload_to_cpu(self, param: torch.Tensor, key: str, src: torch.Tensor) -> torch.Tensor:
         if not self.pin_memory:
@@ -62,6 +80,9 @@ def step(self, closure=None):
             self._initialized = True
             return result
 
+        if closure is None and self.pin_memory and isinstance(self.optimizer, AdamW):
+            return self._chunked_step()
+
         # Move states to GPU
         self._move_states("cuda")
 
@@ -73,6 +94,91 @@ 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 _step_params(self, params: set[torch.Tensor]):
+        groups = self.optimizer.param_groups
+        all_params = [group["params"] for group in groups]
+        for group in groups:
+            group["params"] = [param for param in group["params"] if param in params]
+        try:
+            return self.optimizer.step()
+        finally:
+            for group, group_params in zip(groups, all_params):
+                group["params"] = group_params
+
+    def _chunked_step(self):
+        """Step chunk by chunk through two staging buffers: while chunk i updates on the compute stream,
+        chunk i+1 loads on one side stream and chunk i-1 writes back on another.
+
+        The step counters stay on CPU, as in plain torch AdamW, so the update never syncs the device."""
+        if self._streams is None:
+            self._streams = (torch.cuda.Stream(), torch.cuda.Stream())
+        h2d, d2h = self._streams
+        compute = torch.cuda.current_stream()
+        chunks, buffer_bytes = self._plan_chunks()
+        state = self.optimizer.state
+        # Allocated on the compute stream, so the caching allocator reuses the same blocks every step.
+        buffers = [torch.empty(buffer_bytes, dtype=torch.uint8, device="cuda") for _ in range(2)]
+        h2d.wait_stream(compute)
+        written_back: list[torch.cuda.Event | None] = [None, None]
+        cpu_values: list[dict[tuple[torch.Tensor, str], torch.Tensor]] = [{}, {}]
+
+        def load(i: int) -> torch.cuda.Event:
+            buffer, slot = buffers[i % 2], cpu_values[i % 2]
+            with torch.cuda.stream(h2d):
+                if written_back[i % 2] is not None:
+                    h2d.wait_event(written_back[i % 2])
+                slot.clear()
+                for param, key, offset in 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 h2d.record_event()
+
+        result = None
+        loaded = load(0) if chunks else None
+        for i, (params, _) in enumerate(chunks):
+            next_loaded = load(i + 1) if i + 1 < len(chunks) else None
+            compute.wait_event(loaded)
+            result = self._step_params(params)
+            d2h.wait_stream(compute)
+            with torch.cuda.stream(d2h):
+                for (param, key), cpu_value in cpu_values[i % 2].items():
+                    _local(cpu_value).copy_(_local(state[param][key]), non_blocking=True)
+                    state[param][key] = cpu_value
+                written_back[i % 2] = d2h.record_event()
+            loaded = next_loaded
+        if not chunks:
+            result = self.optimizer.step()
+        compute.wait_stream(d2h)
+        torch.cuda.synchronize()
+        return result
+
     def zero_grad(self, set_to_none: bool = True):
         self.optimizer.zero_grad(set_to_none=set_to_none)
 

From 37cba282ca007a06b3449e9a928208bb908fc18d Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 11 Oct 2026 15:43:53 +0000
Subject: [PATCH 02/12] fix(optim): step parameters that have no offloaded
 state yet through the full path

---
 src/prime_rl/trainer/optim/state_offload.py | 21 ++++++++++++---------
 tests/unit/train/test_state_offload.py      | 19 +++++++++++++++++++
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py
index 4500fefc6a7..329b193c4ad 100644
--- a/src/prime_rl/trainer/optim/state_offload.py
+++ b/src/prime_rl/trainer/optim/state_offload.py
@@ -80,7 +80,13 @@ def step(self, closure=None):
             self._initialized = True
             return result
 
-        if closure is None and self.pin_memory and isinstance(self.optimizer, AdamW):
+        # A parameter without state yet (no grad on any earlier step) is initialized by the full path below.
+        if (
+            closure is None
+            and self.pin_memory
+            and 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
@@ -122,7 +128,7 @@ def _step_params(self, params: set[torch.Tensor]):
         for group in groups:
             group["params"] = [param for param in group["params"] if param in params]
         try:
-            return self.optimizer.step()
+            self.optimizer.step()
         finally:
             for group, group_params in zip(groups, all_params):
                 group["params"] = group_params
@@ -137,6 +143,8 @@ def _chunked_step(self):
         h2d, d2h = self._streams
         compute = torch.cuda.current_stream()
         chunks, buffer_bytes = self._plan_chunks()
+        if not chunks:
+            return
         state = self.optimizer.state
         # Allocated on the compute stream, so the caching allocator reuses the same blocks every step.
         buffers = [torch.empty(buffer_bytes, dtype=torch.uint8, device="cuda") for _ in range(2)]
@@ -160,12 +168,11 @@ def load(i: int) -> torch.cuda.Event:
                     state[param][key] = _with_local(cpu_value, gpu_local)
                 return h2d.record_event()
 
-        result = None
-        loaded = load(0) if chunks else None
+        loaded = load(0)
         for i, (params, _) in enumerate(chunks):
             next_loaded = load(i + 1) if i + 1 < len(chunks) else None
             compute.wait_event(loaded)
-            result = self._step_params(params)
+            self._step_params(params)
             d2h.wait_stream(compute)
             with torch.cuda.stream(d2h):
                 for (param, key), cpu_value in cpu_values[i % 2].items():
@@ -173,11 +180,7 @@ def load(i: int) -> torch.cuda.Event:
                     state[param][key] = cpu_value
                 written_back[i % 2] = d2h.record_event()
             loaded = next_loaded
-        if not chunks:
-            result = self.optimizer.step()
-        compute.wait_stream(d2h)
         torch.cuda.synchronize()
-        return result
 
     def zero_grad(self, set_to_none: bool = True):
         self.optimizer.zero_grad(set_to_none=set_to_none)
diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py
index 68ba4015623..55d2933fbd9 100644
--- a/tests/unit/train/test_state_offload.py
+++ b/tests/unit/train/test_state_offload.py
@@ -112,3 +112,22 @@ 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}"
+
+
+def test_parameter_without_grad_on_the_first_step_is_stepped_later() -> None:
+    def train(wrap: bool) -> list[torch.Tensor]:
+        model = make_model(0)
+        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)
+        if wrap:
+            optimizer = CPUOffloadOptimizer(optimizer)
+        torch.manual_seed(1)
+        for step in range(3):
+            hidden = model[0](torch.randn(16, 64, device="cuda"))
+            loss = hidden.square().mean() if step == 0 else model[2](hidden).square().mean()
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+        return [param.detach() for param in model.parameters()]
+
+    for reference, offloaded in zip(train(False), train(True)):
+        assert torch.equal(reference, offloaded)

From 9fc3f89a31845b6dd7472a314a8582a9608a5cb9 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Mon, 12 Oct 2026 18:28:03 +0000
Subject: [PATCH 03/12] refactor(optim): drop the always-true pin_memory and
 closure clauses of the chunked-step gate

---
 src/prime_rl/trainer/optim/state_offload.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py
index 329b193c4ad..31fd1f40634 100644
--- a/src/prime_rl/trainer/optim/state_offload.py
+++ b/src/prime_rl/trainer/optim/state_offload.py
@@ -32,16 +32,13 @@ 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):
         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] | None = None
 
     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)
@@ -81,11 +78,8 @@ def step(self, closure=None):
             return result
 
         # A parameter without state yet (no grad on any earlier step) is initialized by the full path below.
-        if (
-            closure is None
-            and self.pin_memory
-            and 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)
+        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()
 

From b9809f3aeb02c28f7da9820836a79481a125e418 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 01:12:38 +0000
Subject: [PATCH 04/12] perf(optim): overlap the offloaded AdamW step with the
 next forward

The chunked CPUOffloadOptimizer step returns before its chunks run. They
update on a side stream while the next forward runs, enqueued as the forward
reaches them (8 up front, then 4 ahead of the forward), so the CPU never fills
the CUDA launch queue. Each FSDP parameter group waits for the chunk holding
its parameters before its all-gather. The step's hyperparameters and gradients
are captured at step time, so the math is unchanged.

model.offload_overlap_forward (default true) turns it on for AdamW. Anything
else that reads the parameters calls wait_for_optimizer_step first: weight
broadcasts in the SFT and RL loops; checkpoint saves already wait through
prepare_checkpoint_save.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/scaling.md                               |   2 +-
 .../src/prime_rl/configs/trainer.py           |   3 +
 src/prime_rl/trainer/optim/__init__.py        |  10 +-
 src/prime_rl/trainer/optim/state_offload.py   | 257 ++++++++++++++----
 src/prime_rl/trainer/rl/train.py              |   4 +-
 src/prime_rl/trainer/train.py                 |   5 +-
 6 files changed, 224 insertions(+), 57 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 99877cd6db7..3e7f2cd5929 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -156,7 +156,7 @@ Activation offloading still applies to tensors saved by autograd, but tensors re
 
 ### 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, and by default (`model.offload_overlap_forward = true`) those copies and updates run alongside the next forward, each FSDP module waiting only for its own parameters. 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
 
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 c54aabc2266..149c038bf54 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -297,6 +297,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 = True
+    """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. The math is unchanged; the step's gradients stay allocated until their chunk has updated."""
+
     reshard_after_forward: bool = True
     """Reshard the model after each forward pass."""
 
diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py
index 08316ec71e0..7daf3c35535 100644
--- a/src/prime_rl/trainer/optim/__init__.py
+++ b/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]],
diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py
index 31fd1f40634..f25a14e416a 100644
--- a/src/prime_rl/trainer/optim/state_offload.py
+++ b/src/prime_rl/trainer/optim/state_offload.py
@@ -1,15 +1,24 @@
 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 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:
@@ -24,6 +33,110 @@ def _with_local(value: torch.Tensor, local: torch.Tensor) -> torch.Tensor:
     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):
     """Wraps an optimizer to keep states on CPU, moving to GPU only for step().
 
@@ -32,11 +145,81 @@ class CPUOffloadOptimizer(OffloadOptimizer):
     at the same time: peak memory becomes max(activations, opt_states) instead of sum.
     """
 
-    def __init__(self, optimizer: Optimizer):
+    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._initialized = False
         self._cpu_buffers: dict[tuple[torch.Tensor, str], torch.Tensor] = {}
-        self._streams: tuple[torch.cuda.Stream, torch.cuda.Stream] | None = None
+        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:
         buffer = self._cpu_buffers.get((param, key))
@@ -49,6 +232,7 @@ 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():
@@ -116,67 +300,34 @@ def _plan_chunks(self) -> tuple[list[tuple[set[torch.Tensor], list[tuple[torch.T
             chunks.append((params, entries))
         return chunks, max_used
 
-    def _step_params(self, params: set[torch.Tensor]):
-        groups = self.optimizer.param_groups
-        all_params = [group["params"] for group in groups]
-        for group in groups:
-            group["params"] = [param for param in group["params"] if param in params]
-        try:
-            self.optimizer.step()
-        finally:
-            for group, group_params in zip(groups, all_params):
-                group["params"] = group_params
-
     def _chunked_step(self):
-        """Step chunk by chunk through two staging buffers: while chunk i updates on the compute stream,
-        chunk i+1 loads on one side stream and chunk i-1 writes back on another.
+        """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())
-        h2d, d2h = self._streams
-        compute = torch.cuda.current_stream()
+            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
-        state = self.optimizer.state
-        # Allocated on the compute stream, so the caching allocator reuses the same blocks every step.
-        buffers = [torch.empty(buffer_bytes, dtype=torch.uint8, device="cuda") for _ in range(2)]
-        h2d.wait_stream(compute)
-        written_back: list[torch.cuda.Event | None] = [None, None]
-        cpu_values: list[dict[tuple[torch.Tensor, str], torch.Tensor]] = [{}, {}]
-
-        def load(i: int) -> torch.cuda.Event:
-            buffer, slot = buffers[i % 2], cpu_values[i % 2]
-            with torch.cuda.stream(h2d):
-                if written_back[i % 2] is not None:
-                    h2d.wait_event(written_back[i % 2])
-                slot.clear()
-                for param, key, offset in 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 h2d.record_event()
-
-        loaded = load(0)
+        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):
-            next_loaded = load(i + 1) if i + 1 < len(chunks) else None
-            compute.wait_event(loaded)
-            self._step_params(params)
-            d2h.wait_stream(compute)
-            with torch.cuda.stream(d2h):
-                for (param, key), cpu_value in cpu_values[i % 2].items():
-                    _local(cpu_value).copy_(_local(state[param][key]), non_blocking=True)
-                    state[param][key] = cpu_value
-                written_back[i % 2] = d2h.record_event()
-            loaded = next_loaded
-        torch.cuda.synchronize()
+            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):
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index d5cfbec4df7..63f5bf6f029 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -16,7 +16,7 @@
 import torch.distributed as dist
 from torch.profiler import profile, ProfilerActivity, record_function
 from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager
-from prime_rl.trainer.optim import setup_optimizer
+from prime_rl.trainer.optim import setup_optimizer, wait_for_optimizer_step
 from prime_rl.trainer.scheduler import setup_scheduler
 from prime_rl.configs.trainer import TrainerConfig
 from prime_rl.trainer.rl.data import (
@@ -185,6 +185,7 @@ def train(config: TrainerConfig):
         list(model.named_parameters()),
         parallel_dims,
         offload=config.model.offload,
+        offload_overlap_forward=config.model.offload_overlap_forward,
         model=model,
         full_offload_dtype_policy=(get_full_offload_dtype_policy(model) if config.model.offload == "full" else None),
     )
@@ -681,6 +682,7 @@ def train(config: TrainerConfig):
             # so a still-running kernel holding a cached block (e.g. the
             # optimizer step's tail) faults with an illegal memory access
             # once its block is freed under it.
+            wait_for_optimizer_step(optimizer)
             torch.cuda.synchronize()
             torch.cuda.empty_cache()
             weight_sender.broadcast(model, step=progress.step)
diff --git a/src/prime_rl/trainer/train.py b/src/prime_rl/trainer/train.py
index 693d3cf740c..a9f210d5228 100644
--- a/src/prime_rl/trainer/train.py
+++ b/src/prime_rl/trainer/train.py
@@ -22,7 +22,7 @@
 from prime_rl.trainer.lora import get_lora_state
 from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX
 from prime_rl.utils.logger import format_time, setup_logger
-from prime_rl.trainer.optim import setup_optimizer
+from prime_rl.trainer.optim import setup_optimizer, wait_for_optimizer_step
 from prime_rl.trainer.scheduler import setup_scheduler
 from prime_rl.trainer.model import (
     forward,
@@ -163,6 +163,7 @@ def train(config: SFTConfig | PretrainConfig):
         list(model.named_parameters()),
         parallel_dims,
         offload=config.model.offload,
+        offload_overlap_forward=config.model.offload_overlap_forward,
         model=model,
         full_offload_dtype_policy=(get_full_offload_dtype_policy(model) if config.model.offload == "full" else None),
     )
@@ -554,6 +555,7 @@ def is_online_eval_step(step: int) -> bool:
         if weight_sender is not None and not is_last_step and is_online_eval_step(progress.step):
             logger.info(f"Broadcasting weights at step {progress.step}")
             broadcast_start_time = time.perf_counter()
+            wait_for_optimizer_step(optimizer)
             weight_sender.broadcast(model, step=progress.step)
             broadcast_weights_time = time.perf_counter() - broadcast_start_time
 
@@ -690,6 +692,7 @@ def is_online_eval_step(step: int) -> bool:
     # run has no final epoch.
     if weight_sender is not None and not paused:
         logger.info("Broadcasting final weights")
+        wait_for_optimizer_step(optimizer)
         weight_sender.broadcast(model, step=progress.step)
 
     if gradient_manager is not None:

From 52139d9a9a5834114cf0c26b1b29af2c875803d7 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 02:06:03 +0000
Subject: [PATCH 05/12] perf(optim): chunked offloaded optimizer step for Muon
 via a subset-step interface

CPUOffloadOptimizer now chunks the step of any optimizer that can update a
subset of its parameters at a time with the math of a full step
(subset_step.SupportsSubsetStep). torch AdamW and dion's Muon get adapters;
Muon's units are dion's own batches, so its all-to-alls and compiled updates
are unchanged, and its param-group step counters advance once per step.
Forward overlap applies to every such optimizer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/trainer/optim/__init__.py      |   2 +-
 src/prime_rl/trainer/optim/state_offload.py | 121 +++++++++++--------
 src/prime_rl/trainer/optim/subset_step.py   | 124 ++++++++++++++++++++
 tests/unit/train/test_state_offload.py      |  26 ++++
 4 files changed, 221 insertions(+), 52 deletions(-)
 create mode 100644 src/prime_rl/trainer/optim/subset_step.py

diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py
index 7daf3c35535..8c28ae2c66f 100644
--- a/src/prime_rl/trainer/optim/__init__.py
+++ b/src/prime_rl/trainer/optim/__init__.py
@@ -95,7 +95,7 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        overlap_model = model if offload_overlap_forward and config.type == "adamw" else None
+        overlap_model = model if offload_overlap_forward else None
         return CPUOffloadOptimizer(optimizer, overlap_model=overlap_model), 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 f25a14e416a..ada63c402d3 100644
--- a/src/prime_rl/trainer/optim/state_offload.py
+++ b/src/prime_rl/trainer/optim/state_offload.py
@@ -1,13 +1,15 @@
 import copy
+import math
 
 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 AdamW, Optimizer
+from torch.optim import Optimizer
 
 from prime_rl.trainer.optim.base import OffloadOptimizer
+from prime_rl.trainer.optim.subset_step import SubsetStep, subset_step_for
 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.
@@ -33,24 +35,37 @@ def _with_local(value: torch.Tensor, local: torch.Tensor) -> torch.Tensor:
     return new_dtensor
 
 
+def _nominal_bytes(value: torch.Tensor) -> int:
+    """Size of a tensor's local shard, the same on every rank: an even split of each sharded mesh dim."""
+    numel = value.numel()
+    if isinstance(value, DTensor):
+        shards = math.prod(
+            value.device_mesh.size(dim) for dim, placement in enumerate(value.placements) if placement.is_shard()
+        )
+        numel = -(-numel // shards)
+    return numel * value.element_size()
+
+
 class _ChunkedStep:
-    """One chunked AdamW step. Chunks stream through two GPU staging buffers: while chunk i updates,
+    """One chunked optimizer 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
+    Chunks are enqueued on demand by ``advance``. ``subset_step`` captured the step's hyperparameters
+    and gradients 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]]]],
+        subset_step: SubsetStep,
+        chunks: list[tuple[list[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.subset_step = subset_step
         self.chunks = chunks
         self.h2d, self.d2h, self.update = h2d, d2h, update
         compute = torch.cuda.current_stream()
@@ -61,7 +76,6 @@ def __init__(
         # 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]] = [{}, {}]
@@ -89,23 +103,6 @@ def _load(self, i: int) -> torch.cuda.Event:
                 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
@@ -115,7 +112,7 @@ def advance(self, until: int) -> None:
             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.subset_step.step_params(params)
                 self.updated.append(self.update.record_event())
             self.d2h.wait_stream(self.update)
             with torch.cuda.stream(self.d2h):
@@ -146,27 +143,34 @@ class CPUOffloadOptimizer(OffloadOptimizer):
     """
 
     def __init__(self, optimizer: Optimizer, overlap_model: nn.Module | None = None):
-        """With ``overlap_model`` (FSDP-sharded), a chunked AdamW step returns before its chunks
+        """The step is chunked for optimizers that can update a subset of their parameters at a time
+        (``subset_step.SupportsSubsetStep``, e.g. AdamW and Muon); others move all states at once.
+
+        With ``overlap_model`` (FSDP-sharded), a chunked 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._subset_step = subset_step_for(optimizer)
         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] = {}
+        # Forward order of the FSDP groups.
+        self._group_order: dict[FSDPParamGroup, int] = {}
         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:
+        if overlap_model is not None and self._subset_step is not None:
             self._install_unshard_waits(overlap_model)
 
     def _install_unshard_waits(self, model: nn.Module) -> None:
-        fsdp_groups = {}
+        fsdp_groups, group_order = {}, {}
         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:
+                    group_order.setdefault(group, len(group_order))
                     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"]]
@@ -177,7 +181,7 @@ def _install_unshard_waits(self, model: nn.Module) -> None:
                 "the offloaded optimizer step will not overlap the next forward"
             )
             return
-        self._fsdp_groups = fsdp_groups
+        self._fsdp_groups, self._group_order = fsdp_groups, group_order
         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)
@@ -262,7 +266,7 @@ def step(self, closure=None):
             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(
+        if self._subset_step is not None 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()
@@ -278,23 +282,31 @@ 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:
+    def _plan_chunks(
+        self, units: list[list[torch.Tensor]]
+    ) -> tuple[list[tuple[list[torch.Tensor], list[tuple[torch.Tensor, str, int]]]], int]:
+        """Pack whole units into chunks; each streamed state gets an offset into a staging buffer.
+
+        Chunk boundaries follow ``_nominal_bytes``, so every rank cuts the same chunks and a step that
+        communicates runs the same collectives everywhere."""
+        chunks, params, entries, used, nominal, max_used = [], [], [], 0, 0, 0
+        for unit in units:
+            unit_entries, unit_bytes, unit_nominal = [], 0, 0
+            for param in unit:
+                for key, value in self.optimizer.state[param].items():
+                    if not isinstance(value, torch.Tensor) or value.ndim == 0:
+                        continue
+                    local = _local(value)
+                    unit_entries.append((param, key, unit_bytes))
+                    unit_bytes += -(-local.numel() * local.element_size() // ALIGN_BYTES) * ALIGN_BYTES
+                    unit_nominal += _nominal_bytes(value)
+            if params and nominal + unit_nominal > 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
+                params, entries, used, nominal = [], [], 0, 0
+            params.extend(unit)
+            entries.extend((param, key, used + offset) for param, key, offset in unit_entries)
+            used += unit_bytes
+            nominal += unit_nominal
             max_used = max(max_used, used)
         if params:
             chunks.append((params, entries))
@@ -302,23 +314,30 @@ def _plan_chunks(self) -> tuple[list[tuple[set[torch.Tensor], list[tuple[torch.T
 
     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 compute stream and the step waits for the GPU; with one they update on a side stream in
+        the forward order of their FSDP groups, enqueued as the next forward reaches them.
 
-        The step counters stay on CPU, as in plain torch AdamW, so the update never syncs the device."""
+        Scalar step counters stay where they are (on CPU for 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()
+        subset_step = self._subset_step.begin_subset_step()
+        units = subset_step.units
+        if self._fsdp_groups:
+            units = sorted(units, key=lambda unit: min(self._group_order[self._fsdp_groups[p]] for p in unit))
+        chunks, buffer_bytes = self._plan_chunks(units)
         if not chunks:
             return
         if not self._fsdp_groups:
-            step = _ChunkedStep(self.optimizer, chunks, buffer_bytes, h2d, d2h, torch.cuda.current_stream())
+            step = _ChunkedStep(
+                self.optimizer, subset_step, 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)
+        self._in_flight = _ChunkedStep(self.optimizer, subset_step, chunks, buffer_bytes, h2d, d2h, side)
         for i, (params, _) in enumerate(chunks):
             for param in params:
                 self._group_chunks[self._fsdp_groups[param]] = i
diff --git a/src/prime_rl/trainer/optim/subset_step.py b/src/prime_rl/trainer/optim/subset_step.py
new file mode 100644
index 00000000000..a7cb15dc8c5
--- /dev/null
+++ b/src/prime_rl/trainer/optim/subset_step.py
@@ -0,0 +1,124 @@
+"""Optimizer steps that update a subset of the parameters at a time.
+
+``CPUOffloadOptimizer`` streams offloaded optimizer state through the GPU in chunks, so it updates one
+chunk of parameters at a time. An optimizer opts in by implementing ``SupportsSubsetStep``; torch
+``AdamW`` and dion's ``Muon`` get adapters here (``subset_step_for``).
+
+The contract:
+
+- ``begin_subset_step()`` replaces ``step()`` for one step. It captures everything the step reads that
+  the trainer may change before the last partial update runs (group hyperparameters, gradients) and
+  advances step counters, so afterwards the param groups look as after ``step()``.
+- ``units`` partitions the parameters the step updates into the smallest sets that must update
+  together (e.g. a batch sharing one collective). A step that communicates must return the same
+  units, in the same order, on every rank it communicates with: decide from global shapes, never from
+  local shard sizes.
+- ``step_params(params)`` updates exactly ``params``, a union of whole units, bitwise as the full step
+  would. It reads each state tensor from ``optimizer.state`` at call time (the caller swaps them for
+  GPU copies), enqueues its work on the current stream without syncing the host, and is called at
+  most once per unit. Every rank calls it with the same units in the same order.
+- The streamed state is every tensor in ``optimizer.state[param]`` with at least one dimension (a
+  DTensor's local tensor); scalar step counters stay where they are. ``step_params`` never creates
+  state: the caller runs a plain ``step()`` while a parameter with a gradient has none.
+"""
+
+from collections.abc import Collection
+from typing import Protocol, runtime_checkable
+
+import dion
+import torch
+from dion.opt_utils import create_param_batches
+from torch.optim import AdamW, Optimizer
+
+
+class SubsetStep(Protocol):
+    """One optimizer step, run as several partial updates."""
+
+    units: list[list[torch.Tensor]]
+
+    def step_params(self, params: Collection[torch.Tensor]) -> None: ...
+
+
+@runtime_checkable
+class SupportsSubsetStep(Protocol):
+    def begin_subset_step(self) -> SubsetStep: ...
+
+
+class GroupSubsetStep:
+    """Partial updates of a torch optimizer whose ``step()`` updates exactly the parameters of its param
+    groups that have a gradient: each update narrows every group to the requested parameters and
+    reinstates the hyperparameters and gradients captured at step time around one ``step()``."""
+
+    def __init__(self, optimizer: Optimizer, units: list[list[torch.Tensor]]):
+        self.optimizer = optimizer
+        self.units = units
+        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 unit in units for param in unit}
+
+    def step_params(self, params: Collection[torch.Tensor]) -> None:
+        params = set(params)
+        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.pop(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)
+
+
+class AdamWSubsetStep:
+    """torch AdamW updates every parameter on its own and keeps its step counters in the state."""
+
+    def __init__(self, optimizer: AdamW):
+        self.optimizer = optimizer
+
+    def begin_subset_step(self) -> SubsetStep:
+        groups = self.optimizer.param_groups
+        units = [[param] for group in groups for param in group["params"] if param.grad is not None]
+        return GroupSubsetStep(self.optimizer, units)
+
+
+class DionMuonSubsetStep:
+    """dion's Muon: a unit is one of dion's own Muon batches (matrices of one shape that share an
+    all-to-all), or a whole element-wise group, whose compiled foreach update must keep seeing the
+    whole group. dion advances each group's step counter on every ``step()`` call: the captured
+    counter is reinstated before each partial update, which advances it to this step's value."""
+
+    def __init__(self, optimizer: dion.Muon):
+        self.optimizer = optimizer
+
+    def begin_subset_step(self) -> SubsetStep:
+        optimizer = self.optimizer
+        units = []
+        for group in optimizer.param_groups:
+            params = [param for param in group["params"] if param.grad is not None]
+            if not params:
+                continue
+            if group["algorithm"] == "muon":
+                batch_size = optimizer._get_group_mesh_info(group)[1]
+                units.extend(create_param_batches(params, batch_size, matrix_partitions=optimizer._matrix_partitions))
+            else:
+                units.append(params)
+        subset_step = GroupSubsetStep(optimizer, units)
+        for group in optimizer.param_groups:
+            group["step"] += 1
+        return subset_step
+
+
+def subset_step_for(optimizer: Optimizer) -> SupportsSubsetStep | None:
+    """The optimizer itself if it implements partial steps, an adapter for known optimizers, else None."""
+    if isinstance(optimizer, SupportsSubsetStep):
+        return optimizer
+    if isinstance(optimizer, AdamW):
+        return AdamWSubsetStep(optimizer)
+    if isinstance(optimizer, dion.Muon):
+        return DionMuonSubsetStep(optimizer)
+    return None
diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py
index f1a5d1bf95c..ecf9fd2d70f 100644
--- a/tests/unit/train/test_state_offload.py
+++ b/tests/unit/train/test_state_offload.py
@@ -8,6 +8,7 @@
 from torch.optim import Optimizer
 
 from prime_rl.trainer.ckpt import AppState
+from prime_rl.trainer.optim import state_offload
 from prime_rl.trainer.optim.state_offload import CPUOffloadOptimizer
 
 pytestmark = [pytest.mark.gpu]
@@ -159,3 +160,28 @@ def setup():
     assert [group["step"] for group in resumed_optimizer.param_groups] == [3, 3]
     run_steps(resumed_model, resumed_optimizer, num_steps=2, seed=2)
     assert_params_equal(model, resumed_model)
+
+
+@pytest.mark.parametrize("chunk_bytes", [1 << 30, 1 << 12])
+def test_chunked_muon_matches_plain_muon(monkeypatch, chunk_bytes: int) -> None:
+    """dion's Muon steps chunk by chunk (a Muon batch or a whole AdamW group per unit) and advances its
+    param-group step counters once per step."""
+    monkeypatch.setattr(state_offload, "CHUNK_BYTES", chunk_bytes)
+
+    def train(wrap: bool) -> tuple[torch.nn.Module, Optimizer]:
+        model = make_model(0)
+        groups = [
+            dict(params=[p for p in model.parameters() if p.ndim == 2], algorithm="muon"),
+            dict(params=[p for p in model.parameters() if p.ndim == 1], algorithm="adamw"),
+        ]
+        optimizer = dion.Muon(groups, lr=1e-2, fsdp_mesh_dim=0, world_mesh=None, adjust_lr="rms_norm")
+        if wrap:
+            optimizer = CPUOffloadOptimizer(optimizer)
+        run_steps(model, optimizer, num_steps=4, seed=1)
+        return model, optimizer
+
+    reference_model, _ = train(False)
+    offloaded_model, offloaded_optimizer = train(True)
+    assert_params_equal(reference_model, offloaded_model)
+    assert_states_offloaded(offloaded_optimizer)
+    assert [group["step"] for group in offloaded_optimizer.param_groups] == [4, 4]

From 747aff432574906a10a4cfb9376898ea6fbbc917 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 02:52:37 +0000
Subject: [PATCH 06/12] perf(optim): keep forward overlap to AdamW; document
 chunked Muon offload

dion's Muon step syncs the host and launches hundreds of kernels per chunk,
so running it alongside the forward measured slower than the synchronous
chunked step (40.1k vs 41.2k tok/s on Qwen3-30B-A3B).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/scaling.md                        | 2 +-
 src/prime_rl/trainer/optim/__init__.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 3e7f2cd5929..163a672cfef 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -156,7 +156,7 @@ Activation offloading still applies to tensors saved by autograd, but tensors re
 
 ### Optimizer Offloading
 
-State-only optimizer offload is enabled by default (`model.offload = "optimizer"`). With AdamW, the step streams the states through the GPU in chunks, and by default (`model.offload_overlap_forward = true`) those copies and updates run alongside the next forward, each FSDP module waiting only for its own parameters. 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 and Muon, the step streams the states through the GPU in chunks, copying one chunk in and another out while a third updates. With AdamW, by default (`model.offload_overlap_forward = true`), those copies and updates also run alongside the next forward, each FSDP module waiting only for its own parameters. 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
 
diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py
index 8c28ae2c66f..a5606da4c54 100644
--- a/src/prime_rl/trainer/optim/__init__.py
+++ b/src/prime_rl/trainer/optim/__init__.py
@@ -95,7 +95,9 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        overlap_model = model if offload_overlap_forward else None
+        # dion's Muon step syncs the host and launches hundreds of kernels per chunk; run alongside the
+        # forward, that delays the forward by more than the copies it hides, so Muon steps synchronously.
+        overlap_model = model if offload_overlap_forward and config.type == "adamw" else None
         return CPUOffloadOptimizer(optimizer, overlap_model=overlap_model), None
 
     return optimizer, None

From 5b384c683181c41e3c04468cd1f6eaf1f4d4a216 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 04:21:41 +0000
Subject: [PATCH 07/12] feat(optim): native Muon, replacing dion
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A prime-rl Muon in trainer/optim/muon.py, numerically dion's Muon (with dion's Newton-Schulz
plugged in, every parameter is bitwise equal to dion over 20 steps on 8 GPUs), without dion:

- Expert stacks sharded only on the expert dim orthogonalize in place, batched.
- FSDP-sharded matrices gather onto owner ranks with uneven all-to-alls (no padding), owners
  balanced by Newton-Schulz FLOPs, in layer-ordered buckets; bucket k+1 gathers while bucket k
  orthogonalizes, on Muon's own NCCL communicators (4 CTAs).
- Newton-Schulz in bf16 on QuACK's symmetric GEMM with fused epilogues, batched per shape;
  tall matrices iterate on XᵀX without transposing; schedules are data (default, deepseek_v4,
  polar_express, step3p5).
- Element-wise updates are compiled once per unit; no host syncs in the step.
- Implements the subset-step interface for chunked optimizer-state offload.
- Same state keys and param groups as dion's Muon: its checkpoints resume as they are.

Adds optim.update_rms and optim.newton_schulz. Removes the dion dependency.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/scaling.md                               |  11 +
 .../src/prime_rl/configs/trainer.py           |  12 +-
 pyproject.toml                                |   3 -
 src/prime_rl/trainer/ckpt.py                  |   2 +-
 src/prime_rl/trainer/optim/__init__.py        |  97 +--
 src/prime_rl/trainer/optim/muon.py            | 699 ++++++++++++++++++
 tests/unit/train/test_muon.py                 | 129 ++++
 tests/unit/train/test_state_offload.py        |   6 +-
 uv.lock                                       |  14 -
 9 files changed, 866 insertions(+), 107 deletions(-)
 create mode 100644 src/prime_rl/trainer/optim/muon.py
 create mode 100644 tests/unit/train/test_muon.py

diff --git a/docs/scaling.md b/docs/scaling.md
index 99877cd6db7..dda2aba83bb 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -15,6 +15,7 @@ This page covers how to scale `prime-rl` from a single GPU to a 1000-GPU cluster
   - [Context Parallelism](#context-parallelism)
   - [Activation Checkpointing and Offloading](#activation-checkpointing-and-offloading)
   - [Optimizer Offloading](#optimizer-offloading)
+  - [Muon](#muon)
   - [LM Head Chunking](#lm-head-chunking)
 - [Memory-Tight Recipe](#memory-tight-recipe)
 - [Max Trainer Throughput Recipe](#max-trainer-throughput-recipe)
@@ -158,6 +159,16 @@ Activation offloading still applies to tensors saved by autograd, but tensors re
 
 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).
 
+### Muon
+
+`optim.type = "muon"` updates matrices with the orthogonalized (Newton-Schulz) Nesterov momentum and the remaining parameters with AdamW. Orthogonalization needs whole matrices, and the optimizer gets them without any setting:
+
+- Expert stacks sharded only along the expert dim (EP, or FSDP over the experts) already hold whole matrices on every rank and are orthogonalized in place, batched over the experts.
+- Matrices that FSDP shards along a row or column dim are gathered onto one rank of their shard group with uneven all-to-alls, orthogonalized there and scattered back. Ranks own matrices by Newton-Schulz FLOPs, matrices move in buckets of a few layers, and the next bucket gathers while the current one is orthogonalized. The all-to-alls run on Muon's own NCCL communicators, next to FSDP's.
+- Newton-Schulz runs in bf16, batched over all matrices of one shape, on QuACK's symmetric GEMM where available (`quack` extra, SM80+).
+
+Packed parameters (runtime fusions) are orthogonalized per logical matrix. `optim.newton_schulz` picks the iteration schedule and `optim.update_rms` the RMS of each update (0.2, which matches AdamW's learning rates).
+
 ### 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 `8192` for RL training:
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 c54aabc2266..080c346f09a 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -466,13 +466,19 @@ class MuonConfig(BaseOptimizerConfig):
     """Nesterov momentum for the Muon updates."""
 
     betas1: float = Field(0.9, ge=0)
-    """β1 for the AdamW/Lion sub-optimizer used on non-Muon params."""
+    """β1 for the AdamW sub-optimizer used on non-Muon params."""
 
     betas2: float = Field(0.95, ge=0)
-    """β2 for the AdamW/Lion sub-optimizer used on non-Muon params."""
+    """β2 for the AdamW sub-optimizer used on non-Muon params."""
 
     eps: float = Field(1e-8, gt=0)
-    """Epsilon for the AdamW sub-optimizer used on non-Muon params."""
+    """Epsilon for the AdamW sub-optimizer used on non-Muon params, and for the Newton-Schulz normalization."""
+
+    update_rms: float = Field(0.2, gt=0)
+    """RMS of each Muon update matrix, so Muon reuses the AdamW learning rate. DeepSeek-V4.1 uses 0.18."""
+
+    newton_schulz: Literal["default", "deepseek_v4", "polar_express", "step3p5"] = "default"
+    """Newton-Schulz schedule. ``default``: 5 quintic steps with per-step coefficients. ``deepseek_v4``: 8 steps of (3.4445, -4.7750, 2.0315), then 2 steps of (2, -1.5, 0.5) that pin the singular values at 1. ``polar_express``: 5 steps of Polar Express (arXiv 2505.16932) in bf16, with its 1.01 safety factor. ``step3p5``: 6 Polar Express steps in fp16, as in Step 3.5 Flash."""
 
 
 class SignSGDConfig(BaseOptimizerConfig):
diff --git a/pyproject.toml b/pyproject.toml
index 8e6a33a70e0..5394d8289c4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -60,7 +60,6 @@ gpu = [
     "vllm>=0.32.0 ; sys_platform == 'linux'",
     "mooncake-transfer-engine>=0.3.10.post2 ; sys_platform == 'linux'",
     "torchtitan ; sys_platform == 'linux'",
-    "dion ; sys_platform == 'linux'",
     "tilelang>=0.1.8 ; sys_platform == 'linux'",
     # DSA sparse MLA backward kernel (CuTe DSL); >=1.29 matches the pinned CUTLASS DSL 4.7.
     "nvidia-cudnn-frontend>=1.30.0 ; sys_platform == 'linux'",
@@ -247,7 +246,6 @@ renderers = false
 prime-pydantic-config = false
 # Trusted git/URL sources pinned by rev or wheel URL (defensive: belt-and-suspenders)
 vllm-router = false
-dion = false
 deep-ep = false
 deep-gemm = false
 nixl-cu13 = false
@@ -271,7 +269,6 @@ torchao = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/torchao-0.18.0+git5f2baf9d5-cp310-abi3-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
-dion = { git = "https://github.com/samsja/dion.git", rev = "0c801b0" }
 sonic-moe = { git = "https://github.com/S1ro1/sonic-moe", rev = "0bf06ac" }
 quack-kernels = { git = "https://github.com/S1ro1/quack", rev = "f62009f" }
 flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", subdirectory = "flash_attn/cute", rev = "0251105" }
diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py
index 8514be4eeac..b415892a526 100644
--- a/src/prime_rl/trainer/ckpt.py
+++ b/src/prime_rl/trainer/ckpt.py
@@ -139,7 +139,7 @@ def load_state_dict(self, state_dict: dict[str, Any]):
             # The template only aliases a packed parameter's state where the packing
             # dimension is unsharded, so the loaded logical entries are packed back in.
             write_back_loaded_packed_optimizer_state(self.model, checkpoint_optimizers, state_dict["optimizers"])
-            # The in-place load reaches tensors only: restore the param groups' values too (dion's Muon keeps its
+            # The in-place load reaches tensors only: restore the param groups' values too (Muon keeps its
             # step count there, so AdamW's bias correction would restart).
             loaded_groups = iter(state_dict["optimizers"]["param_groups"])
             for optimizer in checkpoint_optimizers:
diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py
index 08316ec71e0..21f9dc23c58 100644
--- a/src/prime_rl/trainer/optim/__init__.py
+++ b/src/prime_rl/trainer/optim/__init__.py
@@ -1,8 +1,5 @@
 import torch
-import torch.distributed as dist
-from dion import Muon
 from torch import nn
-from torch.distributed.device_mesh import DeviceMesh
 from torch.optim import AdamW, Optimizer
 
 from prime_rl.configs.trainer import OffloadMode, OptimizerConfig
@@ -10,6 +7,7 @@
 from prime_rl.trainer.models.layers.ngram import NGramTable
 from prime_rl.trainer.optim.base import OffloadOptimizer as OffloadOptimizer
 from prime_rl.trainer.optim.base import OptimizerLike
+from prime_rl.trainer.optim.muon import Muon
 from prime_rl.trainer.optim.offload import (
     FullCPUOffloadOptimizer,
     GradientOffloadManager,
@@ -21,33 +19,6 @@
 from prime_rl.utils.logger import get_logger
 
 
-def _warmup_muon_mesh(mesh: DeviceMesh) -> None:
-    """Establish NCCL peer connections before Muon's first optimizer step.
-
-    This is a correctness workaround, not an optional performance warm-up:
-    multi-node GLM-Air training can deadlock on Muon's first bulk all-to-all
-    without it. Optimizer refactors must preserve this initialization.
-    """
-    # get_group() without a mesh dim is only valid for 1-D meshes; the
-    # replicate-only world mesh is 2-D and has no single group to warm.
-    if mesh.ndim != 1:
-        return
-    group = mesh.get_group()
-    size = dist.get_world_size(group)
-    if size <= 1 or dist.get_backend(group) != "nccl":
-        return
-
-    get_logger().info(f"Warming Muon all-to-all connections on group {group.group_name} ({size} ranks)")
-    device = torch.device("cuda", torch.cuda.current_device())
-    # Use a bulk list all-to-all, matching Dion's path rather than a scalar
-    # collective that may leave bulk-transfer peer channels uninitialized.
-    inputs = [torch.zeros(3 * 1024 * 1024, dtype=torch.bfloat16, device=device) for _ in range(size)]
-    outputs = [torch.empty_like(tensor) for tensor in inputs]
-    dist.all_to_all(outputs, inputs, group=group, async_op=True).wait()
-    torch.cuda.synchronize()
-    get_logger().info(f"Finished warming Muon all-to-all connections on group {group.group_name}")
-
-
 def setup_optimizer(
     config: OptimizerConfig,
     named_params: list[tuple[str, nn.Parameter]],
@@ -77,7 +48,6 @@ def setup_optimizer(
     optimizer = _create_optimizer(
         config,
         optimizer_named_params,
-        parallel_dims,
         fused_adamw=config.type == "adamw" and not cpu_offload,
         model=model,
     )
@@ -102,7 +72,6 @@ def setup_optimizer(
 def _create_optimizer(
     config: OptimizerConfig,
     named_params: list[tuple[str, nn.Parameter]],
-    parallel_dims: ParallelDims,
     lr: float | None = None,
     fused_adamw: bool = False,
     model: nn.Module | None = None,
@@ -126,7 +95,7 @@ def _create_optimizer(
                 fused=fused_adamw,
             )
         case "muon":
-            return _create_muon_optimizer(config, named_params, parallel_dims, model, lr)
+            return _create_muon_optimizer(config, named_params, model, lr)
         case "sign_sgd":
             return SignSGD(
                 params=trainable_params,
@@ -138,7 +107,6 @@ def _create_optimizer(
 def _create_muon_optimizer(
     config: OptimizerConfig,
     named_params: list[tuple[str, nn.Parameter]],
-    parallel_dims: ParallelDims,
     model: nn.Module | None,
     lr: float | None = None,
 ) -> Optimizer:
@@ -179,42 +147,13 @@ def muon_enabled(n, p):
         else:
             pass
 
-    param_groups = []
-
-    param_groups.append(
-        dict(params=muon_params, algorithm="muon", lr=lr, weight_decay=config.weight_decay, adjust_lr="rms_norm")
-    )
+    # dion's groups and keys, so that checkpoints of runs on dion's Muon load as they are
+    param_groups = [dict(params=muon_params, algorithm="muon")]
     if expert_params:
-        experts_mesh_name = None
-        if parallel_dims.ep_enabled:
-            experts_mesh_name = "dp_shard_mod_ep"
-        param_groups.append(
-            dict(
-                params=expert_params,
-                algorithm="muon",
-                lr=lr,
-                weight_decay=config.weight_decay,
-                adjust_lr="rms_norm",
-                distributed_mesh_name=experts_mesh_name,
-            )
-        )
+        param_groups.append(dict(params=expert_params, algorithm="muon"))
     if router_params:
-        param_groups.append(
-            dict(
-                params=router_params,
-                algorithm="muon",
-                lr=lr,
-                weight_decay=config.weight_decay,
-                adjust_lr="rms_norm",
-            )
-        )
-
-    param_groups.append(dict(params=adamw_params, algorithm="adamw", lr=lr, weight_decay=config.weight_decay))
-
-    if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled:
-        distributed_mesh = parallel_dims.get_mesh("dp_shard_cp")
-    else:
-        distributed_mesh = parallel_dims.world_mesh
+        param_groups.append(dict(params=router_params, algorithm="muon"))
+    param_groups.append(dict(params=adamw_params, algorithm="adamw"))
 
     # Runtime fusions pack several logical matrices into one physical parameter. Muon
     # orthogonalizes each of them separately, so a packed parameter trains exactly as the
@@ -227,23 +166,15 @@ def muon_enabled(n, p):
             if partitions is not None and packed_info.parameter in muon_params:
                 matrix_partitions[packed_info.parameter] = partitions
 
-    optimizer = Muon(
-        params=param_groups,
-        matrix_partitions=matrix_partitions,
+    return Muon(
+        param_groups,
         lr=lr,
         mu=config.mu,
-        nesterov=config.nesterov,
         betas=(config.betas1, config.betas2),
-        epsilon=config.eps,
         weight_decay=config.weight_decay,
-        adjust_lr="rms_norm",
-        distributed_mesh=distributed_mesh,
-        world_mesh=parallel_dims.world_mesh,
-        fsdp_mesh_dim=1 if parallel_dims.dp_replicate_enabled else 0,
+        epsilon=config.eps,
+        nesterov=config.nesterov,
+        update_rms=config.update_rms,
+        newton_schulz=config.newton_schulz,
+        matrix_partitions=matrix_partitions,
     )
-    # Keep both warm-ups after Muon construction and before its first step. The
-    # main and expert groups establish independent NCCL peer connections.
-    _warmup_muon_mesh(distributed_mesh)
-    if expert_params and parallel_dims.ep_enabled:
-        _warmup_muon_mesh(parallel_dims.get_mesh("dp_shard_mod_ep"))
-    return optimizer
diff --git a/src/prime_rl/trainer/optim/muon.py b/src/prime_rl/trainer/optim/muon.py
new file mode 100644
index 00000000000..55931dd7829
--- /dev/null
+++ b/src/prime_rl/trainer/optim/muon.py
@@ -0,0 +1,699 @@
+"""Muon (https://kellerjordan.github.io/posts/muon/) for FSDP2 and expert-parallel models.
+
+A Muon parameter is updated with its orthogonalized Nesterov momentum, scaled to an RMS of ``update_rms``;
+every other parameter gets AdamW. Orthogonalization needs whole matrices:
+
+- A parameter whose matrix dims are not sharded (an expert stack sharded only along the expert dim, a
+  replicated or a single-device parameter) is orthogonalized where it lives, as one batch of matrices.
+- A parameter sharded along a matrix dim is gathered onto one rank of its shard group, orthogonalized there
+  and scattered back, with uneven all-to-alls and no padding. Ranks own matrices by Newton-Schulz FLOPs, and
+  the matrices move in buckets: bucket k+1 gathers while bucket k is orthogonalized.
+
+Matrices packed into one parameter (runtime fusions, per-head partitions) are orthogonalized each on its own,
+and every Newton-Schulz runs batched over all matrices of one shape. The state ("momentum", plus "variance"
+for AdamW) and the param group keys are those of dion's Muon, so its checkpoints load as they are.
+"""
+
+import functools
+import heapq
+import importlib.util
+import math
+from collections import defaultdict
+from collections.abc import Collection, Iterable, Mapping, Sequence
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Shard
+from torch.distributed.tensor.placement_types import _StridedShard
+from torch.optim import Optimizer
+
+from prime_rl.utils.logger import get_logger
+
+# Whole matrices each rank receives per bucket. Bounds the step's transient memory to a few buckets.
+BUCKET_BYTES = 256 * 2**20
+# CTAs per NCCL collective on Muon's communicators
+MAX_CTAS = 4
+
+
+@dataclass(frozen=True)
+class NewtonSchulz:
+    """Odd polynomial iterations X <- aX + (bA + cA²)X with A = XXᵀ, one (a, b, c) per iteration, on
+    X = G / (safety·‖G‖_F + epsilon), normalized in bf16 and iterated in ``dtype``."""
+
+    coefficients: tuple[tuple[float, float, float], ...]
+    safety: float = 1.0
+    dtype: torch.dtype = torch.bfloat16
+
+
+# Polar Express (arXiv 2505.16932), divided by the safety factor except for the last polynomial, which repeats
+POLAR_EXPRESS = tuple(
+    (a / 1.01, b / 1.01**3, c / 1.01**5)
+    for a, b, c in (
+        (8.28721201814563, -23.595886519098837, 17.300387312530933),
+        (4.107059111542203, -2.9478499167379106, 0.5448431082926601),
+        (3.9486908534822946, -2.908902115962949, 0.5518191394370137),
+        (3.3184196573706015, -2.488488024314874, 0.51004894012372),
+        (2.300652019954817, -1.6689039845747493, 0.4188073119525673),
+        (1.891301407787398, -1.2679958271945868, 0.37680408948524835),
+        (1.8750014808534479, -1.2500016453999487, 0.3750001645474248),
+    )
+) + ((1.875, -1.25, 0.375),)
+
+NEWTON_SCHULZ = {
+    # Per-step quintic coefficients
+    "default": NewtonSchulz(
+        (
+            (4.0848, -6.8946, 2.9270),
+            (3.9505, -6.3029, 2.6377),
+            (3.7418, -5.5913, 2.3037),
+            (2.8769, -3.1427, 1.2046),
+            (2.8366, -3.0525, 1.2012),
+        )
+    ),
+    # DeepSeek-V4: 8 fast-converging steps, then 2 that pin the singular values at 1
+    "deepseek_v4": NewtonSchulz(((3.4445, -4.7750, 2.0315),) * 8 + ((2.0, -1.5, 0.5),) * 2),
+    "polar_express": NewtonSchulz(POLAR_EXPRESS[:5], safety=1.01),
+    # Step 3.5 Flash: 6 Polar Express steps in fp16 (bf16 spiked)
+    "step3p5": NewtonSchulz(POLAR_EXPRESS[:6], safety=1.01, dtype=torch.float16),
+}
+
+
+@functools.cache
+def _quack_gemm_symmetric():
+    """QuACK's symmetric GEMM (computes one triangle and mirrors it), or None without QuACK or before SM80."""
+    if importlib.util.find_spec("quack") is None or torch.cuda.get_device_capability()[0] < 8:
+        return None
+    from quack.gemm_interface import gemm_symmetric
+
+    return gemm_symmetric
+
+
+@torch.compile(fullgraph=True, dynamic=True)
+def _normalize(G: Tensor, safety: float, epsilon: float) -> Tensor:
+    """G / (safety·‖G‖_F + epsilon) per matrix, in bf16, as one fused reduction."""
+    X = G.to(torch.bfloat16)
+    return X / (X.norm(dim=(-2, -1), keepdim=True) * safety + epsilon)
+
+
+def newton_schulz(G: Tensor, schedule: NewtonSchulz, epsilon: float) -> Tensor:
+    """Orthogonalize every matrix of the batch ``G`` [B, m, n]. Returns [B, m, n] in ``schedule.dtype``.
+
+    The Gram matrix A is the smaller of XXᵀ and XᵀX. For tall matrices this iterates X <- aX + X(bA + cA²) on
+    A = XᵀX, the transpose of the wide iteration, without transposing X."""
+    X = _normalize(G, schedule.safety, epsilon).to(schedule.dtype).contiguous()
+    tall = X.size(-2) > X.size(-1)
+    # QuACK needs 16-byte aligned rows
+    gemm_symmetric = _quack_gemm_symmetric() if X.is_cuda and X.size(-1) % 8 == 0 and X.size(-2) % 8 == 0 else None
+    gram = gemm_symmetric or torch.bmm
+    for a, b, c in schedule.coefficients:
+        A = gram(X.mT, X) if tall else gram(X, X.mT)
+        if gemm_symmetric is not None:
+            B = gemm_symmetric(A, A, C=A, alpha=c, beta=b)
+        else:
+            B = torch.baddbmm(A, A, A, beta=b, alpha=c)
+        X = torch.baddbmm(X, X, B, beta=a) if tall else torch.baddbmm(X, B, X, beta=a)
+    return X
+
+
+def _batch(stacks: list[Tensor]) -> Tensor:
+    """Concatenate stacks of matrices [B_i, m, n], as a view when they lie back to back in memory."""
+    first = stacks[0]
+    adjacent = all(
+        x.is_contiguous() and x.untyped_storage().data_ptr() == first.untyped_storage().data_ptr() for x in stacks
+    )
+    for previous, x in zip(stacks, stacks[1:]):
+        adjacent = adjacent and x.storage_offset() == previous.storage_offset() + previous.numel()
+    if not adjacent:
+        return torch.cat(stacks)
+    size = (sum(x.size(0) for x in stacks), *first.shape[1:])
+    return first.as_strided(size, (first[0].numel(), first.size(-1), 1))
+
+
+def orthogonalize(
+    matrices: Sequence[Tensor],
+    partitions: Sequence[tuple[int, ...]],
+    schedule: NewtonSchulz,
+    epsilons: Sequence[float],
+) -> list[Tensor]:
+    """Orthogonalize tensors [..., M, n] whose dim -2 packs matrices of the given row counts, one Newton-Schulz
+    per matrix shape over all of them. Leading dims are batches of matrices."""
+    outputs: list[Tensor | None] = [None] * len(matrices)
+    classes = defaultdict(list)
+    for i, (x, parts, epsilon) in enumerate(zip(matrices, partitions, epsilons)):
+        x = x.reshape(-1, *x.shape[-2:])
+        if len(set(parts)) == 1:
+            # Equal partitions are a reshape: [B, P·r, n] -> [B·P, r, n]
+            classes[parts[0], x.size(-1), epsilon].append((i, None, x.reshape(-1, parts[0], x.size(-1))))
+        else:
+            for start, rows, part in zip(_offsets(parts), parts, x.split(parts, dim=-2)):
+                classes[rows, x.size(-1), epsilon].append((i, (start, rows), part))
+    for (_, _, epsilon), members in classes.items():
+        Y = newton_schulz(_batch([x for _, _, x in members]), schedule, epsilon)
+        for (i, rows, x), y in zip(members, Y.split([x.size(0) for _, _, x in members])):
+            if rows is None:
+                outputs[i] = y.reshape(matrices[i].shape)
+                continue
+            if outputs[i] is None:
+                outputs[i] = torch.empty(matrices[i].shape, dtype=y.dtype, device=y.device)
+            outputs[i].view(-1, *matrices[i].shape[-2:]).narrow(-2, *rows).copy_(y)
+    return outputs
+
+
+def plan_buckets(
+    costs: Sequence[float], nbytes: Sequence[int], world_size: int, bucket_bytes: int
+) -> list[list[tuple[int, int]]]:
+    """Cut the matrices, in order, into buckets of about ``bucket_bytes`` per rank, and give each matrix an owner
+    rank: (matrix, owner) pairs per bucket. Within a bucket, owners balance the Newton-Schulz cost, longest
+    first onto the least loaded rank, so every bucket keeps all ranks equally busy. Keeping the order keeps a
+    bucket to a few consecutive layers."""
+    buckets, bucket, size = [], [], 0
+    for i, n in enumerate(nbytes):
+        if bucket and size + n > world_size * bucket_bytes:
+            buckets.append(bucket)
+            bucket, size = [], 0
+        bucket.append(i)
+        size += n
+    buckets += [bucket] * bool(bucket)
+    owned = []
+    for bucket in buckets:
+        loads = [(0.0, rank) for rank in range(world_size)]
+        pairs = []
+        for i in sorted(bucket, key=lambda i: (-costs[i], i)):
+            load, rank = heapq.heappop(loads)
+            pairs.append((i, rank))
+            heapq.heappush(loads, (load + costs[i], rank))
+        owned.append(sorted(pairs))
+    return owned
+
+
+def newton_schulz_cost(shape: Sequence[int], partitions: Sequence[int]) -> float:
+    """FLOPs of one Newton-Schulz iteration over a tensor [..., M, n] packing matrices of ``partitions`` rows."""
+    batch = math.prod(shape[:-2])
+    return batch * sum(
+        4 * min(rows, shape[-1]) ** 2 * max(rows, shape[-1]) + 2 * min(rows, shape[-1]) ** 3 for rows in partitions
+    )
+
+
+def _offsets(sizes: Sequence[int]) -> list[int]:
+    offsets, total = [], 0
+    for size in sizes:
+        offsets.append(total)
+        total += size
+    return offsets
+
+
+def _local(tensor: Tensor) -> Tensor:
+    # The attribute, not to_local(): that is an autograd function, and the step reads thousands of local tensors
+    return tensor._local_tensor if isinstance(tensor, DTensor) else tensor
+
+
+def _by_group(matrices: list["_Matrix"], tensors: list[Tensor]) -> list[tuple[dict, list[tuple["_Matrix", Tensor]]]]:
+    """Pair matrices with their tensors, per param group, leaving out empty shards."""
+    groups = {}
+    for matrix, tensor in zip(matrices, tensors):
+        if tensor.numel():
+            groups.setdefault(id(matrix.group), (matrix.group, []))[1].append((matrix, tensor))
+    return list(groups.values())
+
+
+def _shape_class(matrix: "_Matrix") -> tuple:
+    """Matrices with equal keys orthogonalize as one batch (see orthogonalize)."""
+    return _local(matrix.param).shape[-1], matrix.partitions
+
+
+def _num_shards(param: Tensor) -> int:
+    if not isinstance(param, DTensor):
+        return 1
+    mesh = param.device_mesh
+    return math.prod(mesh.size(d) for d, p in enumerate(param.placements) if isinstance(p, (Shard, _StridedShard)))
+
+
+@dataclass(frozen=True)
+class _Sharding:
+    """How a parameter's shard group splits one of its matrix dims, in torch.chunk layout."""
+
+    dim: int
+    group: ProcessGroup
+    rank: int
+    sizes: tuple[int, ...]
+
+    @property
+    def offsets(self) -> list[int]:
+        return _offsets(self.sizes)
+
+
+_MUON_GROUPS: dict[tuple[tuple[int, ...], ...], ProcessGroup] = {}
+
+
+def _muon_group(mesh: DeviceMesh, mesh_dim: int) -> ProcessGroup:
+    """This rank's group along ``mesh_dim``, as a communicator of Muon's own: on FSDP's, Muon's all-to-alls
+    would queue behind (or ahead of) FSDP's all-gathers. Every rank creates every group, in the same order."""
+    layout = tuple(map(tuple, mesh.mesh.movedim(mesh_dim, -1).reshape(-1, mesh.size(mesh_dim)).tolist()))
+    if layout not in _MUON_GROUPS:
+        options = None
+        if dist.get_backend() == "nccl":
+            # Muon moves little data: few channels keep the communicator's buffers small (2.8 GiB per rank on
+            # 8 GPUs with NCCL's default) and leave the SMs to Newton-Schulz
+            options = dist.ProcessGroupNCCL.Options()
+            options.config.max_ctas = MAX_CTAS
+        groups = [list(ranks) for ranks in layout]
+        _MUON_GROUPS[layout], _ = dist.new_subgroups_by_enumeration(groups, pg_options=options)
+        _warmup_all_to_all(_MUON_GROUPS[layout])
+    return _MUON_GROUPS[layout]
+
+
+def _matrix_sharding(param: Tensor) -> _Sharding | None:
+    """The sharding of ``param``'s matrix dims, or None when every rank holds whole matrices."""
+    if not isinstance(param, DTensor):
+        return None
+    mesh = param.device_mesh
+    shards = [
+        (mesh_dim, placement)
+        for mesh_dim, placement in enumerate(param.placements)
+        if isinstance(placement, (Shard, _StridedShard)) and mesh.size(mesh_dim) > 1 and placement.dim >= param.ndim - 2
+    ]
+    if not shards:
+        return None
+    if len(shards) > 1 or isinstance(shards[0][1], _StridedShard):
+        raise NotImplementedError(f"Muon gathers matrices over one plain Shard, got placements {param.placements}")
+    mesh_dim, placement = shards[0]
+    group = _muon_group(mesh, mesh_dim)
+    # The all-to-all splits are in group rank order, DTensor's shards in mesh order
+    coordinate = list(mesh.get_coordinate())
+    coordinate[mesh_dim] = slice(None)
+    if mesh.mesh[tuple(coordinate)].tolist() != dist.get_process_group_ranks(group):
+        raise ValueError(f"Muon needs the shard ranks of {param.placements} in ascending order")
+    size, world_size = param.shape[placement.dim], mesh.size(mesh_dim)
+    chunk = -(-size // world_size)
+    sizes = tuple(max(0, min(chunk, size - rank * chunk)) for rank in range(world_size))
+    rank = mesh.get_local_rank(mesh_dim)
+    assert param.to_local().shape[placement.dim] == sizes[rank]
+    return _Sharding(placement.dim, group, rank, sizes)
+
+
+@dataclass(frozen=True)
+class _Matrix:
+    """A Muon parameter: its matrix partitions along dim -2, sharding and update scale."""
+
+    param: Tensor
+    group: dict
+    partitions: tuple[int, ...]
+    sharding: _Sharding | None
+    lr_ratio: float | Tensor  # per local row for packed parameters
+
+    def local_shape(self, rank: int) -> list[int]:
+        """The shape of ``rank``'s shard."""
+        shape = list(_local(self.param).shape)
+        shape[self.sharding.dim] = self.sharding.sizes[rank]
+        return shape
+
+    def whole_shape(self) -> list[int]:
+        """The shape of the gathered matrices: whole along the sharded dim, local along batch dims."""
+        shape = list(_local(self.param).shape)
+        shape[self.sharding.dim] = sum(self.sharding.sizes)
+        return shape
+
+
+@dataclass(frozen=True)
+class _Bucket:
+    """Sharded matrices that move together. Shards are laid out by owner rank on the sender's side (``sent``)
+    and by sender rank on the owner's side (``owned``), the same way in both directions."""
+
+    group: ProcessGroup
+    sent: list[_Matrix]
+    owned: list[_Matrix]
+    send_splits: list[int]
+    recv_splits: list[int]
+
+
+@dataclass(frozen=True)
+class _Plan:
+    """The work of one step, in units that update together: buckets of sharded matrices, chunks of local
+    matrices orthogonalized as one batch, and AdamW groups, each one compiled update."""
+
+    buckets: list[_Bucket]
+    local: list[list[_Matrix]]
+    adamw: list[tuple[dict, list[Tensor]]]
+
+    @property
+    def units(self) -> list[list[Tensor]]:
+        return (
+            [[m.param for m in bucket.sent] for bucket in self.buckets]
+            + [[m.param for m in chunk] for chunk in self.local]
+            + [params for _, params in self.adamw]
+        )
+
+    def restrict(self, params: Iterable[Tensor]) -> "_Plan":
+        """The units of ``params``, a union of whole units."""
+        selected = set(map(id, params))
+        return _Plan(
+            buckets=[bucket for bucket in self.buckets if id(bucket.sent[0].param) in selected],
+            local=[chunk for chunk in self.local if id(chunk[0].param) in selected],
+            adamw=[(group, params) for group, params in self.adamw if id(params[0]) in selected],
+        )
+
+
+@dataclass
+class _Step:
+    """One Muon step as partial updates (``SubsetStep``): the hyperparameters and gradients are captured when
+    the step begins, so a unit updated later updates as it would have then."""
+
+    optimizer: "Muon"
+    plan: _Plan
+    hyperparameters: dict[int, dict]
+    grads: dict[Tensor, Tensor]
+
+    @property
+    def units(self) -> list[list[Tensor]]:
+        return self.plan.units
+
+    def step_params(self, params: Collection[Tensor]) -> None:
+        self.optimizer._run(self.plan.restrict(params), self.hyperparameters, self.grads)
+
+
+# The element-wise updates take a whole unit's tensors per call: one compiled call costs ~50 µs of host time.
+@torch.compile(fullgraph=True, dynamic=True)
+def _momentum_(M: list[Tensor], G: list[Tensor], U: list[Tensor], mu: Tensor, nesterov: bool) -> None:
+    """M <- μM + G and U <- μM + G (Nesterov) or M."""
+    for m, g, u in zip(M, G, U):
+        g = g.to(m.dtype)
+        m.mul_(mu).add_(g)
+        u.copy_(m * mu + g if nesterov else m)
+
+
+@torch.compile(fullgraph=True, dynamic=True)
+def _muon_update_(
+    X: list[Tensor], U: list[Tensor], lr: Tensor, adjusted_lr: list[Tensor], weight_decay: Tensor
+) -> None:
+    for x, u, a in zip(X, U, adjusted_lr):
+        x.mul_(1 - lr * weight_decay)
+        x.sub_(u * a)
+
+
+@torch.compile(fullgraph=True, dynamic=True)
+def _adamw_(
+    X: list[Tensor],
+    G: list[Tensor],
+    M: list[Tensor],
+    V: list[Tensor],
+    lr: Tensor,
+    beta1: Tensor,
+    beta2: Tensor,
+    weight_decay: Tensor,
+    step: Tensor,
+    epsilon: Tensor,
+) -> None:
+    """dion's AdamW, op for op: X <- (1 - lr·λ)X - lr/(1 - β1ᵗ) · M / (√V / √(1 - β2ᵗ) + ε)."""
+    G = [g.to(M[0].dtype) for g in G]
+    torch._foreach_lerp_(M, G, [1 - beta1] * len(X))
+    torch._foreach_lerp_(V, [g.to(V[0].dtype) for g in torch._foreach_mul(G, G)], [1 - beta2] * len(X))
+    denom = torch._foreach_sqrt(V)
+    torch._foreach_div_(denom, (1 - beta2**step).sqrt())
+    torch._foreach_add_(denom, [epsilon] * len(X))
+    torch._foreach_mul_(X, 1 - lr * weight_decay)
+    update = torch._foreach_div(M, denom)
+    torch._foreach_mul_(update, lr / (1 - beta1**step))
+    torch._foreach_sub_(X, update)
+
+
+def _hyperparameters(group: dict) -> dict:
+    # 0-d CPU tensors, so the compiled updates do not recompile when the scheduler changes the lr
+    keys = ("lr", "mu", "beta1", "beta2", "weight_decay", "epsilon", "step")
+    return {key: torch.tensor(group[key]) for key in keys} | {"nesterov": group["nesterov"]}
+
+
+def _warmup_all_to_all(group: ProcessGroup) -> None:
+    """Establish NCCL peer connections before the first optimizer step.
+
+    This is a correctness workaround, not an optional performance warm-up: multi-node GLM-Air training can
+    deadlock on Muon's first bulk all-to-all without it.
+    """
+    size = dist.get_world_size(group)
+    if size <= 1 or dist.get_backend(group) != "nccl":
+        return
+    get_logger().info(f"Warming Muon all-to-all connections on group {group.group_name} ({size} ranks)")
+    # A bulk transfer, like the step's, rather than a scalar collective that may leave bulk channels unset
+    inputs = torch.zeros(size * 3 * 2**20, dtype=torch.bfloat16, device="cuda")
+    dist.all_to_all_single(torch.empty_like(inputs), inputs, group=group)
+    torch.cuda.synchronize()
+
+
+class Muon(Optimizer):
+    """Muon for the param groups with ``algorithm="muon"`` (matrices, ndim >= 2) and AdamW for those with
+    ``algorithm="adamw"``.
+
+    Muon: M <- μM + G, U <- orthogonalize(μM + G) (Nesterov) or orthogonalize(M), and
+    X <- (1 - lr·λ)X - lr·update_rms·√max(m, n)·U for every matrix m x n. ``matrix_partitions`` maps a
+    parameter to the row counts of the matrices packed along its dim -2, each orthogonalized and scaled as its
+    own matrix.
+
+    ``step()`` runs ``begin_subset_step()`` and updates all of its units, so an offloaded step that updates
+    one chunk of units at a time (``SupportsSubsetStep``) is bitwise the same step. Units are identical on
+    every rank, and the step never synchronizes with the host.
+    """
+
+    def __init__(
+        self,
+        params: Iterable[dict],
+        *,
+        lr: float,
+        mu: float = 0.95,
+        betas: tuple[float, float] = (0.9, 0.95),
+        weight_decay: float = 0.01,
+        epsilon: float = 1e-8,
+        nesterov: bool = True,
+        update_rms: float = 0.2,
+        newton_schulz: str = "default",
+        matrix_partitions: Mapping[Tensor, Sequence[int]] | None = None,
+        bucket_bytes: int = BUCKET_BYTES,
+    ):
+        defaults = dict(
+            algorithm="muon",
+            lr=lr,
+            mu=mu,
+            beta1=betas[0],
+            beta2=betas[1],
+            weight_decay=weight_decay,
+            epsilon=epsilon,
+            nesterov=nesterov,
+            step=0,
+        )
+        super().__init__(params, defaults)
+        self.update_rms = update_rms
+        self.schedule = NEWTON_SCHULZ[newton_schulz]
+        self.bucket_bytes = bucket_bytes
+        self._plans: dict[tuple[int, ...], _Plan] = {}
+
+        partitions = {param: tuple(sizes) for param, sizes in (matrix_partitions or {}).items()}
+        self._matrices: dict[Tensor, _Matrix] = {}
+        for group in self.param_groups:
+            if group["algorithm"] not in ("muon", "adamw"):
+                raise ValueError(f"Unknown algorithm {group['algorithm']}")
+            if group["algorithm"] == "muon":
+                for param in group["params"]:
+                    self._matrices[param] = self._matrix(param, group, partitions.pop(param, (param.size(-2),)))
+        if partitions:
+            raise ValueError("matrix_partitions has parameters that are not Muon parameters of this optimizer")
+
+    def _matrix(self, param: Tensor, group: dict, partitions: tuple[int, ...]) -> _Matrix:
+        if param.ndim < 2:
+            raise ValueError(f"Muon parameters must be matrices, got shape {tuple(param.shape)}")
+        if sum(partitions) != param.size(-2) or min(partitions) <= 0:
+            raise ValueError(f"Partitions {partitions} do not split dim -2 of shape {tuple(param.shape)}")
+        sharding = _matrix_sharding(param)
+        ratios = [self.update_rms * math.sqrt(max(rows, param.size(-1))) for rows in partitions]
+        lr_ratio = ratios[0]
+        if len(partitions) > 1:
+            # Per row, on the device, so the step scales packed matrices without a host-to-device copy
+            rows = torch.tensor(ratios).repeat_interleave(torch.tensor(partitions))
+            if sharding is not None and sharding.dim == param.ndim - 2:
+                rows = rows.narrow(0, sharding.offsets[sharding.rank], sharding.sizes[sharding.rank])
+            lr_ratio = rows.unsqueeze(-1).to(_local(param).device)
+        return _Matrix(param, group, partitions, sharding, lr_ratio)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        step = self.begin_subset_step()
+        step.step_params([param for unit in step.units for param in unit])
+        return loss
+
+    def begin_subset_step(self) -> _Step:
+        """Start a step: advance the step counters and capture the hyperparameters and gradients. The returned
+        step updates any union of its ``units`` with ``step_params``."""
+        for group in self.param_groups:
+            group["step"] += 1
+        params = [param for group in self.param_groups for param in group["params"] if param.grad is not None]
+        key = tuple(map(id, params))
+        if key not in self._plans:
+            self._plans[key] = self._plan(params)
+        hyperparameters = {id(group): _hyperparameters(group) for group in self.param_groups}
+        return _Step(self, self._plans[key], hyperparameters, {param: _local(param.grad) for param in params})
+
+    def _plan(self, params: list[Tensor]) -> _Plan:
+        matrices = [self._matrices[param] for param in params if param in self._matrices]
+        by_group = defaultdict(list)
+        for matrix in matrices:
+            if matrix.sharding is not None:
+                by_group[matrix.sharding.group].append(matrix)
+        buckets = []
+        for group, members in by_group.items():
+            world_size, rank = dist.get_world_size(group), dist.get_rank(group)
+            costs = [newton_schulz_cost(m.param.shape, m.partitions) for m in members]
+            nbytes = [2 * m.param.numel() for m in members]
+            for pairs in plan_buckets(costs, nbytes, world_size, self.bucket_bytes):
+                # Every owner's matrices grouped by shape, in the same order on the sender's and owner's side
+                owned = [
+                    sorted((members[i] for i, owner in pairs if owner == r), key=_shape_class)
+                    for r in range(world_size)
+                ]
+                send_splits = [sum(math.prod(m.local_shape(rank)) for m in owned[r]) for r in range(world_size)]
+                recv_splits = [sum(math.prod(m.local_shape(q)) for m in owned[rank]) for q in range(world_size)]
+                sent = [m for r in range(world_size) for m in owned[r]]
+                buckets.append(_Bucket(group, sent, owned[rank], send_splits, recv_splits))
+
+        # Chunked by shard sizes that are the same on every rank, so that every rank has the same units
+        local, chunk, chunk_bytes = [], [], 0
+        for matrix in (m for m in matrices if m.sharding is None):
+            nbytes = 2 * matrix.param.numel() // _num_shards(matrix.param)
+            if chunk and chunk_bytes + nbytes > self.bucket_bytes:
+                local.append(chunk)
+                chunk, chunk_bytes = [], 0
+            chunk.append(matrix)
+            chunk_bytes += nbytes
+        local += [chunk] * bool(chunk)
+
+        selected = set(map(id, params))
+        adamw = [
+            (group, [param for param in group["params"] if id(param) in selected])
+            for group in self.param_groups
+            if group["algorithm"] == "adamw"
+        ]
+        return _Plan(buckets, local, [(group, params) for group, params in adamw if params])
+
+    @torch.no_grad()
+    def _run(self, plan: _Plan, hyperparameters: dict[int, dict], grads: dict[Tensor, Tensor]) -> None:
+        # Bucket k+1 gathers and bucket k-1 scatters while bucket k is orthogonalized. Local chunks fill in.
+        slots = len(plan.buckets) + 1
+        local = [plan.local[i * len(plan.local) // slots : (i + 1) * len(plan.local) // slots] for i in range(slots)]
+        gathering = self._gather(plan.buckets[0], hyperparameters, grads) if plan.buckets else None
+        scattering = None
+        for i, bucket in enumerate(plan.buckets):
+            gathered = gathering
+            if i + 1 < len(plan.buckets):
+                gathering = self._gather(plan.buckets[i + 1], hyperparameters, grads)
+            for chunk in local[i]:
+                self._update_local(chunk, hyperparameters, grads)
+            orthogonalized = self._orthogonalize_and_scatter(bucket, *gathered)
+            if scattering is not None:
+                self._apply(*scattering, hyperparameters)
+            scattering = (bucket, *orthogonalized)
+        for chunk in local[-1]:
+            self._update_local(chunk, hyperparameters, grads)
+        if scattering is not None:
+            self._apply(*scattering, hyperparameters)
+
+        for group, params in plan.adamw:
+            self._adamw(params, hyperparameters[id(group)], grads)
+
+    def _state(self, param: Tensor, *keys: str) -> list[Tensor]:
+        """State tensors, read at every use: an offloaded step swaps them for GPU copies between updates."""
+        state = self.state[param]
+        for key in keys:
+            if key not in state:
+                state[key] = torch.zeros_like(param)
+        return [_local(state[key]) for key in keys]
+
+    def _momentum(
+        self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: dict, grads: dict[Tensor, Tensor]
+    ) -> None:
+        """M <- μM + G, and the input of the orthogonalization into ``updates``, one call per param group."""
+        for group, members in _by_group(matrices, updates):
+            h = hyperparameters[id(group)]
+            M = [self._state(m.param, "momentum")[0] for m, _ in members]
+            G = [grads[m.param] for m, _ in members]
+            _momentum_(M, G, [U for _, U in members], h["mu"], h["nesterov"])
+
+    def _update(self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: dict) -> None:
+        for group, members in _by_group(matrices, updates):
+            h = hyperparameters[id(group)]
+            X = [_local(m.param) for m, _ in members]
+            adjusted_lr = [m.lr_ratio * h["lr"] for m, _ in members]
+            _muon_update_(X, [U for _, U in members], h["lr"], adjusted_lr, h["weight_decay"])
+
+    def _update_local(self, chunk: list[_Matrix], hyperparameters: dict, grads: dict[Tensor, Tensor]) -> None:
+        """Orthogonalize matrices that every rank holds whole, as one batch per matrix shape. The chunk's inputs
+        share one buffer in which matrices of one shape lie back to back, so each batch is a view of it."""
+        matrices = sorted((m for m in chunk if _local(m.param).numel()), key=_shape_class)
+        if not matrices:
+            return
+        device = _local(matrices[0].param).device
+        buffer = torch.empty(sum(_local(m.param).numel() for m in matrices), dtype=torch.bfloat16, device=device)
+        updates = self._shard_views(buffer, matrices)
+        self._momentum(matrices, updates, hyperparameters, grads)
+        epsilons = [m.group["epsilon"] for m in matrices]
+        updates = orthogonalize(updates, [m.partitions for m in matrices], self.schedule, epsilons)
+        self._update(matrices, updates, hyperparameters)
+
+    def _gather(
+        self, bucket: _Bucket, hyperparameters: dict, grads: dict[Tensor, Tensor]
+    ) -> tuple[dist.Work, Tensor, Tensor]:
+        """Momentum into this rank's shards, sent to their owners."""
+        device = _local(bucket.sent[0].param).device
+        send = torch.empty(sum(bucket.send_splits), dtype=torch.bfloat16, device=device)
+        self._momentum(bucket.sent, self._shard_views(send, bucket.sent), hyperparameters, grads)
+        recv = send.new_empty(sum(bucket.recv_splits))
+        work = dist.all_to_all_single(
+            recv, send, bucket.recv_splits, bucket.send_splits, group=bucket.group, async_op=True
+        )
+        return work, recv, send
+
+    def _orthogonalize_and_scatter(
+        self, bucket: _Bucket, work: dist.Work, recv: Tensor, send: Tensor
+    ) -> tuple[dist.Work, Tensor, Tensor]:
+        """Orthogonalize the matrices this rank owns and send each rank its shards back."""
+        work.wait()
+        world_size = len(bucket.recv_splits)
+        pieces = iter(recv.split([math.prod(m.local_shape(q)) for q in range(world_size) for m in bucket.owned]))
+        shards = [[next(pieces) for _ in bucket.owned] for _ in range(world_size)]
+        # Whole matrices back to back, grouped by shape (see _plan), so each Newton-Schulz batch is a view
+        shapes = [m.whole_shape() for m in bucket.owned]
+        buffer = recv.new_empty(sum(math.prod(shape) for shape in shapes))
+        matrices = [piece.view(shape) for piece, shape in zip(buffer.split([math.prod(s) for s in shapes]), shapes)]
+        for j, (m, matrix) in enumerate(zip(bucket.owned, matrices)):
+            torch.cat([shards[q][j].view(m.local_shape(q)) for q in range(world_size)], dim=m.sharding.dim, out=matrix)
+        epsilons = [m.group["epsilon"] for m in bucket.owned]
+        outputs = orthogonalize(matrices, [m.partitions for m in bucket.owned], self.schedule, epsilons)
+        scattered = [
+            output.narrow(m.sharding.dim, m.sharding.offsets[q], m.sharding.sizes[q]).reshape(-1)
+            for q in range(world_size)
+            for output, m in zip(outputs, bucket.owned)
+        ]
+        send = torch.cat(scattered).to(torch.bfloat16) if scattered else recv.new_empty(0)
+        recv = send.new_empty(sum(bucket.send_splits))
+        work = dist.all_to_all_single(
+            recv, send, bucket.send_splits, bucket.recv_splits, group=bucket.group, async_op=True
+        )
+        return work, recv, send
+
+    def _apply(self, bucket: _Bucket, work: dist.Work, recv: Tensor, send: Tensor, hyperparameters: dict) -> None:
+        work.wait()
+        self._update(bucket.sent, self._shard_views(recv, bucket.sent), hyperparameters)
+
+    @staticmethod
+    def _shard_views(buffer: Tensor, matrices: list[_Matrix]) -> list[Tensor]:
+        shapes = [_local(m.param).shape for m in matrices]
+        return [piece.view(shape) for piece, shape in zip(buffer.split([math.prod(s) for s in shapes]), shapes)]
+
+    def _adamw(self, params: list[Tensor], h: dict, grads: dict[Tensor, Tensor]) -> None:
+        """One AdamW over the whole group, as dion: its compiled update rounds differently for other lists."""
+        M, V = zip(*(self._state(param, "momentum", "variance") for param in params))
+        X, G = [_local(param) for param in params], [grads[param] for param in params]
+        _adamw_(X, G, list(M), list(V), h["lr"], h["beta1"], h["beta2"], h["weight_decay"], h["step"], h["epsilon"])
diff --git a/tests/unit/train/test_muon.py b/tests/unit/train/test_muon.py
new file mode 100644
index 00000000000..923bb382ff1
--- /dev/null
+++ b/tests/unit/train/test_muon.py
@@ -0,0 +1,129 @@
+import math
+
+import pytest
+import torch
+
+from prime_rl.trainer.optim.muon import NEWTON_SCHULZ, Muon, newton_schulz, orthogonalize, plan_buckets
+
+DEFAULT = NEWTON_SCHULZ["default"]
+
+
+def newton_schulz_fp32(G: torch.Tensor) -> torch.Tensor:
+    X = G.mT if G.size(-2) > G.size(-1) else G
+    X = X / X.norm(dim=(-2, -1), keepdim=True)
+    for a, b, c in DEFAULT.coefficients:
+        A = X @ X.mT
+        X = a * X + (b * A + c * A @ A) @ X
+    return X.mT if G.size(-2) > G.size(-1) else X
+
+
+def newton_schulz_bf16(G: torch.Tensor) -> torch.Tensor:
+    """The plain bf16 Newton-Schulz Muon has used so far (no fused GEMM epilogues)."""
+    X = G.bfloat16()
+    X = X.mT if G.size(-2) > G.size(-1) else X
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-8)
+    for a, b, c in DEFAULT.coefficients:
+        A = X @ X.mT
+        X = a * X + (b * A + c * (A @ A)) @ X
+    return X.mT if G.size(-2) > G.size(-1) else X
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("shape", [(1, 1024, 3072), (1, 3072, 1024), (8, 768, 2048), (4, 100, 60)])
+def test_newton_schulz_is_as_accurate_as_bf16(shape):
+    torch.manual_seed(0)
+    G = torch.randn(shape, device="cuda")
+    reference = newton_schulz_fp32(G)
+
+    def error(X: torch.Tensor) -> float:
+        return ((X.float() - reference).norm() / reference.norm()).item()
+
+    assert error(newton_schulz(G, DEFAULT, 1e-8)) <= 1.25 * error(newton_schulz_bf16(G))
+
+
+def test_orthogonalize_packed_matrices_each_on_their_own():
+    torch.manual_seed(0)
+    packed = torch.randn(3, 48 + 16 + 16, 32)
+    equal = torch.randn(2, 4 * 8, 32)
+    out_packed, out_equal = orthogonalize([packed, equal], [(48, 16, 16), (8,) * 4], DEFAULT, [1e-8, 1e-8])
+    for out, x, parts in ((out_packed, packed, (48, 16, 16)), (out_equal, equal, (8,) * 4)):
+        expected = torch.cat([newton_schulz(part, DEFAULT, 1e-8) for part in x.split(parts, dim=-2)], dim=-2)
+        torch.testing.assert_close(out, expected, atol=1e-2, rtol=0)
+
+
+def test_plan_buckets_keeps_order_and_balances_every_bucket():
+    # Layers of a dense model: q, k, v, o, gate, up, down
+    costs = [2.0, 0.1, 0.1, 2.0, 4.8, 4.8, 4.8] * 36
+    nbytes = [32, 8, 8, 32, 96, 96, 96] * 36
+    buckets = plan_buckets(costs, nbytes, world_size=8, bucket_bytes=100)
+    assert [i for bucket in buckets for i, _ in bucket] == list(range(len(costs)))
+    for bucket in buckets:
+        assert sum(nbytes[i] for i, _ in bucket) <= 8 * 100
+        loads = [sum(costs[i] for i, owner in bucket if owner == rank) for rank in range(8)]
+        assert max(loads) - min(loads) <= max(costs[i] for i, _ in bucket)
+
+
+def reference_muon_step(params, grads, states, partitions, lr=0.02, mu=0.95, weight_decay=0.1):
+    """Muon with Nesterov momentum and an update RMS of 0.2, one matrix at a time."""
+    for name, X in params.items():
+        M = states[name]
+        M.mul_(mu).add_(grads[name])
+        U = (M * mu + grads[name]).bfloat16()
+        parts = partitions.get(name, (X.size(-2),))
+        updates = []
+        for part in U.split(parts, dim=-2):
+            O = newton_schulz(part.reshape(-1, *part.shape[-2:]), DEFAULT, 1e-8).reshape(part.shape)
+            updates.append(O.float() * lr * 0.2 * math.sqrt(max(part.shape[-2:])))
+        X.mul_(1 - lr * weight_decay).sub_(torch.cat(updates, dim=-2))
+
+
+def test_muon_matches_reference():
+    shapes = {"wide": (16, 48), "tall": (40, 24), "qkv": (48, 16), "experts": (3, 32, 16)}
+    partitions = {"qkv": (32, 8, 8), "experts": (16, 16)}
+    torch.manual_seed(0)
+    params = {name: torch.nn.Parameter(torch.randn(shape)) for name, shape in shapes.items()}
+    reference = {name: p.detach().clone() for name, p in params.items()}
+    states = {name: torch.zeros(shape) for name, shape in shapes.items()}
+    optimizer = Muon(
+        [dict(params=list(params.values()), algorithm="muon")],
+        lr=0.02,
+        weight_decay=0.1,
+        matrix_partitions={params[name]: parts for name, parts in partitions.items()},
+    )
+    for _ in range(3):
+        grads = {name: torch.randn(shape) for name, shape in shapes.items()}
+        for name, p in params.items():
+            p.grad = grads[name]
+        optimizer.step()
+        reference_muon_step(reference, grads, states, partitions)
+    for name, p in params.items():
+        torch.testing.assert_close(p.detach(), reference[name], atol=1e-5, rtol=1e-5)
+
+
+def test_subset_steps_match_the_full_step():
+    """An offloaded step updates one chunk of units at a time, after the trainer has moved on (zero_grad)."""
+
+    def train(chunked: bool) -> list[torch.Tensor]:
+        torch.manual_seed(0)
+        params = [torch.nn.Parameter(torch.randn(shape)) for shape in [(16, 48), (40, 24), (3, 32, 16), (24,)]]
+        optimizer = Muon(
+            [dict(params=params[:3], algorithm="muon"), dict(params=params[3:], algorithm="adamw")],
+            lr=0.02,
+            matrix_partitions={params[2]: (16, 16)},
+            bucket_bytes=2048,
+        )
+        for step in range(3):
+            for p in params:
+                p.grad = torch.randn(p.shape)
+            if chunked and step > 0:
+                subset_step = optimizer.begin_subset_step()
+                optimizer.zero_grad()
+                for unit in reversed(subset_step.units):
+                    subset_step.step_params(unit)
+            else:
+                optimizer.step()
+        assert [group["step"] for group in optimizer.param_groups] == [3, 3]
+        return [p.detach() for p in params]
+
+    for reference, chunked in zip(train(False), train(True)):
+        assert torch.equal(reference, chunked)
diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py
index 4c3b3abca86..2fe82757417 100644
--- a/tests/unit/train/test_state_offload.py
+++ b/tests/unit/train/test_state_offload.py
@@ -1,6 +1,5 @@
 from typing import NamedTuple
 
-import dion
 import pytest
 import torch
 from torch.distributed.checkpoint import load as dcp_load
@@ -8,6 +7,7 @@
 from torch.optim import Optimizer
 
 from prime_rl.trainer.ckpt import AppState
+from prime_rl.trainer.optim.muon import Muon
 from prime_rl.trainer.optim.state_offload import CPUOffloadOptimizer
 
 pytestmark = [pytest.mark.gpu]
@@ -119,7 +119,7 @@ def test_offload_reuses_pinned_buffers_across_steps() -> None:
 
 
 def test_checkpoint_resume_restores_param_groups(tmp_path) -> None:
-    """dion's Muon keeps its step count in the param groups, which an in-place offloaded load does not reach."""
+    """Muon keeps its step count in the param groups, which an in-place offloaded load does not reach."""
 
     def setup():
         model = make_model(0)
@@ -127,7 +127,7 @@ def setup():
             dict(params=[p for p in model.parameters() if p.ndim == 2], algorithm="muon"),
             dict(params=[p for p in model.parameters() if p.ndim == 1], algorithm="adamw"),
         ]
-        optimizer = dion.Muon(groups, lr=1e-2, fsdp_mesh_dim=0, world_mesh=None, adjust_lr="rms_norm")
+        optimizer = Muon(groups, lr=1e-2)
         return model, CPUOffloadOptimizer(optimizer)
 
     model, optimizer = setup()
diff --git a/uv.lock b/uv.lock
index 9d16078862d..dd82959f213 100644
--- a/uv.lock
+++ b/uv.lock
@@ -21,7 +21,6 @@ certifi = false
 deep-ep = false
 deep-gemm = false
 deep-select = false
-dion = false
 fastokens = false
 flash-attn-3 = false
 flash-mla = false
@@ -1464,16 +1463,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/50/3d/9373ad9c56321fdab5b41197068e1d8c25883b3fea29dd361f9b55116869/dill-0.4.0-py3-none-any.whl", hash = "sha256:44f54bf6412c2c8464c14e8243eb163690a9800dbe2c367330883b19c7561049", size = 119668, upload-time = "2025-04-16T00:41:47.671Z" },
 ]
 
-[[package]]
-name = "dion"
-version = "0.1.0"
-source = { git = "https://github.com/samsja/dion.git?rev=0c801b0#0c801b0ae7c3a98433255bff78994e43e859c738" }
-dependencies = [
-    { name = "numpy" },
-    { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } },
-    { name = "triton" },
-]
-
 [[package]]
 name = "dirhash"
 version = "0.5.0"
@@ -5238,7 +5227,6 @@ all = [
     { name = "deep-gemm", version = "2.5.0+891d57b", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "deep-select", version = "1.0.0+bfa4507.20261026.121543", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "deep-select", version = "1.0.0+bfa4507.20261026.121630", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "dion", marker = "sys_platform == 'linux'" },
     { name = "flash-attn", version = "2.8.3+cu130torch2.13", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.47/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "flash-attn", version = "2.8.3+cu130torch2.13", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "flash-attn-3", marker = "sys_platform == 'linux'" },
@@ -5300,7 +5288,6 @@ flash-attn-cute = [
 gpu = [
     { name = "deep-select", version = "1.0.0+bfa4507.20261026.121543", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" },
     { name = "deep-select", version = "1.0.0+bfa4507.20261026.121630", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
-    { name = "dion", marker = "sys_platform == 'linux'" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux'" },
     { name = "flash-mla", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "flashinfer-cubin", marker = "sys_platform == 'linux'" },
@@ -5361,7 +5348,6 @@ requires-dist = [
     { name = "deep-select", marker = "platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121543-cp312-cp312-linux_aarch64.whl" },
     { name = "deep-select", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'gpu'" },
     { name = "deep-select", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'gpu'", url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-b52bd9d/deep_select-1.0.0+bfa4507.20261026.121630-cp312-cp312-linux_x86_64.whl" },
-    { name = "dion", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/samsja/dion.git?rev=0c801b0" },
     { name = "fastapi", specifier = ">=0.115" },
     { name = "fastapi", marker = "extra == 'dashboard'", specifier = ">=0.115" },
     { name = "flash-attn", marker = "platform_machine == 'aarch64' and extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.48/flash_attn-2.8.3+cu130torch2.13-cp312-cp312-linux_aarch64.whl" },

From 7a1ab0f20b409dfb617425ef5b0d42df72fba047 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 04:31:07 +0000
Subject: [PATCH 08/12] fix(optim): per-tensor element-wise updates and flat
 normalization, so compiled graphs stay few

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/trainer/optim/muon.py | 65 +++++++++++++-----------------
 1 file changed, 28 insertions(+), 37 deletions(-)

diff --git a/src/prime_rl/trainer/optim/muon.py b/src/prime_rl/trainer/optim/muon.py
index 55931dd7829..86f5cc2e5ee 100644
--- a/src/prime_rl/trainer/optim/muon.py
+++ b/src/prime_rl/trainer/optim/muon.py
@@ -93,10 +93,10 @@ def _quack_gemm_symmetric():
 
 
 @torch.compile(fullgraph=True, dynamic=True)
-def _normalize(G: Tensor, safety: float, epsilon: float) -> Tensor:
-    """G / (safety·‖G‖_F + epsilon) per matrix, in bf16, as one fused reduction."""
-    X = G.to(torch.bfloat16)
-    return X / (X.norm(dim=(-2, -1), keepdim=True) * safety + epsilon)
+def _normalize(X: Tensor, safety: float, epsilon: float) -> Tensor:
+    """X / (safety·‖X‖ + epsilon) per row of X [B, m·n], as one fused reduction. Flat rows keep the compiled
+    graphs few: one per device and batch size class."""
+    return X / (X.norm(dim=-1, keepdim=True) * safety + epsilon)
 
 
 def newton_schulz(G: Tensor, schedule: NewtonSchulz, epsilon: float) -> Tensor:
@@ -104,7 +104,8 @@ def newton_schulz(G: Tensor, schedule: NewtonSchulz, epsilon: float) -> Tensor:
 
     The Gram matrix A is the smaller of XXᵀ and XᵀX. For tall matrices this iterates X <- aX + X(bA + cA²) on
     A = XᵀX, the transpose of the wide iteration, without transposing X."""
-    X = _normalize(G, schedule.safety, epsilon).to(schedule.dtype).contiguous()
+    X = G.to(torch.bfloat16).reshape(G.size(0), -1)
+    X = _normalize(X, schedule.safety, epsilon).reshape(G.shape).to(schedule.dtype)
     tall = X.size(-2) > X.size(-1)
     # QuACK needs 16-byte aligned rows
     gemm_symmetric = _quack_gemm_symmetric() if X.is_cuda and X.size(-1) % 8 == 0 and X.size(-2) % 8 == 0 else None
@@ -211,15 +212,6 @@ def _local(tensor: Tensor) -> Tensor:
     return tensor._local_tensor if isinstance(tensor, DTensor) else tensor
 
 
-def _by_group(matrices: list["_Matrix"], tensors: list[Tensor]) -> list[tuple[dict, list[tuple["_Matrix", Tensor]]]]:
-    """Pair matrices with their tensors, per param group, leaving out empty shards."""
-    groups = {}
-    for matrix, tensor in zip(matrices, tensors):
-        if tensor.numel():
-            groups.setdefault(id(matrix.group), (matrix.group, []))[1].append((matrix, tensor))
-    return list(groups.values())
-
-
 def _shape_class(matrix: "_Matrix") -> tuple:
     """Matrices with equal keys orthogonalize as one batch (see orthogonalize)."""
     return _local(matrix.param).shape[-1], matrix.partitions
@@ -375,23 +367,20 @@ def step_params(self, params: Collection[Tensor]) -> None:
         self.optimizer._run(self.plan.restrict(params), self.hyperparameters, self.grads)
 
 
-# The element-wise updates take a whole unit's tensors per call: one compiled call costs ~50 µs of host time.
+# One tensor per call, flattened where the arithmetic allows: a compiled function recompiles for every new list
+# length or tensor rank, and with fullgraph=True running out of recompiles is an error.
 @torch.compile(fullgraph=True, dynamic=True)
-def _momentum_(M: list[Tensor], G: list[Tensor], U: list[Tensor], mu: Tensor, nesterov: bool) -> None:
+def _momentum_(M: Tensor, G: Tensor, U: Tensor, mu: Tensor, nesterov: bool) -> None:
     """M <- μM + G and U <- μM + G (Nesterov) or M."""
-    for m, g, u in zip(M, G, U):
-        g = g.to(m.dtype)
-        m.mul_(mu).add_(g)
-        u.copy_(m * mu + g if nesterov else m)
+    G = G.to(M.dtype)
+    M.mul_(mu).add_(G)
+    U.copy_(M * mu + G if nesterov else M)
 
 
 @torch.compile(fullgraph=True, dynamic=True)
-def _muon_update_(
-    X: list[Tensor], U: list[Tensor], lr: Tensor, adjusted_lr: list[Tensor], weight_decay: Tensor
-) -> None:
-    for x, u, a in zip(X, U, adjusted_lr):
-        x.mul_(1 - lr * weight_decay)
-        x.sub_(u * a)
+def _muon_update_(X: Tensor, U: Tensor, lr: Tensor, adjusted_lr: Tensor, weight_decay: Tensor) -> None:
+    X.mul_(1 - lr * weight_decay)
+    X.sub_(U * adjusted_lr)
 
 
 @torch.compile(fullgraph=True, dynamic=True)
@@ -614,19 +603,21 @@ def _state(self, param: Tensor, *keys: str) -> list[Tensor]:
     def _momentum(
         self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: dict, grads: dict[Tensor, Tensor]
     ) -> None:
-        """M <- μM + G, and the input of the orthogonalization into ``updates``, one call per param group."""
-        for group, members in _by_group(matrices, updates):
-            h = hyperparameters[id(group)]
-            M = [self._state(m.param, "momentum")[0] for m, _ in members]
-            G = [grads[m.param] for m, _ in members]
-            _momentum_(M, G, [U for _, U in members], h["mu"], h["nesterov"])
+        """M <- μM + G, and the input of the orthogonalization into ``updates``."""
+        for matrix, U in zip(matrices, updates):
+            if U.numel():
+                (M,) = self._state(matrix.param, "momentum")
+                h = hyperparameters[id(matrix.group)]
+                _momentum_(M.view(-1), grads[matrix.param].reshape(-1), U.view(-1), h["mu"], h["nesterov"])
 
     def _update(self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: dict) -> None:
-        for group, members in _by_group(matrices, updates):
-            h = hyperparameters[id(group)]
-            X = [_local(m.param) for m, _ in members]
-            adjusted_lr = [m.lr_ratio * h["lr"] for m, _ in members]
-            _muon_update_(X, [U for _, U in members], h["lr"], adjusted_lr, h["weight_decay"])
+        for matrix, U in zip(matrices, updates):
+            if U.numel():
+                X, h = _local(matrix.param), hyperparameters[id(matrix.group)]
+                if isinstance(matrix.lr_ratio, Tensor):  # per row
+                    _muon_update_(X, U, h["lr"], matrix.lr_ratio * h["lr"], h["weight_decay"])
+                else:
+                    _muon_update_(X.view(-1), U.reshape(-1), h["lr"], matrix.lr_ratio * h["lr"], h["weight_decay"])
 
     def _update_local(self, chunk: list[_Matrix], hyperparameters: dict, grads: dict[Tensor, Tensor]) -> None:
         """Orthogonalize matrices that every rank holds whole, as one batch per matrix shape. The chunk's inputs

From 418a894d11f7e7bcf6cfe82da5cf185025594f56 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 04:35:01 +0000
Subject: [PATCH 09/12] fix(optim): compiled Muon helpers fall back to eager
 past the recompile limit

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/trainer/optim/muon.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/prime_rl/trainer/optim/muon.py b/src/prime_rl/trainer/optim/muon.py
index 86f5cc2e5ee..7e27b48bbec 100644
--- a/src/prime_rl/trainer/optim/muon.py
+++ b/src/prime_rl/trainer/optim/muon.py
@@ -92,7 +92,7 @@ def _quack_gemm_symmetric():
     return gemm_symmetric
 
 
-@torch.compile(fullgraph=True, dynamic=True)
+@torch.compile(dynamic=True)
 def _normalize(X: Tensor, safety: float, epsilon: float) -> Tensor:
     """X / (safety·‖X‖ + epsilon) per row of X [B, m·n], as one fused reduction. Flat rows keep the compiled
     graphs few: one per device and batch size class."""
@@ -367,9 +367,9 @@ def step_params(self, params: Collection[Tensor]) -> None:
         self.optimizer._run(self.plan.restrict(params), self.hyperparameters, self.grads)
 
 
-# One tensor per call, flattened where the arithmetic allows: a compiled function recompiles for every new list
-# length or tensor rank, and with fullgraph=True running out of recompiles is an error.
-@torch.compile(fullgraph=True, dynamic=True)
+# The compiled helpers take one tensor each, flattened where the arithmetic allows, so they compile few graphs. They
+# have no graph breaks, but are not fullgraph: past dynamo's recompile limit they run eagerly instead of raising.
+@torch.compile(dynamic=True)
 def _momentum_(M: Tensor, G: Tensor, U: Tensor, mu: Tensor, nesterov: bool) -> None:
     """M <- μM + G and U <- μM + G (Nesterov) or M."""
     G = G.to(M.dtype)
@@ -377,13 +377,13 @@ def _momentum_(M: Tensor, G: Tensor, U: Tensor, mu: Tensor, nesterov: bool) -> N
     U.copy_(M * mu + G if nesterov else M)
 
 
-@torch.compile(fullgraph=True, dynamic=True)
+@torch.compile(dynamic=True)
 def _muon_update_(X: Tensor, U: Tensor, lr: Tensor, adjusted_lr: Tensor, weight_decay: Tensor) -> None:
     X.mul_(1 - lr * weight_decay)
     X.sub_(U * adjusted_lr)
 
 
-@torch.compile(fullgraph=True, dynamic=True)
+@torch.compile(dynamic=True)
 def _adamw_(
     X: list[Tensor],
     G: list[Tensor],

From b7214e387a574cb086e30bd6d61015ef1c936873 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 04:58:59 +0000
Subject: [PATCH 10/12] fix(optim): refer to Muon param groups by index, as
 load_state_dict replaces them

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/trainer/optim/muon.py | 40 +++++++++++++++---------------
 tests/unit/train/test_muon.py      | 32 ++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/src/prime_rl/trainer/optim/muon.py b/src/prime_rl/trainer/optim/muon.py
index 7e27b48bbec..c6f15a150a3 100644
--- a/src/prime_rl/trainer/optim/muon.py
+++ b/src/prime_rl/trainer/optim/muon.py
@@ -292,7 +292,7 @@ class _Matrix:
     """A Muon parameter: its matrix partitions along dim -2, sharding and update scale."""
 
     param: Tensor
-    group: dict
+    group: int  # index into param_groups, whose dicts load_state_dict replaces
     partitions: tuple[int, ...]
     sharding: _Sharding | None
     lr_ratio: float | Tensor  # per local row for packed parameters
@@ -329,7 +329,7 @@ class _Plan:
 
     buckets: list[_Bucket]
     local: list[list[_Matrix]]
-    adamw: list[tuple[dict, list[Tensor]]]
+    adamw: list[tuple[int, list[Tensor]]]  # (param group index, params)
 
     @property
     def units(self) -> list[list[Tensor]]:
@@ -356,7 +356,7 @@ class _Step:
 
     optimizer: "Muon"
     plan: _Plan
-    hyperparameters: dict[int, dict]
+    hyperparameters: list[dict]  # per param group
     grads: dict[Tensor, Tensor]
 
     @property
@@ -479,16 +479,16 @@ def __init__(
 
         partitions = {param: tuple(sizes) for param, sizes in (matrix_partitions or {}).items()}
         self._matrices: dict[Tensor, _Matrix] = {}
-        for group in self.param_groups:
+        for index, group in enumerate(self.param_groups):
             if group["algorithm"] not in ("muon", "adamw"):
                 raise ValueError(f"Unknown algorithm {group['algorithm']}")
             if group["algorithm"] == "muon":
                 for param in group["params"]:
-                    self._matrices[param] = self._matrix(param, group, partitions.pop(param, (param.size(-2),)))
+                    self._matrices[param] = self._matrix(param, index, partitions.pop(param, (param.size(-2),)))
         if partitions:
             raise ValueError("matrix_partitions has parameters that are not Muon parameters of this optimizer")
 
-    def _matrix(self, param: Tensor, group: dict, partitions: tuple[int, ...]) -> _Matrix:
+    def _matrix(self, param: Tensor, group: int, partitions: tuple[int, ...]) -> _Matrix:
         if param.ndim < 2:
             raise ValueError(f"Muon parameters must be matrices, got shape {tuple(param.shape)}")
         if sum(partitions) != param.size(-2) or min(partitions) <= 0:
@@ -523,7 +523,7 @@ def begin_subset_step(self) -> _Step:
         key = tuple(map(id, params))
         if key not in self._plans:
             self._plans[key] = self._plan(params)
-        hyperparameters = {id(group): _hyperparameters(group) for group in self.param_groups}
+        hyperparameters = [_hyperparameters(group) for group in self.param_groups]
         return _Step(self, self._plans[key], hyperparameters, {param: _local(param.grad) for param in params})
 
     def _plan(self, params: list[Tensor]) -> _Plan:
@@ -561,14 +561,14 @@ def _plan(self, params: list[Tensor]) -> _Plan:
 
         selected = set(map(id, params))
         adamw = [
-            (group, [param for param in group["params"] if id(param) in selected])
-            for group in self.param_groups
+            (index, [param for param in group["params"] if id(param) in selected])
+            for index, group in enumerate(self.param_groups)
             if group["algorithm"] == "adamw"
         ]
         return _Plan(buckets, local, [(group, params) for group, params in adamw if params])
 
     @torch.no_grad()
-    def _run(self, plan: _Plan, hyperparameters: dict[int, dict], grads: dict[Tensor, Tensor]) -> None:
+    def _run(self, plan: _Plan, hyperparameters: list[dict], grads: dict[Tensor, Tensor]) -> None:
         # Bucket k+1 gathers and bucket k-1 scatters while bucket k is orthogonalized. Local chunks fill in.
         slots = len(plan.buckets) + 1
         local = [plan.local[i * len(plan.local) // slots : (i + 1) * len(plan.local) // slots] for i in range(slots)]
@@ -590,7 +590,7 @@ def _run(self, plan: _Plan, hyperparameters: dict[int, dict], grads: dict[Tensor
             self._apply(*scattering, hyperparameters)
 
         for group, params in plan.adamw:
-            self._adamw(params, hyperparameters[id(group)], grads)
+            self._adamw(params, hyperparameters[group], grads)
 
     def _state(self, param: Tensor, *keys: str) -> list[Tensor]:
         """State tensors, read at every use: an offloaded step swaps them for GPU copies between updates."""
@@ -601,25 +601,25 @@ def _state(self, param: Tensor, *keys: str) -> list[Tensor]:
         return [_local(state[key]) for key in keys]
 
     def _momentum(
-        self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: dict, grads: dict[Tensor, Tensor]
+        self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: list[dict], grads: dict[Tensor, Tensor]
     ) -> None:
         """M <- μM + G, and the input of the orthogonalization into ``updates``."""
         for matrix, U in zip(matrices, updates):
             if U.numel():
                 (M,) = self._state(matrix.param, "momentum")
-                h = hyperparameters[id(matrix.group)]
+                h = hyperparameters[matrix.group]
                 _momentum_(M.view(-1), grads[matrix.param].reshape(-1), U.view(-1), h["mu"], h["nesterov"])
 
-    def _update(self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: dict) -> None:
+    def _update(self, matrices: list[_Matrix], updates: list[Tensor], hyperparameters: list[dict]) -> None:
         for matrix, U in zip(matrices, updates):
             if U.numel():
-                X, h = _local(matrix.param), hyperparameters[id(matrix.group)]
+                X, h = _local(matrix.param), hyperparameters[matrix.group]
                 if isinstance(matrix.lr_ratio, Tensor):  # per row
                     _muon_update_(X, U, h["lr"], matrix.lr_ratio * h["lr"], h["weight_decay"])
                 else:
                     _muon_update_(X.view(-1), U.reshape(-1), h["lr"], matrix.lr_ratio * h["lr"], h["weight_decay"])
 
-    def _update_local(self, chunk: list[_Matrix], hyperparameters: dict, grads: dict[Tensor, Tensor]) -> None:
+    def _update_local(self, chunk: list[_Matrix], hyperparameters: list[dict], grads: dict[Tensor, Tensor]) -> None:
         """Orthogonalize matrices that every rank holds whole, as one batch per matrix shape. The chunk's inputs
         share one buffer in which matrices of one shape lie back to back, so each batch is a view of it."""
         matrices = sorted((m for m in chunk if _local(m.param).numel()), key=_shape_class)
@@ -629,12 +629,12 @@ def _update_local(self, chunk: list[_Matrix], hyperparameters: dict, grads: dict
         buffer = torch.empty(sum(_local(m.param).numel() for m in matrices), dtype=torch.bfloat16, device=device)
         updates = self._shard_views(buffer, matrices)
         self._momentum(matrices, updates, hyperparameters, grads)
-        epsilons = [m.group["epsilon"] for m in matrices]
+        epsilons = [self.param_groups[m.group]["epsilon"] for m in matrices]
         updates = orthogonalize(updates, [m.partitions for m in matrices], self.schedule, epsilons)
         self._update(matrices, updates, hyperparameters)
 
     def _gather(
-        self, bucket: _Bucket, hyperparameters: dict, grads: dict[Tensor, Tensor]
+        self, bucket: _Bucket, hyperparameters: list[dict], grads: dict[Tensor, Tensor]
     ) -> tuple[dist.Work, Tensor, Tensor]:
         """Momentum into this rank's shards, sent to their owners."""
         device = _local(bucket.sent[0].param).device
@@ -660,7 +660,7 @@ def _orthogonalize_and_scatter(
         matrices = [piece.view(shape) for piece, shape in zip(buffer.split([math.prod(s) for s in shapes]), shapes)]
         for j, (m, matrix) in enumerate(zip(bucket.owned, matrices)):
             torch.cat([shards[q][j].view(m.local_shape(q)) for q in range(world_size)], dim=m.sharding.dim, out=matrix)
-        epsilons = [m.group["epsilon"] for m in bucket.owned]
+        epsilons = [self.param_groups[m.group]["epsilon"] for m in bucket.owned]
         outputs = orthogonalize(matrices, [m.partitions for m in bucket.owned], self.schedule, epsilons)
         scattered = [
             output.narrow(m.sharding.dim, m.sharding.offsets[q], m.sharding.sizes[q]).reshape(-1)
@@ -674,7 +674,7 @@ def _orthogonalize_and_scatter(
         )
         return work, recv, send
 
-    def _apply(self, bucket: _Bucket, work: dist.Work, recv: Tensor, send: Tensor, hyperparameters: dict) -> None:
+    def _apply(self, bucket: _Bucket, work: dist.Work, recv: Tensor, send: Tensor, hyperparameters: list[dict]) -> None:
         work.wait()
         self._update(bucket.sent, self._shard_views(recv, bucket.sent), hyperparameters)
 
diff --git a/tests/unit/train/test_muon.py b/tests/unit/train/test_muon.py
index 923bb382ff1..8883ac92f15 100644
--- a/tests/unit/train/test_muon.py
+++ b/tests/unit/train/test_muon.py
@@ -1,3 +1,4 @@
+import copy
 import math
 
 import pytest
@@ -127,3 +128,34 @@ def train(chunked: bool) -> list[torch.Tensor]:
 
     for reference, chunked in zip(train(False), train(True)):
         assert torch.equal(reference, chunked)
+
+
+def test_resume_from_state_dict_matches_continuous_training():
+    """load_state_dict replaces the param group dicts: the optimizer must keep working on the new ones."""
+
+    def make():
+        torch.manual_seed(0)
+        params = [torch.nn.Parameter(torch.randn(shape)) for shape in [(16, 48), (3, 32, 16), (24,)]]
+        groups = [dict(params=params[:2], algorithm="muon"), dict(params=params[2:], algorithm="adamw")]
+        return params, Muon(groups, lr=0.02, matrix_partitions={params[1]: (16, 16)})
+
+    def train(params, optimizer, steps):
+        for step in steps:
+            torch.manual_seed(step)
+            for p in params:
+                p.grad = torch.randn(p.shape)
+            optimizer.step()
+
+    params, optimizer = make()
+    train(params, optimizer, range(2))
+    saved = ([p.detach().clone() for p in params], copy.deepcopy(optimizer.state_dict()))
+    train(params, optimizer, range(2, 4))
+
+    resumed_params, resumed = make()
+    with torch.no_grad():
+        for p, value in zip(resumed_params, saved[0]):
+            p.copy_(value)
+    resumed.load_state_dict(saved[1])
+    train(resumed_params, resumed, range(2, 4))
+    for p, q in zip(params, resumed_params):
+        assert torch.equal(p, q)

From 9e91c0265c7bbe44eaa79145a508f533b01fe5e9 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sun, 8 Nov 2026 16:23:36 +0000
Subject: [PATCH 11/12] perf(optim): make forward-overlapped offload step
 opt-in

model.offload_overlap_forward now defaults to false. The overlapped step and the
pre-broadcast wait have not been exercised in an RL run yet, nor with sparse weight
deltas; enable it explicitly until they have.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/scaling.md                                           | 2 +-
 packages/prime-rl-configs/src/prime_rl/configs/trainer.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/scaling.md b/docs/scaling.md
index 3e7f2cd5929..50fd73d8e52 100644
--- a/docs/scaling.md
+++ b/docs/scaling.md
@@ -156,7 +156,7 @@ Activation offloading still applies to tensors saved by autograd, but tensors re
 
 ### Optimizer Offloading
 
-State-only optimizer offload is enabled by default (`model.offload = "optimizer"`). With AdamW, the step streams the states through the GPU in chunks, and by default (`model.offload_overlap_forward = true`) those copies and updates run alongside the next forward, each FSDP module waiting only for its own parameters. 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
 
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 149c038bf54..cc7192873c0 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -297,8 +297,8 @@ 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 = True
-    """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. The math is unchanged; the step's gradients stay allocated until their chunk has updated."""
+    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."""

From bfb7dc43cb2a86ef6459ff52ae7641a6874e8ec1 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Tue, 10 Nov 2026 01:35:46 +0000
Subject: [PATCH 12/12] test(optim): one parameter per offload chunk in the
 parity tests; reuse the DTensor local helpers in _move_states

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/prime_rl/trainer/optim/state_offload.py | 22 ++++++++-------------
 tests/unit/train/test_state_offload.py      |  5 ++++-
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py
index 31fd1f40634..bc8c3bb0a1c 100644
--- a/src/prime_rl/trainer/optim/state_offload.py
+++ b/src/prime_rl/trainer/optim/state_offload.py
@@ -52,20 +52,14 @@ def _move_states(self, device: str):
         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()
 
diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py
index f1a5d1bf95c..331e37d894f 100644
--- a/tests/unit/train/test_state_offload.py
+++ b/tests/unit/train/test_state_offload.py
@@ -8,6 +8,7 @@
 from torch.optim import Optimizer
 
 from prime_rl.trainer.ckpt import AppState
+from prime_rl.trainer.optim import state_offload
 from prime_rl.trainer.optim.state_offload import CPUOffloadOptimizer
 
 pytestmark = [pytest.mark.gpu]
@@ -68,7 +69,9 @@ def state_data_ptrs(optimizer: CPUOffloadOptimizer) -> dict[tuple[int, str], int
 
 
 @pytest.fixture
-def trained_pair() -> TrainedPair:
+def trained_pair(monkeypatch: pytest.MonkeyPatch) -> TrainedPair:
+    # Every parameter gets a chunk of its own, so the tests cross chunk boundaries and reuse both staging slots.
+    monkeypatch.setattr(state_offload, "CHUNK_BYTES", 0)
     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)