perf(optim): chunked optimizer-state offload for Muon via a subset-step interface (builds on #4427)

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)
@@ -165,7 +166,17 @@ Every chunk runs the MoE's EP all-to-alls, so all ranks of an EP group must run
 
 ### Optimizer Offloading
 
-State-only optimizer offload is enabled by default (`model.offload = "optimizer"`). For full offload, set `model.offload = "full"`; this keeps BF16 compute weights on GPU and runs CPU optimizer chunks as gradients become ready during backward. Full offload only supports AdamW and SignSGD (`optim.type = "sign_sgd"`) and disables gradient clipping. SignSGD is stateless, so it halves the host RAM footprint versus AdamW (8 instead of 16 bytes per parameter: FP32 master + FP32 accumulated gradient, no moments).
+State-only optimizer offload is enabled by default (`model.offload = "optimizer"`). With AdamW 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, set `model.offload_overlap_forward = true` to also 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).
+
+### 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
 

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

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

pyproject.toml

@@ -61,7 +61,6 @@ gpu = [
     "vllm>=0.33.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). vLLM 0.33 requires >=1.31; 1.32 changed the backward's
     # signature (preallocated dq/dkv/d_sink and workspace), which our call sites and `cudnn_dsa_patch` predate.
@@ -246,7 +245,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
 nixl-cu13 = false
 torchao = false
@@ -268,7 +266,6 @@ torchao = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-86f0ad0/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.0-86f0ad0/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" }

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:

src/prime_rl/trainer/optim/__init__.py

@@ -1,15 +1,13 @@
 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
 from prime_rl.trainer.models.fusions import get_model_packed_parameters
 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,38 +19,12 @@
 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]],
     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]:
@@ -77,7 +49,6 @@ def setup_optimizer(
     optimizer = _create_optimizer(
         config,
         optimizer_named_params,
-        parallel_dims,
         fused_adamw=config.type == "adamw" and not cpu_offload,
         model=model,
     )
@@ -94,15 +65,23 @@ def setup_optimizer(
 
     if cpu_offload:
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
-        return CPUOffloadOptimizer(optimizer), None
+        # Muon steps synchronously: its step's host time would land on the forward's launch thread, and
+        # overlapping it with the forward has not been measured to pay off.
+        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]],
-    parallel_dims: ParallelDims,
     lr: float | None = None,
     fused_adamw: bool = False,
     model: nn.Module | None = None,
@@ -126,7 +105,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 +117,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 +157,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 +176,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