diff --git a/docs/scaling.md b/docs/scaling.md index 30e8fb45d1d..57aed173269 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) @@ -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 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 bb93d255ee5..bbee0252a53 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/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): diff --git a/pyproject.toml b/pyproject.toml index be974a44bcf..5cbd4bec95b 100644 --- a/pyproject.toml +++ b/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" } diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py index e0f74f73a90..1cbfbb92edc 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..96da5057ef2 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,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 diff --git a/src/prime_rl/trainer/optim/muon.py b/src/prime_rl/trainer/optim/muon.py new file mode 100644 index 00000000000..c6f15a150a3 --- /dev/null +++ b/src/prime_rl/trainer/optim/muon.py @@ -0,0 +1,690 @@ +"""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(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.""" + return X / (X.norm(dim=-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 = 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 + 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 _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: 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 + + 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[int, list[Tensor]]] # (param group index, params) + + @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: list[dict] # per param group + 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 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) + M.mul_(mu).add_(G) + U.copy_(M * mu + G if nesterov else M) + + +@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(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 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, 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: 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: + 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 = [_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 = [ + (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: 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)] + 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[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: 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[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: list[dict]) -> None: + for matrix, U in zip(matrices, updates): + if U.numel(): + 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: 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) + 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 = [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: 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 + 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 = [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) + 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: list[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/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py index 15a7f1eca2c..3f2ab42d74c 100644 --- a/src/prime_rl/trainer/optim/state_offload.py +++ b/src/prime_rl/trainer/optim/state_offload.py @@ -1,10 +1,137 @@ 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 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. +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: + 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 + + +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 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``. ``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, + 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() + 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.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 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.subset_step.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): @@ -15,15 +142,90 @@ 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, overlap_model: nn.Module | None = None): + """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.pin_memory = pin_memory + 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 and self._subset_step is not None: + self._install_unshard_waits(overlap_model) + + def _install_unshard_waits(self, model: nn.Module) -> None: + 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"]] + 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, 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) + + 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: - 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) @@ -34,23 +236,18 @@ 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(): - 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() @@ -62,6 +259,12 @@ def step(self, closure=None): self._initialized = True return result + # A parameter without state yet (no grad on any earlier step) is initialized by the full path below. + 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() + # Move states to GPU self._move_states("cuda") @@ -73,7 +276,71 @@ def step(self, closure=None): return result + 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, 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)) + return chunks, max_used + + 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 in + the forward order of their FSDP groups, enqueued as the next forward reaches them. + + 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 + 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, 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, 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 + 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/optim/subset_step.py b/src/prime_rl/trainer/optim/subset_step.py new file mode 100644 index 00000000000..831cb1bddcd --- /dev/null +++ b/src/prime_rl/trainer/optim/subset_step.py @@ -0,0 +1,93 @@ +"""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`` (e.g. +prime-rl's ``Muon``); torch ``AdamW`` gets an adapter 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 torch +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) + + +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) + return None diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 16542da0758..82236262567 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), ) @@ -683,6 +684,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 bde4c6295ef..845f347ef37 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), ) @@ -556,6 +557,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 @@ -692,6 +694,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: diff --git a/tests/unit/train/test_muon.py b/tests/unit/train/test_muon.py new file mode 100644 index 00000000000..8883ac92f15 --- /dev/null +++ b/tests/unit/train/test_muon.py @@ -0,0 +1,161 @@ +import copy +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) + + +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) diff --git a/tests/unit/train/test_state_offload.py b/tests/unit/train/test_state_offload.py index 4c3b3abca86..23abcca3b6c 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,8 @@ 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.muon import Muon 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) @@ -118,8 +121,27 @@ def test_offload_reuses_pinned_buffers_across_steps() -> None: 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) + + 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 +149,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() @@ -140,3 +162,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: + """Muon steps chunk by chunk (a bucket of sharded matrices, a batch of local matrices 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 = Muon(groups, lr=1e-2) + 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] diff --git a/uv.lock b/uv.lock index eb232a6a7a5..8a405ef9038 100644 --- a/uv.lock +++ b/uv.lock @@ -20,7 +20,6 @@ exclude-newer-span = "P7D" certifi = false deep-ep = false deep-select = false -dion = false fastokens = false flash-attn-3 = false flashinfer-python = false @@ -1469,16 +1468,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.14.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, - { name = "triton" }, -] - [[package]] name = "dirhash" version = "0.5.0" @@ -5276,7 +5265,6 @@ all = [ { name = "deep-ep", version = "1.2.1+29d31c0", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-86f0ad0/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "deep-select", version = "1.0.0+bfa4507.20261109.170213", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-86f0ad0/deep_select-1.0.0+bfa4507.20261109.170213-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, { name = "deep-select", version = "1.0.0+bfa4507.20261109.170237", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-86f0ad0/deep_select-1.0.0+bfa4507.20261109.170237-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.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "flash-attn", version = "2.8.3+cu130torch2.14", source = { url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, { name = "flash-attn-3", marker = "sys_platform == 'linux'" }, @@ -5335,7 +5323,6 @@ flash-attn-cute = [ gpu = [ { name = "deep-select", version = "1.0.0+bfa4507.20261109.170213", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-86f0ad0/deep_select-1.0.0+bfa4507.20261109.170213-cp312-cp312-linux_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, { name = "deep-select", version = "1.0.0+bfa4507.20261109.170237", source = { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-86f0ad0/deep_select-1.0.0+bfa4507.20261109.170237-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 = "flashinfer-cubin", marker = "sys_platform == 'linux'" }, { name = "flashinfer-jit-cache", version = "0.7.2.post1+cu130", source = { url = "https://github.com/flashinfer-ai/flashinfer/releases/download/v0.7.2.post1/flashinfer_jit_cache-0.7.2.post1%2Bcu130-cp39-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, @@ -5392,7 +5379,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-86f0ad0/deep_select-1.0.0+bfa4507.20261109.170213-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-86f0ad0/deep_select-1.0.0+bfa4507.20261109.170237-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.10.1/flash_attn-2.8.3+cu130torch2.14-cp312-cp312-linux_aarch64.whl" },