From 44303a0bba47b963afed8d8074f57cdfe12e4f46 Mon Sep 17 00:00:00 2001 From: samsja <55492238+samsja@users.noreply.github.com> Date: Mon, 28 Sep 2026 19:47:18 -0700 Subject: [PATCH] fix(optim): stream Muon gradients and state per batch --- examples/advanced/glm-5.3/sft/h200/base.toml | 4 +- .../src/prime_rl/configs/trainer.py | 6 ++ src/prime_rl/trainer/optim/__init__.py | 11 ++- src/prime_rl/trainer/optim/muon.py | 71 ++++++++++++++++++ src/prime_rl/trainer/optim/state_offload.py | 73 ++++++++++++++++++- 5 files changed, 159 insertions(+), 6 deletions(-) create mode 100644 src/prime_rl/trainer/optim/muon.py diff --git a/examples/advanced/glm-5.3/sft/h200/base.toml b/examples/advanced/glm-5.3/sft/h200/base.toml index 10a305bf181..84ad618a955 100644 --- a/examples/advanced/glm-5.3/sft/h200/base.toml +++ b/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" 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 70100bee117..62ccedc8b24 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/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.""" diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 78496d443c5..f83a840321f 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/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. diff --git a/src/prime_rl/trainer/optim/muon.py b/src/prime_rl/trainer/optim/muon.py new file mode 100644 index 00000000000..b621c9c9135 --- /dev/null +++ b/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) diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py index 15a7f1eca2c..cf312029aa7 100644 --- a/src/prime_rl/trainer/optim/state_offload.py +++ b/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