fix(optim): stream Muon gradients and state for 8-node GLM-5.3 SFT

examples/advanced/glm-5.3/sft/h200/base.toml

@@ -19,9 +19,11 @@ freq = 1
 max_inflight_activations = 5
 
 [optim]
-type = "sign_sgd"
+type = "muon"
 lr = 5e-5
 weight_decay = 0.01
+max_concurrent_tasks = 1
+stream_cpu_offload = true
 
 [scheduler]
 type = "linear"

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

@@ -527,6 +527,12 @@ class MuonConfig(BaseOptimizerConfig):
     mu: float = Field(0.95, ge=0)
     """Momentum factor for the Muon algorithm."""
 
+    max_concurrent_tasks: int = Field(3, ge=1)
+    """Maximum number of concurrent Muon update tasks. Lower values reduce optimizer-step peak memory."""
+
+    stream_cpu_offload: bool = False
+    """Stream gradients and optimizer state through CPU one update batch at a time."""
+
     betas1: float = Field(0.9, ge=0)
     """β1 for the AdamW/Lion sub-optimizer used on non-Muon params."""
 

src/prime_rl/trainer/optim/__init__.py

@@ -1,6 +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 SGD, AdamW, Optimizer
@@ -9,12 +8,13 @@
 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
 from prime_rl.trainer.optim.offload import (
     FullCPUOffloadOptimizer,
     GradientOffloadManager,
     _create_cpu_master_weights,
 )
-from prime_rl.trainer.optim.state_offload import CPUOffloadOptimizer
+from prime_rl.trainer.optim.state_offload import CPUOffloadOptimizer, StreamingMuonCPUOffloadOptimizer
 from prime_rl.trainer.parallel_dims import ParallelDims
 from prime_rl.trainer.sign_sgd import SignSGD
 from prime_rl.utils.logger import get_logger
@@ -98,7 +98,13 @@ def setup_optimizer(
         )
         return optimizer, optimizer._gradient_manager
 
+    if config.type == "muon" and config.stream_cpu_offload and not cpu_offload:
+        raise ValueError("Muon state streaming requires optim_cpu_offload=true")
+
     if cpu_offload:
+        if config.type == "muon" and config.stream_cpu_offload:
+            get_logger().info("Streaming Muon gradients and optimizer state from CPU one update batch at a time")
+            return StreamingMuonCPUOffloadOptimizer(optimizer), None
         get_logger().info("Wrapping optimizer with CPUOffloadOptimizer for optimizer state CPU offloading")
         return CPUOffloadOptimizer(optimizer), None
 
@@ -240,6 +246,7 @@ def muon_enabled(n, p):
         distributed_mesh=distributed_mesh,
         world_mesh=parallel_dims.world_mesh,
         fsdp_mesh_dim=1 if parallel_dims.dp_replicate_enabled else 0,
+        max_concurrent_tasks=config.max_concurrent_tasks,
     )
     # Keep both warm-ups after Muon construction and before its first step. The
     # main and expert groups establish independent NCCL peer connections.

src/prime_rl/trainer/optim/muon.py

