From e67deb34e019b98c8a2b8441261d91c9a42acd38 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sun, 4 Oct 2026 18:28:49 +0000 Subject: [PATCH] feat(optim): symmetric-GEMM Newton-Schulz and in-place expert Muon updates --- src/prime_rl/trainer/optim/__init__.py | 3 +- src/prime_rl/trainer/optim/muon.py | 101 +++++++++++++++++++++++++ tests/unit/train/test_muon.py | 30 ++++++++ 3 files changed, 133 insertions(+), 1 deletion(-) create mode 100644 src/prime_rl/trainer/optim/muon.py create mode 100644 tests/unit/train/test_muon.py diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 33755114091..47a886c296e 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,6 +8,7 @@ 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, newton_schulz_symmetric from prime_rl.trainer.optim.offload import ( FullCPUOffloadOptimizer, GradientOffloadManager, @@ -235,6 +235,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, + newton_schulz_func=newton_schulz_symmetric, ) # 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..cad70b0f3a3 --- /dev/null +++ b/src/prime_rl/trainer/optim/muon.py @@ -0,0 +1,101 @@ +import torch +from dion import Muon as DionMuon +from dion.muon import muon_update_batch_async, zeropower_via_newtonschulz5 +from dion.opt_utils import AsyncTask +from torch import Tensor +from torch.distributed.tensor import DTensor + +# The quintic coefficients of Dion's dense reference, zeropower_via_newtonschulz5. +NS_COEFFICIENTS = ( + (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), +) + + +def get_quack_gemm_symmetric(): + """QuACK's symmetric GEMM, or None when it is not installed or the GPU is unsupported (it needs SM80+).""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8: + return None + try: + from quack.gemm_interface import gemm_symmetric + except ImportError: + return None + return gemm_symmetric + + +@torch.no_grad() +def newton_schulz_symmetric(G: Tensor, epsilon: float | Tensor = 1e-7) -> Tensor: + """Dion's Newton-Schulz, with the two symmetric products (X Xᵀ and A·A) on QuACK's symmetric GEMM. + + A symmetric GEMM computes one triangle of the output and mirrors it, so it does about half the + work of a dense one. Falls back to the dense reference where QuACK cannot run. + """ + gemm_symmetric = get_quack_gemm_symmetric() if G.is_cuda else None + # QuACK needs 16-byte aligned rows, i.e. both matrix dimensions divisible by 8 in bf16. + if gemm_symmetric is None or G.size(-1) % 8 or G.size(-2) % 8: + return zeropower_via_newtonschulz5(G, epsilon=epsilon) + + X = G.to(dtype=torch.bfloat16) + transposed = G.size(-2) > G.size(-1) + if transposed: + X = X.mT + X = (X / (X.norm(dim=(-2, -1), keepdim=True) + epsilon)).contiguous() + addmm = torch.baddbmm if X.ndim == 3 else torch.addmm + for a, b, c in NS_COEFFICIENTS: + A = gemm_symmetric(X, X.mT) + B = gemm_symmetric(A, A, C=A, alpha=c, beta=b) + X = addmm(X, B, X, beta=a) + return X.mT if transposed else X + + +def holds_whole_matrices(param: Tensor, flatten: bool) -> bool: + """Whether every shard of ``param`` is a stack of whole matrices: a 3D+ batch sharded only on dim 0.""" + return ( + isinstance(param, DTensor) + and param.ndim >= 3 + and not flatten + and all(placement.dim == 0 for placement in param.placements if placement.is_shard()) + ) + + +class Muon(DionMuon): + """Dion's Muon, except that sharded parameters whose shards are whole matrices update in place. + + Dion all-to-alls every sharded parameter so that one rank holds the whole tensor (and, under EP, + zero-pads the rank's expert shard to the full expert count). An expert stack + ``[experts, fan_out, fan_in]`` sharded only over the expert axis already holds whole matrices on + every rank, and Newton-Schulz treats the leading axis as a batch, so it needs no communication. + """ + + def _create_muon_tasks(self, param_groups: list[dict], algo_name: str = "muon"): + def split(group: dict, local: bool) -> dict: + params = [p for p in group["params"] if holds_whole_matrices(p, group["flatten"]) == local] + return {**group, "params": params} + + for group in param_groups: + for param in split(group, local=True)["params"]: + if param.grad is None: + continue + state = self._get_or_initialize_state(param, algo_name) + yield AsyncTask( + muon_update_batch_async( + X=[param], + G=[param.grad], + M=[state["momentum"]], + lr=torch.tensor(group["lr"]), + momentum=torch.tensor(group["mu"]), + weight_decay=torch.tensor(group["weight_decay"]), + epsilon=torch.tensor(group["epsilon"]), + nesterov=group["nesterov"], + flatten=group["flatten"], + adjust_lr=group["adjust_lr"], + device_rank=0, + world_size=1, + newton_schulz_func=self._newton_schulz_func, + partitions=self._matrix_partitions.get(param), + ) + ) + yield from super()._create_muon_tasks([split(group, local=False) for group in param_groups], algo_name) diff --git a/tests/unit/train/test_muon.py b/tests/unit/train/test_muon.py new file mode 100644 index 00000000000..e009aa8a257 --- /dev/null +++ b/tests/unit/train/test_muon.py @@ -0,0 +1,30 @@ +import pytest +import torch +from dion.muon import zeropower_via_newtonschulz5 + +from prime_rl.trainer.optim.muon import NS_COEFFICIENTS, get_quack_gemm_symmetric, newton_schulz_symmetric + +pytestmark = [pytest.mark.gpu] + + +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 NS_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.parametrize("shape", [(1024, 3072), (3072, 1024), (8, 768, 2048)]) +def test_newton_schulz_symmetric_is_as_accurate_as_dense(shape): + if get_quack_gemm_symmetric() is None: + pytest.skip("QuACK symmetric GEMM is unavailable") + 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_symmetric(G)) <= 1.25 * error(zeropower_via_newtonschulz5(G))