feat(optim): own Muon implementation, drop the dion dependency

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

@@ -437,10 +437,16 @@ 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."""
+
+    newton_schulz_coefficients: Literal["muon", "deepseek_v4"] = "muon"
+    """Newton-Schulz schedule. ``muon``: 5 steps tuned for speed (Keller Jordan's, as in dion). ``deepseek_v4``: 8 steps of (3.4445, -4.7750, 2.0315), then 2 of (2, -1.5, 0.5) (DeepSeek-V4 report, section 2.4); twice the cost, singular values closer to 1."""
+
+    update_rms: float = Field(0.2, gt=0)
+    """RMS of a Muon update per unit of learning rate: each matrix's update is scaled by ``update_rms * sqrt(max(fan_out, fan_in))`` (Moonlight's matching of AdamW: 0.2; DeepSeek-V4: 0.18)."""
 
 
 class SignSGDConfig(BaseOptimizerConfig):

pyproject.toml

@@ -59,7 +59,6 @@ gpu = [
     "mooncake-transfer-engine>=0.3.10.post2 ; sys_platform == 'linux'",
     "ring-flash-attn>=0.1.8 ; 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'",
@@ -252,7 +251,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
@@ -275,7 +273,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" }

src/prime_rl/trainer/optim/__init__.py

@@ -1,14 +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.optim.base import OffloadOptimizer as OffloadOptimizer
 from prime_rl.trainer.optim.base import OptimizerLike
+from prime_rl.trainer.optim.muon import Muon, shard_group
 from prime_rl.trainer.optim.offload import (
     FullCPUOffloadOptimizer,
     GradientOffloadManager,
@@ -20,25 +19,20 @@
 from prime_rl.utils.logger import get_logger
 
 
-def _warmup_muon_mesh(mesh: DeviceMesh) -> None:
+def _warmup_muon_group(group: dist.ProcessGroup) -> 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
+    # Use a bulk list all-to-all, matching Muon'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]
@@ -124,7 +118,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,
@@ -136,7 +130,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:
@@ -171,70 +164,36 @@ 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")
-    )
-    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,
-            )
-        )
-    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
+    # The group layout matches dion's Muon, whose checkpoints this optimizer loads.
+    param_groups = [dict(params=muon_params)]
+    param_groups += [dict(params=params) for params in (expert_params, router_params) if params]
+    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
     # matrices it replaces would. Packed biases and frozen parameters are not Muon's.
     matrix_partitions = {}
     if model is not None:
-        muon_params = {p for group in param_groups if group["algorithm"] == "muon" for p in group["params"]}
+        matrices = {p for group in param_groups[:-1] for p in group["params"]}
         for packed_info in get_model_packed_parameters(model):
             partitions = packed_info.spec.muon_matrix_partitions(packed_info.parameter)
-            if partitions is not None and packed_info.parameter in muon_params:
+            if partitions is not None and packed_info.parameter in matrices:
                 matrix_partitions[packed_info.parameter] = partitions
 
     optimizer = Muon(
-        params=param_groups,
-        matrix_partitions=matrix_partitions,
+        param_groups,
         lr=lr,
         mu=config.mu,
         nesterov=config.nesterov,
         betas=(config.betas1, config.betas2),
         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,
+        matrix_partitions=matrix_partitions,
+        newton_schulz_coefficients=config.newton_schulz_coefficients,
+        update_rms=config.update_rms,
     )
-    # 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"))
+    # Before Muon's first step: warm every group its all-to-alls run over.
+    groups = dict.fromkeys(shard_group(p)[1] for group in param_groups[:-1] for p in group["params"])
+    for group in groups:
+        if group is not None:
+            _warmup_muon_group(group)
     return optimizer