@@ -0,0 +1,71 @@
+from itertools import chain
+
+import torch
+from dion import Muon as DionMuon
+from dion.opt_utils import AsyncRuntime, create_param_batches
+
+from prime_rl.trainer.optim.state_offload import StreamingMuonCPUOffloadOptimizer
+
+
+class Muon(DionMuon):
+    """Dion Muon with configurable optimizer-task concurrency."""
+
+    def __init__(self, *args, max_concurrent_tasks: int = 3, **kwargs):
+        super().__init__(*args, **kwargs)
+        if max_concurrent_tasks < 1:
+            raise ValueError("max_concurrent_tasks must be at least 1")
+        self.max_concurrent_tasks = max_concurrent_tasks
+
+    @torch.no_grad()
+    def step(self, closure=None, state_offloader: StreamingMuonCPUOffloadOptimizer | None = None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        muon_groups = []
+        lion_groups = []
+        adamw_groups = []
+        for group in self.param_groups:
+            group["step"] += 1
+            match group["algorithm"]:
+                case "muon":
+                    muon_groups.append(group)
+                case "lion":
+                    lion_groups.append(group)
+                case "adamw":
+                    adamw_groups.append(group)
+                case algorithm:
+                    raise ValueError(f"Unknown algorithm: {algorithm}")
+
+        if state_offloader is None:
+            tasks = chain(
+                self._create_muon_tasks(muon_groups),
+                self._create_lion_tasks(lion_groups),
+                self._create_adamw_tasks(adamw_groups),
+            )
+            AsyncRuntime(tasks, max_concurrent_tasks=self.max_concurrent_tasks).run()
+        else:
+            for group in muon_groups:
+                _, world_size, _, _ = self._get_group_mesh_info(group)
+                params = [param for param in group["params"] if state_offloader.has_gradient(param)]
+                for batch in create_param_batches(params, world_size, self._matrix_partitions):
+                    self._run_streamed_batch(group, batch, self._create_muon_tasks, state_offloader)
+
+            for groups, create_tasks in (
+                (lion_groups, self._create_lion_tasks),
+                (adamw_groups, self._create_adamw_tasks),
+            ):
+                for group in groups:
+                    for param in group["params"]:
+                        if state_offloader.has_gradient(param):
+                            self._run_streamed_batch(group, [param], create_tasks, state_offloader)
+        return loss
+
+    def _run_streamed_batch(self, group, params, create_tasks, state_offloader):
+        state_offloader.load_gradients(params)
+        state_offloader._move_states("cuda", params)
+        batch_group = {**group, "params": params}
+        AsyncRuntime(create_tasks([batch_group]), max_concurrent_tasks=1).run()
+        state_offloader.release_gradients(params)
+        state_offloader._move_states("cpu", params)

src/prime_rl/trainer/optim/state_offload.py

@@ -32,9 +32,13 @@ def _offload_to_cpu(self, param: torch.Tensor, key: str, src: torch.Tensor) -> t
             buffer.copy_(src, non_blocking=True)
         return buffer
 
-    def _move_states(self, device: str):
-        """Move optimizer states to CPU or back to GPU (matching each parameter's device)."""
-        for param in self.optimizer.state:
+    def _move_states(self, device: str, params=None):
+        """Move optimizer states for the selected parameters to CPU or GPU."""
+        if params is None:
+            params = self.optimizer.state
+        for param in params:
+            if param not in self.optimizer.state:
+                continue
             state = self.optimizer.state[param]
             for key, value in state.items():
                 if isinstance(value, DTensor):
@@ -120,3 +124,66 @@ def finish_checkpoint_save(self) -> None:
 
     def finish_checkpoint_load(self) -> None:
         self._initialized = True
+
+
+class StreamingMuonCPUOffloadOptimizer(CPUOffloadOptimizer):
+    """Keep Muon state and completed gradients on CPU between update batches."""
+
+    def __init__(self, optimizer: Optimizer, pin_memory: bool = True):
+        super().__init__(optimizer, pin_memory=pin_memory)
+        self._cpu_gradients: dict[torch.Tensor, torch.Tensor] = {}
+
+    def _offload_gradients(self) -> None:
+        gpu_gradients = []
+        for group in self.optimizer.param_groups:
+            for param in group["params"]:
+                grad = param.grad
+                if grad is None:
+                    continue
+                if isinstance(grad, DTensor):
+                    cpu_grad = copy.copy(grad)
+                    cpu_grad._local_tensor = self._offload_to_cpu(param, "_gradient", grad._local_tensor)
+                else:
+                    cpu_grad = self._offload_to_cpu(param, "_gradient", grad)
+                self._cpu_gradients[param] = cpu_grad
+                param.grad = None
+                gpu_gradients.append(grad)
+        if gpu_gradients and torch.cuda.is_initialized():
+            torch.cuda.synchronize()
+
+    def has_gradient(self, param: torch.Tensor) -> bool:
+        return param in self._cpu_gradients
+
+    def load_gradients(self, params) -> None:
+        for param in params:
+            grad = self._cpu_gradients[param]
+            if isinstance(grad, DTensor):
+                gpu_grad = copy.copy(grad)
+                gpu_grad._local_tensor = grad._local_tensor.to(param.device, non_blocking=True)
+            else:
+                gpu_grad = grad.to(param.device, non_blocking=True)
+            param.grad = gpu_grad
+
+    def release_gradients(self, params) -> None:
+        for param in params:
+            param.grad = None
+            del self._cpu_gradients[param]
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        self._offload_gradients()
+        result = self.optimizer.step(state_offloader=self)
+        self._initialized = True
+        return loss if closure is not None else result
+
+    def state_dict(self):
+        return self.optimizer.state_dict()
+
+    def prepare_checkpoint_save(self) -> None:
+        pass
+
+    def finish_checkpoint_save(self) -> None:
+        pass