From 381d8b7aa059ee51686a959a08ad8926543a22d6 Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Thu, 2 Apr 2026 21:23:16 -0700 Subject: [PATCH 1/7] Add Gram Newton-Schulz iteration for Muon optimizer Integrate Gram Newton-Schulz (Gram NS) as the default orthogonalization method for Muon, with a configurable ns_method switch to fall back to standard NS when needed (e.g., for debugging convergence issues). Gram NS iterates on the small square Gram matrix R = X @ X.T (n x n) instead of the full rectangular X (n x m), reducing FLOPs by ~50% for typical transformer weight matrices (aspect ratio ~5). It uses fp16 instead of bf16 for better numerical precision at the same compute cost, with a restart at iteration 2 for half-precision stability. Benchmark results on A100: - (2048, 11059): 2.25x GPU speedup, 1.85x CPU speedup - (3584, 19353): 2.07x GPU speedup, 1.35x CPU speedup - Falls back to standard NS for square matrices (no FLOP advantage) Usage: set ns_method in DeepSpeed config: {"optimizer": {"type": "muon", "params": {"ns_method": "gram"}}} Use "standard" to disable Gram NS and revert to original behavior. Reference: https://arxiv.org/abs/2503.02022 Signed-off-by: Ma, Guokai --- deepspeed/runtime/engine.py | 2 +- deepspeed/runtime/zero/muon/original_muon.py | 122 +++++++++++++++++-- deepspeed/runtime/zero/stage3.py | 4 +- deepspeed/runtime/zero/stage_1_and_2.py | 6 +- tests/unit/ops/muon/test_muon.py | 91 ++++++++++++++ 5 files changed, 210 insertions(+), 15 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 5fc07b3a7238..1fc71f1a8e9e 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -1754,7 +1754,7 @@ def _configure_basic_optimizer(self, model_parameters): param_groups = [] if muon_params: accepted_parameters = dict() - for key in ["lr", "momentum", "weight_decay", "muon_lr"]: + for key in ["lr", "momentum", "weight_decay", "muon_lr", "ns_method"]: if key in optimizer_parameters: if key == "muon_lr": # muon_lr will override lr accepted_parameters['lr'] = optimizer_parameters[key] diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index f4dc7a0909bb..83c118e82a16 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -30,6 +30,7 @@ import torch import deepspeed.comm as dist # replace torch's distributed package with deepspeed.comm to resolve deepspeed check from deepspeed.runtime import compiler +from deepspeed.accelerator import get_accelerator @compiler.compile() @@ -63,12 +64,89 @@ def zeropower_via_newtonschulz5(G, steps: int): @compiler.compile() -def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True): +def zeropower_via_gram_newtonschulz(G, steps: int): + """ + Gram Newton-Schulz iteration for orthogonalization. + + Mathematically equivalent to standard Newton-Schulz but iterates on the + small square Gram matrix R = X @ X.T (n x n) instead of the full rectangular + X (n x m). This reduces FLOPs significantly when m >> n (typical for + transformer weight matrices with aspect ratio ~5). + + Uses fp16 instead of bf16 for better numerical precision at the same + compute cost. Includes a restart at iteration 2 to maintain stability + in half-precision. + + Falls back to standard Newton-Schulz for square matrices (n == m) + where there is no FLOP advantage. + + Reference: https://arxiv.org/abs/2503.02022 + """ + assert G.ndim >= 2 + a, b, c = (3.4445, -4.7750, 2.0315) + # Use fp16 on GPU for better precision than bf16; fp32 on CPU for stability + compute_dtype = torch.float16 if get_accelerator().on_accelerator(G) else torch.float32 + X = G.to(compute_dtype) + if G.size(-2) > G.size(-1): + X = X.mT + + n, m = X.size(-2), X.size(-1) + + X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) + + # For square matrices, no FLOP advantage; use standard iteration + if m <= n: + for _ in range(steps): + A = X @ X.mT + B = b * A + c * A @ A + X = a * X + B @ X + if G.size(-2) > G.size(-1): + X = X.mT + return X + + # Gram NS: iterate on R = X @ X.T (n x n) instead of X (n x m) + R = X @ X.mT + I = torch.eye(n, device=X.device, dtype=X.dtype) + Q = None + restart_at = 2 + + for i in range(steps): + if i == restart_at and i != 0: + X = Q @ X + R = X @ X.mT + Q = None + + Z = b * R + c * R @ R + + if Q is None: + Q = Z + a * I + else: + Q = a * Q + Z @ Q + + if i < steps - 1 and (i + 1) != restart_at: + RZ = a * R + Z @ R + R = a * RZ + Z @ RZ + + X = Q @ X + + if G.size(-2) > G.size(-1): + X = X.mT + return X + + +NS_METHODS = {"standard", "gram"} + + +@compiler.compile() +def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram"): momentum.lerp_(grad, 1 - beta) update = grad.lerp_(momentum, beta) if nesterov else momentum if update.ndim == 4: # for the case of conv filters update = update.view(len(update), -1) - update = zeropower_via_newtonschulz5(update, steps=ns_steps) + if ns_method == "gram": + update = zeropower_via_gram_newtonschulz(update, steps=ns_steps) + else: + update = zeropower_via_newtonschulz5(update, steps=ns_steps) update *= max(1, grad.size(-2) / grad.size(-1))**0.5 return update @@ -93,10 +171,12 @@ class Muon(torch.optim.Optimizer): lr: The learning rate, in units of spectral norm per update. weight_decay: The AdamW-style weight decay. momentum: The momentum. A value of 0.95 here is usually fine. + ns_method: Newton-Schulz method. "gram" (default) uses Gram NS for ~2x speedup + on rectangular matrices. "standard" uses the original iteration. """ - def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95): - defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum) + def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95, ns_method="gram"): + defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_method=ns_method) assert isinstance(params, list) and len(params) >= 1 and isinstance(params[0], torch.nn.Parameter) params = sorted(params, key=lambda x: x.size(), reverse=True) super().__init__(params, defaults) @@ -122,7 +202,10 @@ def step(self, closure=None): state = self.state[p] if len(state) == 0: state["momentum_buffer"] = torch.zeros_like(p) - update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"]) + update = muon_update(p.grad, + state["momentum_buffer"], + beta=group["momentum"], + ns_method=group.get("ns_method", "gram")) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()], @@ -136,8 +219,8 @@ class SingleDeviceMuon(torch.optim.Optimizer): Muon variant for usage in non-distributed settings. """ - def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95): - defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum) + def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95, ns_method="gram"): + defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_method=ns_method) super().__init__(params, defaults) @torch.no_grad() @@ -156,7 +239,10 @@ def step(self, closure=None): state = self.state[p] if len(state) == 0: state["momentum_buffer"] = torch.zeros_like(p) - update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"]) + update = muon_update(p.grad, + state["momentum_buffer"], + beta=group["momentum"], + ns_method=group.get("ns_method", "gram")) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) @@ -208,7 +294,10 @@ def __init__(self, param_groups): group["lr"] = group.get("lr", 0.02) group["momentum"] = group.get("momentum", 0.95) group["weight_decay"] = group.get("weight_decay", 0) - assert set(group.keys()) == set(["params", "lr", "momentum", "weight_decay", "use_muon"]) + group["ns_method"] = group.get("ns_method", "gram") + assert group[ + "ns_method"] in NS_METHODS, f"ns_method must be one of {NS_METHODS}, got {group['ns_method']}" + assert set(group.keys()) == set(["params", "lr", "momentum", "weight_decay", "use_muon", "ns_method"]) else: # defaults group["lr"] = group.get("lr", 3e-4) @@ -240,7 +329,10 @@ def step(self, closure=None): state = self.state[p] if len(state) == 0: state["momentum_buffer"] = torch.zeros_like(p) - update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"]) + update = muon_update(p.grad, + state["momentum_buffer"], + beta=group["momentum"], + ns_method=group.get("ns_method", "gram")) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()], @@ -277,7 +369,10 @@ def __init__(self, param_groups): group["lr"] = group.get("lr", 0.02) group["momentum"] = group.get("momentum", 0.95) group["weight_decay"] = group.get("weight_decay", 0) - assert set(group.keys()) == set(["params", "lr", "momentum", "weight_decay", "use_muon"]) + group["ns_method"] = group.get("ns_method", "gram") + assert group[ + "ns_method"] in NS_METHODS, f"ns_method must be one of {NS_METHODS}, got {group['ns_method']}" + assert set(group.keys()) == set(["params", "lr", "momentum", "weight_decay", "use_muon", "ns_method"]) else: # defaults group["lr"] = group.get("lr", 3e-4) @@ -304,7 +399,10 @@ def step(self, closure=None): state = self.state[p] if len(state) == 0: state["momentum_buffer"] = torch.zeros_like(p) - update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"]) + update = muon_update(p.grad, + state["momentum_buffer"], + beta=group["momentum"], + ns_method=group.get("ns_method", "gram")) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) else: diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index c4f19f43de4f..34925f1d1bbe 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -791,6 +791,7 @@ def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): if self.use_muon: self.sub_groups_using_muon = [] self.muon_beta = None + self.muon_ns_method = None for idx, param_group in enumerate(fp16_param_groups): if getattr(param_group['params'][0], 'use_muon', False): self.sub_groups_using_muon.extend([True] * len(param_groups[idx])) @@ -799,6 +800,7 @@ def _create_fp16_partitions_with_defragmentation(self, fp16_param_groups): raise ValueError(f"All Muon parameter groups must have the same momentum (beta). " f"Found {self.muon_beta} and {group_beta}.") self.muon_beta = group_beta + self.muon_ns_method = param_group.get('ns_method', 'gram') else: self.sub_groups_using_muon.extend([False] * len(param_groups[idx])) # bookkeeping related to param groups @@ -1547,7 +1549,7 @@ def _apply_distributed_muon_update(self, communication_data_type: torch.dtype, b param = params[base_i + rank] g = param.grad m = gathered_momentums_pad[base_i + rank] - update = muon_update(g, m, beta=self.muon_beta) + update = muon_update(g, m, beta=self.muon_beta, ns_method=getattr(self, 'muon_ns_method', 'gram')) g.data.copy_(update, non_blocking=False) grad_handle = dist.all_gather(grads_pad[base_i:base_i + world_sz], grads_pad[base_i + rank], diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index f3a0352bebfa..6f40a0fbccf6 100755 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -1995,7 +1995,11 @@ def get_flat_partition(self, assert tensor.ndim > 1, f"if use muon, then tensor dim > 1, got {tensor.size()}" buffer = torch.narrow(self.optimizer.state[flatten_copy]["momentum_buffer"], 0, buffer_idx, tensor.numel()).view(tensor.size()) - grad_accum = muon_update(grad_accum, buffer, self.optimizer.param_groups[param_group_idx]['momentum']) + ns_method = self.optimizer.param_groups[param_group_idx].get('ns_method', 'gram') + grad_accum = muon_update(grad_accum, + buffer, + self.optimizer.param_groups[param_group_idx]['momentum'], + ns_method=ns_method) tensor = grad_accum num_elements = tensor.numel() buffer_idx += num_elements diff --git a/tests/unit/ops/muon/test_muon.py b/tests/unit/ops/muon/test_muon.py index 02594941cef0..84b06dd96265 100644 --- a/tests/unit/ops/muon/test_muon.py +++ b/tests/unit/ops/muon/test_muon.py @@ -86,3 +86,94 @@ def test(self, optimizer_type, zero_stage, lr, hidden_dim, nlayer, offload_optim after_training = [p.clone().cpu() for p in model.parameters()] for initial, final in zip(initial_params, after_training): assert not torch.equal(initial.cpu(), final.cpu()), "Parameters should have been updated during training" + + +class TestGramNewtonSchulz(DistributedTest): + """Test Gram Newton-Schulz integration with Muon optimizer.""" + + world_size = 2 + reuse_dist_env = True + + @pytest.mark.parametrize('ns_method', ['gram', 'standard']) + @pytest.mark.parametrize('zero_stage', [1, 2]) + def test_ns_method_training(self, ns_method, zero_stage): + """Verify both ns_method values work end-to-end with DeepSpeed.""" + hidden_dim = 64 + batch_size = 8 + config_dict = { + "train_batch_size": batch_size, + "optimizer": { + "type": "muon", + "params": { + "lr": 0.01, + "ns_method": ns_method, + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + }, + "zero_optimization": { + "stage": zero_stage, + "reduce_scatter": False, + }, + } + + model = SimpleModel(hidden_dim=hidden_dim, nlayers=3) + initial_params = [p.clone().cpu() for p in model.parameters()] + engine, optimizer, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=model.parameters(), + dist_init_required=False, + ) + + for _ in range(3): + x = torch.randn(batch_size, hidden_dim, device=engine.device, dtype=torch.half) + y = torch.randint(0, hidden_dim, (batch_size, ), device=engine.device) + loss = engine(x, y) + engine.backward(loss) + engine.step() + + after_training = [p.clone().cpu() for p in model.parameters()] + for initial, final in zip(initial_params, after_training): + assert not torch.equal(initial, final), "Parameters should have been updated" + + @pytest.mark.parametrize('ns_method', ['gram', 'standard']) + def test_ns_method_stage3(self, ns_method): + """Verify ns_method works with ZeRO Stage 3.""" + hidden_dim = 64 + batch_size = 8 + config_dict = { + "train_batch_size": batch_size, + "optimizer": { + "type": "muon", + "params": { + "lr": 0.01, + "ns_method": ns_method, + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + }, + "zero_optimization": { + "stage": 3, + "reduce_scatter": False, + }, + } + + model = SimpleModel(hidden_dim=hidden_dim, nlayers=3) + engine, optimizer, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=model.parameters(), + dist_init_required=False, + ) + + for _ in range(3): + x = torch.randn(batch_size, hidden_dim, device=engine.device, dtype=torch.half) + y = torch.randint(0, hidden_dim, (batch_size, ), device=engine.device) + loss = engine(x, y) + engine.backward(loss) + engine.step() From e9beb2de1e488e78352111cb214edf1ae256b117 Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Thu, 2 Apr 2026 23:55:07 -0700 Subject: [PATCH 2/7] docs: add ns_method parameter to Muon optimizer documentation Signed-off-by: Ma, Guokai --- docs/_pages/config-json.md | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index f8209c8d8068..543339fad43f 100755 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -41,6 +41,17 @@ toc_label: "Contents" Muon optimizer is supported with ZeRO Stage 1, 2, and 3. To use Muon, set the optimizer name to `Muon`. The parameters applied for Muon are automatically determined by the matrix shape and name. For ZeRO Stage 3 with NVMe offloading, set `save_muon_momentum_buffer_in_memory` to `true` under `zero_optimization` to keep the Muon momentum buffer in GPU/CPU memory instead of swapping to NVMe. +Muon supports the following params: + +| "params" key | Description | Default | +| -------------- | -------------------------------------------------------------------------------------------------------------------- | --------- | +| lr | Learning rate for all parameters. Overridden by `muon_lr` / `adam_lr` if set. | 0.001 | +| momentum | Momentum coefficient for the Muon update. | 0.95 | +| weight\_decay | Weight decay (AdamW-style). | 0.0 | +| muon\_lr | Learning rate override for Muon parameters. Defaults to `lr` if not set. | - | +| adam\_lr | Learning rate override for non-Muon (Adam) parameters. Defaults to `lr` if not set. | - | +| ns\_method | Newton-Schulz orthogonalization method: `"gram"` for Gram NS (~2x faster on rectangular matrices), `"standard"` for the original iteration. Use `"standard"` to fall back if you encounter convergence issues. | `"gram"` | + Example of **optimizer** with Adam ```json @@ -73,7 +84,8 @@ If not set, muon_lr will default to lr. "lr": 0.001, "momentum": 0.9, "weight_decay": 0.0, - "muon_lr": 0.001 + "muon_lr": 0.001, + "ns_method": "gram" } }, "zero_optimization": { From d17212ef8b8cd37b4b2fb40b5039be65e77d2ddf Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Thu, 2 Apr 2026 23:55:13 -0700 Subject: [PATCH 3/7] fix: correct Gram Newton-Schulz reference URL Signed-off-by: Ma, Guokai --- deepspeed/runtime/zero/muon/original_muon.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 83c118e82a16..a0850bccd3fd 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -80,7 +80,7 @@ def zeropower_via_gram_newtonschulz(G, steps: int): Falls back to standard Newton-Schulz for square matrices (n == m) where there is no FLOP advantage. - Reference: https://arxiv.org/abs/2503.02022 + Reference: https://tridao.me/blog/2026/gram-newton-schulz/ """ assert G.ndim >= 2 a, b, c = (3.4445, -4.7750, 2.0315) From 54930203d2fa00c9f7f189dd432218f8b51b74b2 Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Fri, 3 Apr 2026 00:21:07 -0700 Subject: [PATCH 4/7] Use accelerator API for dtype selection in Newton-Schulz iterations Both NS functions now query the accelerator to choose compute dtype instead of hardcoding. Standard NS uses is_bf16_supported() to select bf16 vs fp32; Gram NS uses is_fp16_supported() to select fp16 vs fp32. Signed-off-by: Ma, Guokai --- deepspeed/runtime/zero/muon/original_muon.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index a0850bccd3fd..a48f14b0bcca 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -46,7 +46,9 @@ def zeropower_via_newtonschulz5(G, steps: int): """ assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() + # Use bf16 when hardware supports it; fp32 otherwise + compute_dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32 + X = G.to(compute_dtype) if G.size(-2) > G.size(-1): X = X.mT @@ -84,8 +86,8 @@ def zeropower_via_gram_newtonschulz(G, steps: int): """ assert G.ndim >= 2 a, b, c = (3.4445, -4.7750, 2.0315) - # Use fp16 on GPU for better precision than bf16; fp32 on CPU for stability - compute_dtype = torch.float16 if get_accelerator().on_accelerator(G) else torch.float32 + # Use fp16 for better precision than bf16 when hardware supports it; fp32 otherwise + compute_dtype = torch.float16 if get_accelerator().is_fp16_supported() else torch.float32 X = G.to(compute_dtype) if G.size(-2) > G.size(-1): X = X.mT From e5de42ce5cab4c2700cd640e5d54c7c9b807bfaf Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Fri, 3 Apr 2026 06:08:26 -0700 Subject: [PATCH 5/7] Fix non-contiguous tensor output from Gram NS for tall matrices Gram Newton-Schulz produces non-contiguous tensors via .mT for tall weight matrices (e.g., gate_proj/up_proj in LLaMA). This caused downstream grad norm computation (g.data.double()) to be ~1.8x slower due to strided memory access, adding ~75ms to optimizer step time. Add .contiguous() to the Gram NS return path for tall matrices, and ensure muon_update casts back to the original gradient dtype (Gram NS uses fp16 internally while gradients are bf16). Benchmark (Qwen2.5-3B, 2xA100, ZeRO-2, 3 runs avg): Before fix: 945.1ms/step (optimizer: 229.9ms) After fix: 936.6ms/step (optimizer: 204.3ms) Standard NS baseline: 1054.5ms/step Gram NS speedup: 10.4% -> 11.2% Signed-off-by: Ma, Guokai --- deepspeed/runtime/zero/muon/original_muon.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index a48f14b0bcca..3ef981c23155 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -132,7 +132,7 @@ def zeropower_via_gram_newtonschulz(G, steps: int): X = Q @ X if G.size(-2) > G.size(-1): - X = X.mT + X = X.mT.contiguous() return X @@ -141,6 +141,7 @@ def zeropower_via_gram_newtonschulz(G, steps: int): @compiler.compile() def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram"): + orig_dtype = grad.dtype momentum.lerp_(grad, 1 - beta) update = grad.lerp_(momentum, beta) if nesterov else momentum if update.ndim == 4: # for the case of conv filters @@ -150,6 +151,8 @@ def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method= else: update = zeropower_via_newtonschulz5(update, steps=ns_steps) update *= max(1, grad.size(-2) / grad.size(-1))**0.5 + if update.dtype != orig_dtype: + update = update.to(orig_dtype) return update From a6cf6b69c7f90cf2e9957143c42f8eb1367efa7b Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Fri, 3 Apr 2026 06:55:33 -0700 Subject: [PATCH 6/7] Fold transpose into matmul in Gram NS for tall matrices Replace (Q @ X).mT.contiguous() with X.mT @ Q.mT which produces a contiguous result directly. cuBLAS handles transposed inputs natively via transpose flags, so the matmul cost is identical but the extra memcpy from .contiguous() is eliminated. Benchmark (Qwen2.5-3B, 2xA100, ZeRO-2, 3 runs avg): Before: 936.6ms/step (backward: 628.4ms) After: 931.5ms/step (backward: 612.8ms) Speedup vs standard NS: 11.2% -> 11.7% Signed-off-by: Ma, Guokai --- deepspeed/runtime/zero/muon/original_muon.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 3ef981c23155..ec97cdf85d8d 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -129,10 +129,10 @@ def zeropower_via_gram_newtonschulz(G, steps: int): RZ = a * R + Z @ R R = a * RZ + Z @ RZ - X = Q @ X - if G.size(-2) > G.size(-1): - X = X.mT.contiguous() + X = X.mT @ Q.mT + else: + X = Q @ X return X From 61095611639d8d563c8c7a4b20b378f7724b388c Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Fri, 3 Apr 2026 07:33:53 -0700 Subject: [PATCH 7/7] Use fused addmm and eliminate eye allocation in Gram NS Replace separate scalar-multiply + matmul + add operations with single torch.addmm calls for Q and R updates, reducing kernel launch overhead. Remove torch.eye allocation by using diagonal().add_() instead. Signed-off-by: Ma, Guokai --- deepspeed/runtime/zero/muon/original_muon.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index ec97cdf85d8d..2bb745da4f5f 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -108,7 +108,6 @@ def zeropower_via_gram_newtonschulz(G, steps: int): # Gram NS: iterate on R = X @ X.T (n x n) instead of X (n x m) R = X @ X.mT - I = torch.eye(n, device=X.device, dtype=X.dtype) Q = None restart_at = 2 @@ -121,13 +120,14 @@ def zeropower_via_gram_newtonschulz(G, steps: int): Z = b * R + c * R @ R if Q is None: - Q = Z + a * I + Q = Z.clone() + Q.diagonal().add_(a) else: - Q = a * Q + Z @ Q + Q = torch.addmm(Q, Z, Q, beta=a, alpha=1.0) if i < steps - 1 and (i + 1) != restart_at: - RZ = a * R + Z @ R - R = a * RZ + Z @ RZ + RZ = torch.addmm(R, Z, R, beta=a, alpha=1.0) + R = torch.addmm(RZ, Z, RZ, beta=a, alpha=1.0) if G.size(-2) > G.size(-1): X = X.mT @ Q.mT