From 41e501dadf5509a76c670d059e7ff7c7809634ad Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Tue, 30 Jun 2026 17:28:24 -0700 Subject: [PATCH 1/6] Balance LayerWise LPT by Newton-Schulz compute cost, not shard numel The existing LPT bin-packing sorts params by numel and assigns each to the least-loaded shard by numel. With GTP, this creates severe compute imbalance: a GTP-sharded param with local shape [548, 8192] has the same numel as a non-GTP param of that size, but its Newton-Schulz cost is 64x higher because NS operates on the full all-gathered [35072, 8192] shape. Three changes: 1. _ns_compute_cost() estimates NS cost using the full post-AllGather shape: max(M,N) * min(M,N)^2. 2. Sort and assign by compute cost instead of numel. 3. Persistent shard_compute_loads across buckets (shard_cursors still reset per bucket for correct memory layout) with a per-bucket numel cap (epsilon=0.3) to bound padding. Reduces compute imbalance from 137x to 1.3x in simulation. Measured 19% throughput improvement on Ultra-half (54-layer hybrid Mamba-MoE, 1024 GB200 GPUs, Muon + GTP=64): 2080 -> 1707 ms/iter. Co-Authored-By: Claude Opus 4.6 --- .../core/optimizer/layer_wise_optimizer.py | 44 ++++++++++++++++--- 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 376f9a1f1c0..9d3dfd90289 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -247,6 +247,25 @@ def _compute_per_buffer_param_layout( bucket_id = 0 shard_imbalance_padding_numel = 0 + # Persistent compute loads across buckets so LPT spreads expensive + # (GTP-sharded) params evenly instead of clustering them per bucket. + shard_compute_loads = [0] * dp_size + + def _ns_compute_cost(param): + """Estimate Newton-Schulz compute cost for a 2D parameter. + + For GTP-sharded params, reconstructs the full post-AllGather shape + (GTP always shards along dim 0). Cost ~ max(M,N) * min(M,N)^2, + which is the dominant term in the NS orthogonalization. + """ + if param.dim() != 2: + return param.data.nelement() + m, n = param.data.shape + if getattr(param, 'is_gtp_weight_remat', False): + m = m * getattr(param, 'gtp_remat_size', 1) + big, small = max(m, n), min(m, n) + return big * small * small + def _emit_bucket( chunk_params: List[torch.nn.Parameter], shared_embedding: bool = False ) -> None: @@ -276,17 +295,28 @@ def _emit_bucket( shard_assignments[shard_id].append((None, numel)) shard_cursors[shard_id] = numel else: - # Greedy LPT: largest first, assign to the least-loaded shard. - # The within-shard order is sorted-by-numel, not backprop — - # that is fine because all params in the chunk share the same - # bucket_id, so DDP's backprop-order iteration still sees - # monotonic bucket_ids across the chunk boundary. - for param in sorted(chunk_params, key=lambda p: -p.data.nelement()): + # Compute-balanced LPT: sort by Newton-Schulz compute cost + # (accounts for full post-AllGather shape under GTP), assign to + # the shard with least accumulated compute load. Compute loads + # persist across buckets; numel cursors reset per bucket. + # A per-bucket numel cap prevents excessive padding. + _NUMEL_EPSILON = 0.3 + total_chunk_numel = sum(p.data.nelement() for p in chunk_params) + max_shard_numel = total_chunk_numel / dp_size * (1 + _NUMEL_EPSILON) + for param in sorted(chunk_params, key=lambda p: -_ns_compute_cost(p)): numel = param.data.nelement() - min_shard = min(range(dp_size), key=lambda s: shard_cursors[s]) + candidates = [ + s for s in range(dp_size) + if pad_param_start(shard_cursors[s]) + numel <= max_shard_numel + ] + if candidates: + min_shard = min(candidates, key=lambda s: shard_compute_loads[s]) + else: + min_shard = min(range(dp_size), key=lambda s: shard_cursors[s]) placement = pad_param_start(shard_cursors[min_shard]) shard_assignments[min_shard].append((param, numel)) shard_cursors[min_shard] = placement + numel + shard_compute_loads[min_shard] += _ns_compute_cost(param) padded_shard_size = pad_to_divisor(max(shard_cursors), shard_divisor) bucket_start_index = buffer_cursor From 26030519a72c21f3b74cca852e63152cef856337 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Tue, 30 Jun 2026 12:13:12 -0700 Subject: [PATCH 2/6] Suppress noisy external-library log lines on non-rank-0 processes Add DeprecationWarning to the warning filter list on non-rank-0. Override warnings.showwarning as a fallback for libraries like CUTLASS DSL that use catch_warnings() with simplefilter("always") to bypass filters. Suppress noisy Python loggers (GroupedGemm variants, absl) on non-rank-0 in setup_logging(). Signed-off-by: Deepak Narayanan --- megatron/training/initialize.py | 10 ++++++++++ pretrain_gpt.py | 13 +++++++++++++ pretrain_hybrid.py | 13 +++++++++++++ 3 files changed, 36 insertions(+) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 87d6aa65b03..c3eb2cccb2e 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -580,3 +580,13 @@ def setup_logging() -> None: if is_rank0(): logger.info(f'Setting logging level to {logging_level}') logging.getLogger().setLevel(logging_level) + + if not is_rank0(): + for noisy_logger_name in [ + 'GroupedGemmQuantSm100', + 'GroupedGemmDsreluSm100', + 'GroupedGemmSreluSm100', + 'GroupedGemmWgradSm100', + 'absl', + ]: + logging.getLogger(noisy_logger_name).setLevel(logging.ERROR) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 47c1935eb90..8b979d16d4a 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -17,6 +17,19 @@ if rank != 0: warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", category=DeprecationWarning) + + # Some libraries (e.g., CUTLASS DSL) use warnings.catch_warnings() with + # simplefilter("always"), which overrides the filters above. Override + # showwarning as a fallback to suppress warnings that slip through. + _original_showwarning = warnings.showwarning + + def _rank0_only_showwarning(message, category, filename, lineno, file=None, line=None): + if issubclass(category, (UserWarning, FutureWarning, DeprecationWarning)): + return + _original_showwarning(message, category, filename, lineno, file, line) + + warnings.showwarning = _rank0_only_showwarning from functools import lru_cache, partial from typing import Any, List, Optional, Tuple diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 39bc7f30b57..cf665b36887 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -16,6 +16,19 @@ if rank != 0: warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", category=DeprecationWarning) + + # Some libraries (e.g., CUTLASS DSL) use warnings.catch_warnings() with + # simplefilter("always"), which overrides the filters above. Override + # showwarning as a fallback to suppress warnings that slip through. + _original_showwarning = warnings.showwarning + + def _rank0_only_showwarning(message, category, filename, lineno, file=None, line=None): + if issubclass(category, (UserWarning, FutureWarning, DeprecationWarning)): + return + _original_showwarning(message, category, filename, lineno, file, line) + + warnings.showwarning = _rank0_only_showwarning from functools import lru_cache, partial from typing import Any, List, Optional, Tuple From 39377174e0eb70ac884b4a7d58a5d9dfb68fe0b9 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Thu, 2 Jul 2026 08:10:09 -0700 Subject: [PATCH 3/6] Fix TE version check --- .../core/tensor_parallel/generalized_tensor_parallelism.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index c72cf313dc2..1de1a4fcca0 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -45,13 +45,6 @@ try: import transformer_engine as te # noqa: F401 - _te_version = Version(te.__version__) - if _te_version < _GTP_TE_MIN_VERSION: - raise ImportError( - f"megatron.core.tensor_parallel.gtp_api requires TransformerEngine " - f">= {_GTP_TE_MIN_VERSION} (found {_te_version})." - ) - import transformer_engine_torch as tex from transformer_engine.pytorch.constants import ( MXFP8_BLOCK_SCALING_SIZE, From 43c4c60449942af28dca21ae0c6f4f2a353b43ca Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Sun, 21 Jun 2026 16:03:59 -0700 Subject: [PATCH 4/6] Fill LayerWise optimizer bucket padding with real parameters instead of sharding if possible The layer-wise (Muon) optimizer distributes whole parameter matrices across dp_size shards within each bucket, and each bucket's buffer is sized to dp_size times its largest shard. Cutting a bucket the instant chunk_numel >= bucket_size can leave the shards unevenly loaded, so the shorter shards become padding (e.g., 100 equal expert matrices spread over 32 shards use buffer space for 128, +28% per bucket). Treat bucket_size as a soft minimum instead: once it is met, keep absorbing params while each one fits into the existing shard padding (without growing dp_size * max_shard_cursor), and only close the bucket when the next param would enlarge it. The padding slots get filled with real params, so when params pack evenly across dp_size shards the per-bucket overhead drops to zero. The single-dominating-param case is unchanged in spirit: PADDING_FLOOR still governs the soft minimum. Co-Authored-By: Claude Opus 4.8 --- .../core/optimizer/layer_wise_optimizer.py | 47 ++++++++++++++----- 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 9d3dfd90289..79d50e4fe6c 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -342,17 +342,34 @@ def _emit_bucket( # # Padding floor: the on-buffer bucket size is ``dp_size * # max_shard_cursor``, which is at least ``dp_size * chunk_max_param`` - # because some shard must hold that param whole. If a single param - # dominates the chunk, finalising on ``chunk_numel >= bucket_size`` - # alone would emit a bucket with most of its shards near-empty - # padding. Instead extend the chunk so its raw numel approaches the - # padded buffer size, capping per-bucket overhead at ``1 / - # PADDING_FLOOR - 1`` (~11% at 0.9). Falls back to ``bucket_size`` - # when no single param dominates. + # because some shard must hold that param whole. ``bucket_size`` is a + # soft *minimum*: once it is reached we keep absorbing params as long + # as each one fits into the existing shard padding (i.e., without + # growing ``dp_size * max_shard_cursor``), and only close the bucket + # when the next param would otherwise enlarge it. This fills shard + # padding with real params instead of emitting padded-out buckets. + # When the params pack evenly across ``dp_size`` shards the overhead + # is zero. ``int(dp_size * chunk_max_param * PADDING_FLOOR)`` keeps the + # soft minimum sensible when a single param dominates a shard. PADDING_FLOOR = 0.9 chunk_params: List[torch.nn.Parameter] = [] chunk_numel = 0 chunk_max_param = 0 + # Mirror _emit_bucket's greedy LPT placement incrementally so we can + # decide, per param, whether it still fits in the current bucket. + shard_loads = [0] * dp_size + + def _absorbs(numel: int) -> bool: + """True if ``numel`` fits in the least-loaded shard without growing + the bucket's padded size (``dp_size * padded_shard_size``), i.e., + it fills existing shard padding instead of adding a new row.""" + target = pad_to_divisor(max(shard_loads), shard_divisor) + return pad_param_start(min(shard_loads)) + numel <= target + + def _place(numel: int) -> None: + shard_id = min(range(dp_size), key=lambda s: shard_loads[s]) + shard_loads[shard_id] = pad_param_start(shard_loads[shard_id]) + numel + for param in reversed(params): param_numel = param.data.nelement() if getattr(param, 'shared_embedding', False): @@ -362,18 +379,24 @@ def _emit_bucket( chunk_params = [] chunk_numel = 0 chunk_max_param = 0 + shard_loads[:] = [0] * dp_size _emit_bucket([param], shared_embedding=True) continue - chunk_params.append(param) - chunk_numel += param_numel - chunk_max_param = max(chunk_max_param, param_numel) - if bucket_size is not None: + # Close the bucket once it has met its soft-minimum size *and* this + # param can no longer be absorbed into the existing shard padding + # (adding it would grow the bucket). + if bucket_size is not None and chunk_params: threshold = max(bucket_size, int(dp_size * chunk_max_param * PADDING_FLOOR)) - if chunk_numel >= threshold: + if chunk_numel >= threshold and not _absorbs(param_numel): _emit_bucket(chunk_params) chunk_params = [] chunk_numel = 0 chunk_max_param = 0 + shard_loads[:] = [0] * dp_size + _place(param_numel) + chunk_params.append(param) + chunk_numel += param_numel + chunk_max_param = max(chunk_max_param, param_numel) _emit_bucket(chunk_params) total_buffer_numel = bucket_indices[-1][1] if bucket_indices else 0 From 3baee21ed6b9df8ad6646a65efed4477aab19f5e Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Tue, 23 Jun 2026 18:30:18 -0700 Subject: [PATCH 5/6] GTP+Muon: support 'blockwise' and 'distributed' mode Branch the GTP orthogonalize path on self.tp_mode, mirroring how newton_schulz_tp treats TP-sharded params: - blockwise: local NS on the [M/g, K] shard, no GTP collective - duplicated: all-gather over GTP -> whole-matrix NS -> reshard - distributed: distribute NS over GTP via the small-Gram all-reduce Signed-off-by: Shiqing Fan Signed-off-by: Deepak Narayanan --- .../core/optimizer/emerging_optimizers.py | 72 ++++- .../test_gtp_muon.py | 252 ++++++++++++++++++ .../test_gtp_muon_dcp.py | 7 + 3 files changed, 318 insertions(+), 13 deletions(-) create mode 100644 tests/unit_tests/generalized_tensor_parallel/test_gtp_muon.py diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 53ac956b35c..cf5985f6ff0 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -230,13 +230,24 @@ def scaled_orthogonalize_fn( scaled_orthogonalize_fn=scaled_orthogonalize_fn, ) + @staticmethod + def _all_gather_tensor(t, group, dim): + """All-gather equal-size shards of ``t`` over ``group`` and concat along ``dim``.""" + shards = [torch.empty_like(t) for _ in range(get_pg_size(group))] + torch.distributed.all_gather(shards, t.contiguous(), group) + return torch.cat(shards, dim=dim) + def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_dim): - """All-gather grad along GTP_remat/EGTP_remat dim 0, orthogonalize, then slice back. + """Orthogonalize a (possibly GTP-sharded) momentum, then reshard. + + When GTP is inactive this is a plain passthrough to ``scaled_orthogonalize_fn``. + Otherwise, ``self.tp_mode`` controls how GTP sharding is handled: - GTP_remat shards weights along dim 0 independently of TP's partition_dim. Newton-Schulz - needs the full weight matrix, so we reconstruct the GTP_remat dimension before running - the TP-aware orthogonalization, then extract the local GTP_remat shard from the result. - When GTP_remat is inactive this is a plain passthrough to scaled_orthogonalize_fn. + - **blockwise**: orthogonalize the local GTP shard independently, no collective. + - **duplicated**: all-gather over GTP, run whole-matrix NS (TP-aware), reshard. + - **distributed**: distribute NS over GTP via small-Gram all-reduce. When both + GTP and TP are active, NS is distributed over the larger group to minimize + redundant compute; the smaller group is all-gathered beforehand. """ # TODO: Clean up code that determines if parameter is a MoE layer and which TP group to use is_expert = getattr(p, 'expert_tp', False) @@ -256,21 +267,56 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di return self.scaled_orthogonalize_fn(grad, tp_group, partition_dim) gtp_remat_size = get_pg_size(gtp_remat_group) - gtp_rank = get_pg_rank(gtp_remat_group) - shards = [torch.empty_like(grad) for _ in range(gtp_remat_size)] - torch.distributed.all_gather(shards, grad, gtp_remat_group) - gathered_grad = torch.cat(shards, dim=0) - gathered_grad = self.scaled_orthogonalize_fn(gathered_grad, tp_group, partition_dim) + if self.tp_mode == "blockwise": + # Local block NS on this rank's GTP row-shard (shape [M/gtp_remat_size, K]): + # partition_dim=None makes scaled_orthogonalize_fn run a plain Newton-Schulz on + # the shard with no GTP/TP collective. + return self.scaled_orthogonalize_fn(grad, tp_group, None) + + if self.tp_mode == "duplicated": + # All-gather the full matrix over GTP (dim 0), orthogonalize the whole tensor + # (scaled_orthogonalize_fn handles any TP sharding per tp_mode), reshard dim 0. + gathered_grad = self._all_gather_tensor(grad, gtp_remat_group, 0) + gathered_grad = self.scaled_orthogonalize_fn(gathered_grad, tp_group, partition_dim) + shard_size = gathered_grad.size(0) // gtp_remat_size + gtp_rank = get_pg_rank(gtp_remat_group) + return gathered_grad[gtp_rank * shard_size : (gtp_rank + 1) * shard_size].contiguous() + + # distributed: NS via the small-Gram all-reduce (no redundant full-matrix NS). + is_tp_active = ( + partition_dim is not None and tp_group is not None and get_pg_size(tp_group) > 1 + ) + + if not is_tp_active: + # GTP-only: distribute NS over the GTP group on the local dim-0 row shard. + return self.scaled_orthogonalize_fn(grad, gtp_remat_group, partition_dim=0) + + # GTP + TP: distributed NS can only operate over one (group, dim) at a + # time. Distribute over the larger group so that the NS GEMMs are sharded + # across more ranks (less redundant compute), and all-gather the smaller + # group to eliminate its sharding beforehand. + tp_size = get_pg_size(tp_group) + if gtp_remat_size >= tp_size: + smaller_group, smaller_dim = tp_group, partition_dim + larger_group, larger_dim = gtp_remat_group, 0 + else: + smaller_group, smaller_dim = gtp_remat_group, 0 + larger_group, larger_dim = tp_group, partition_dim - shard_size = gathered_grad.shape[0] // gtp_remat_size - return gathered_grad[gtp_rank * shard_size : (gtp_rank + 1) * shard_size].contiguous() + gathered_grad = self._all_gather_tensor(grad, smaller_group, smaller_dim) + orthogonalized_grad = self.scaled_orthogonalize_fn(gathered_grad, larger_group, larger_dim) + shard_size = orthogonalized_grad.size(smaller_dim) // get_pg_size(smaller_group) + reshard_rank = get_pg_rank(smaller_group) + return orthogonalized_grad.narrow( + smaller_dim, reshard_rank * shard_size, shard_size + ).contiguous() def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: """Orthogonalize the momentum. Args: - p: The parameter tensor. i is necessary to pass param tensor in addition to + p: The parameter tensor. It is necessary to pass param tensor in addition to momentum because a lot of information is only available in the param tensor, attributes for example. grad: The momentum tensor. diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon.py new file mode 100644 index 00000000000..1b2780680a4 --- /dev/null +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon.py @@ -0,0 +1,252 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Parity tests for GTP + Muon Newton-Schulz orthogonalization. + +``TensorParallelMuon.scaled_orthogonalize_fn_with_gtp_remat`` orthogonalizes a GTP-row-sharded +momentum in one of three modes (blockwise, duplicated, distributed) instead of all-gathering +the full matrix and running a redundant full NS on every rank. These tests verify that each +mode produces the correct per-shard result: + + 1. test_gtp_distributed_mode - GTP4, TP1: distribute NS over GTP (dim 0). + 2. test_row_parallel - TP2, GTP2 (TP dim 1, GTP dim 0): gather smaller, distribute larger. + 3. test_col_parallel - TP2, GTP2 (both dim 0): gather smaller, distribute larger. + 4. test_gtp_duplicated_mode - GTP4, TP1: all-gather over GTP, full NS, reshard. + 5. test_gtp_blockwise_mode - GTP4, TP1: local NS on GTP shard, no collective. + +All require world_size == 4. +""" + +import pytest +import torch + +from megatron.core.tensor_parallel.gtp import HAVE_GTP + +if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.17", allow_module_level=True) + +from megatron.core import parallel_state as ps +from megatron.core.optimizer.emerging_optimizers import HAVE_EMERGING_OPTIMIZERS, TensorParallelMuon + +if not HAVE_EMERGING_OPTIMIZERS: + pytest.skip("emerging_optimizers not available", allow_module_level=True) + +from megatron.core.process_groups_config import ProcessGroupCollection +from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import ( + _torchrun_dist_init, # autouse fixture: initializes the torchrun dist group +) +from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import ( + reset_gtp_globals, # autouse fixture: resets GTP class state between tests +) +from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import ( # noqa: F401 + _requires_multi_gpu, + _run_distributed, +) + +# Parity is asserted at num_ns_steps=1: there the distributed Gram all-reduce equals the +# full-matrix Gram to fp32 reduction-order noise (~1e-5). At more steps the (mathematically +# identical) distributed result still matches full NS in exact arithmetic, but the aggressive +# NS coefficients are tuned beyond convergence (see newton_schulz docstring) and amplify the +# ~1e-6 fp difference (~1e-2 by step 5) — that is NS conditioning, not a distribution error, so +# a one-step parity check is the meaningful correctness test. fp32-highest matmul throughout. +_M, _K = 128, 64 +_NS_STEPS = 1 +_ATOL, _RTOL = 1e-4, 1e-4 + + +def _make_muon(pg_collection, tp_mode="distributed"): + """A TensorParallelMuon used only for its orthogonalize helpers (never stepped).""" + placeholder = torch.nn.Parameter(torch.zeros(1, device="cuda")) + return TensorParallelMuon( + params=[placeholder], + lr=0.01, + momentum=0.95, + weight_decay=0.0, + num_ns_steps=_NS_STEPS, + fp32_matmul_prec="highest", + pg_collection=pg_collection, + tp_mode=tp_mode, + ) + + +def _full_weight(): + """Full [M, K] momentum, identical on every rank (rank-0 broadcast).""" + torch.manual_seed(0) + w = torch.randn(_M, _K, dtype=torch.float32, device="cuda") + torch.distributed.broadcast(w, src=0) + return w + + +def _reference_full_orth(opt, w, tp_group): + """Orthogonalize the full matrix (partition_dim=None → plain NS), same scale/coeffs.""" + return opt.scaled_orthogonalize_fn(w.clone(), tp_group, partition_dim=None) + + +def _world_size(group): + return torch.distributed.get_world_size(group=group) + + +def _rank(group): + return torch.distributed.get_rank(group=group) + + +def _worker_gtp_distributed(rank, world_size, port): + """distributed mode (GTP4, TP1): distribute NS over the GTP group on the local dim-0 shard.""" + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=world_size + ) + try: + pgc = ProcessGroupCollection.use_mpu_process_groups() + opt = _make_muon(pgc) + w = _full_weight() + ref = _reference_full_orth(opt, w, pgc.tp) + + gs, gr = _world_size(pgc.gtp_remat), _rank(pgc.gtp_remat) + sp = _M // gs + local = w[gr * sp : (gr + 1) * sp, :].clone() + local.is_gtp_weight_remat = True + + out = opt.scaled_orthogonalize_fn_with_gtp_remat(local, local, pgc.tp, None) + expected = ref[gr * sp : (gr + 1) * sp, :] + torch.testing.assert_close(out, expected, atol=_ATOL, rtol=_RTOL) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + +def _worker_row_parallel(rank, world_size, port): + """RowParallel TP, GTP (TP dim 1, GTP dim 0): gather the smaller group, distribute larger.""" + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=2, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + pgc = ProcessGroupCollection.use_mpu_process_groups() + opt = _make_muon(pgc) + w = _full_weight() + ref = _reference_full_orth(opt, w, pgc.tp) + + gs, gr = _world_size(pgc.gtp_remat), _rank(pgc.gtp_remat) + ts, tr = _world_size(pgc.tp), _rank(pgc.tp) + sp, kt = _M // gs, _K // ts # GTP row block, TP col block + local = w[gr * sp : (gr + 1) * sp, tr * kt : (tr + 1) * kt].clone() + local.is_gtp_weight_remat = True + + out = opt.scaled_orthogonalize_fn_with_gtp_remat(local, local, pgc.tp, 1) + expected = ref[gr * sp : (gr + 1) * sp, tr * kt : (tr + 1) * kt] + torch.testing.assert_close(out, expected, atol=_ATOL, rtol=_RTOL) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + +def _worker_col_parallel(rank, world_size, port): + """ColumnParallel TP, GTP (both dim 0): gather smaller group, distribute larger. + + dim-0 carve is TP-outer / GTP-inner (GTP slices the already-TP-sharded weight), so this + rank owns rows ``tr*(M/ts) + gr*sp : + sp``. + """ + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=2, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + pgc = ProcessGroupCollection.use_mpu_process_groups() + opt = _make_muon(pgc) + w = _full_weight() + ref = _reference_full_orth(opt, w, pgc.tp) + + gs, gr = _world_size(pgc.gtp_remat), _rank(pgc.gtp_remat) + ts, tr = _world_size(pgc.tp), _rank(pgc.tp) + m_tp = _M // ts + sp = m_tp // gs + off = tr * m_tp + gr * sp + local = w[off : off + sp, :].clone() + local.is_gtp_weight_remat = True + + out = opt.scaled_orthogonalize_fn_with_gtp_remat(local, local, pgc.tp, 0) + expected = ref[off : off + sp, :] + torch.testing.assert_close(out, expected, atol=_ATOL, rtol=_RTOL) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + +def _worker_gtp_duplicated(rank, world_size, port): + """duplicated mode (GTP4, TP1): all-gather full matrix over GTP, whole NS, reshard. + + Mathematically identical to the full-matrix reference, so the local block must match. + """ + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=world_size + ) + try: + pgc = ProcessGroupCollection.use_mpu_process_groups() + opt = _make_muon(pgc, tp_mode="duplicated") + w = _full_weight() + ref = _reference_full_orth(opt, w, pgc.tp) + + gs, gr = _world_size(pgc.gtp_remat), _rank(pgc.gtp_remat) + sp = _M // gs + local = w[gr * sp : (gr + 1) * sp, :].clone() + local.is_gtp_weight_remat = True + + out = opt.scaled_orthogonalize_fn_with_gtp_remat(local, local, pgc.tp, None) + expected = ref[gr * sp : (gr + 1) * sp, :] + torch.testing.assert_close(out, expected, atol=_ATOL, rtol=_RTOL) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + +def _worker_gtp_blockwise(rank, world_size, port): + """blockwise mode (GTP4, TP1): local NS on the [M/gtp_size, K] shard, no GTP collective. + + Must equal a plain Newton-Schulz of the local shard, not the full-matrix shard. + """ + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=world_size + ) + try: + pgc = ProcessGroupCollection.use_mpu_process_groups() + opt = _make_muon(pgc, tp_mode="blockwise") + w = _full_weight() + + gs, gr = _world_size(pgc.gtp_remat), _rank(pgc.gtp_remat) + sp = _M // gs + local = w[gr * sp : (gr + 1) * sp, :].clone() + local.is_gtp_weight_remat = True + + out = opt.scaled_orthogonalize_fn_with_gtp_remat(local, local, pgc.tp, None) + # blockwise orthogonalizes the local block independently — no GTP comm. + expected = opt.scaled_orthogonalize_fn(local.clone(), pgc.tp, None) + torch.testing.assert_close(out, expected, atol=_ATOL, rtol=_RTOL) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + +class TestGTPMuonDistributedNS: + """Distributed-NS orthogonalization matches full-matrix NS, per shard.""" + + def test_gtp_distributed_mode(self): + _requires_multi_gpu(4) + _run_distributed(_worker_gtp_distributed, 4) + + def test_row_parallel(self): + _requires_multi_gpu(4) + _run_distributed(_worker_row_parallel, 4) + + def test_col_parallel(self): + _requires_multi_gpu(4) + _run_distributed(_worker_col_parallel, 4) + + def test_gtp_duplicated_mode(self): + _requires_multi_gpu(4) + _run_distributed(_worker_gtp_duplicated, 4) + + def test_gtp_blockwise_mode(self): + _requires_multi_gpu(4) + _run_distributed(_worker_gtp_blockwise, 4) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py index b26d8a974ce..b1ead3faf86 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py @@ -8,6 +8,13 @@ by ``replica_id`` so DCP does not see multiple writers for the same shard. """ +import pytest + +from megatron.core.tensor_parallel.gtp import HAVE_GTP + +if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.17", allow_module_level=True) + import torch from megatron.core.dist_checkpointing import load, save From 38cbbe0e25fa46071042be3404496574781c1ef0 Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Wed, 15 Jul 2026 16:06:24 -0700 Subject: [PATCH 6/6] Add --muon-use-syrk option for Triton SYRK kernel in Newton-Schulz Co-Authored-By: Claude Opus 4.6 Signed-off-by: Deepak Narayanan --- megatron/core/optimizer/emerging_optimizers.py | 16 ++++++++++++++++ megatron/core/optimizer/optimizer_config.py | 3 +++ megatron/training/arguments.py | 3 +++ 3 files changed, 22 insertions(+) diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index cf5985f6ff0..9345bd66088 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -30,6 +30,8 @@ ) from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT, newton_schulz_tp + _NS_TP_SUPPORTS_SYRK = "use_syrk" in inspect.signature(newton_schulz_tp).parameters + # It is necessary to import optimizers for the registry to work. from emerging_optimizers.scalar_optimizers import Lion # pylint: disable=unused-import from emerging_optimizers.soap import SOAP # pylint: disable=unused-import @@ -178,9 +180,17 @@ def __init__( extra_scale_factor: float = 1.0, pg_collection: Optional[ProcessGroupCollection] = None, tp_mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated", + use_syrk: bool = False, ) -> None: if num_ns_steps < 1: raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") + if use_syrk and not _NS_TP_SUPPORTS_SYRK: + log_single_rank( + logger, + logging.WARNING, + "use_syrk requested but the installed emerging_optimizers does not support it; " + "falling back to standard GEMM.", + ) def scaled_orthogonalize_fn( grad: torch.Tensor, @@ -197,6 +207,9 @@ def scaled_orthogonalize_fn( size = [grad.size(-2), grad.size(-1)] if partition_dim is not None: size[partition_dim] *= get_pg_size(tp_group) + ns_kwargs = {} + if use_syrk and _NS_TP_SUPPORTS_SYRK: + ns_kwargs["use_syrk"] = True orth_grad = newton_schulz_tp( grad, steps=num_ns_steps, @@ -204,6 +217,7 @@ def scaled_orthogonalize_fn( tp_group=tp_group, partition_dim=partition_dim, tp_mode="duplicated" if tp_mode == "blockwise" else tp_mode, + **ns_kwargs, ) scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode) return orth_grad * scale_factor * extra_scale_factor @@ -422,6 +436,7 @@ def __init__( extra_scale_factor: float = 1.0, pg_collection: Optional[ProcessGroupCollection] = None, tp_mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated", + use_syrk: bool = False, moment2_method: Literal["adamuon", "normuon"] = "adamuon", beta2: float = 0.95, eps: float = 1e-8, @@ -444,6 +459,7 @@ def __init__( extra_scale_factor=extra_scale_factor, pg_collection=pg_collection, tp_mode=tp_mode, + use_syrk=use_syrk, ) self.moment2_method = moment2_method diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 24f9a032c47..20045108c89 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -282,6 +282,9 @@ class OptimizerConfig: muon_tp_mode: str = "blockwise" """How to perform NS calculation for tensor parallel weights. Defaults to "blockwise".""" + muon_use_syrk: bool = False + """Use the Triton SYRK kernel for the Gram matrix in Newton-Schulz iteration.""" + muon_extra_scale_factor: float = 1.0 """Additional scale factor for the muon update.""" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 623adcad91e..50633aefa26 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2554,6 +2554,9 @@ def _add_regularization_args(parser): group.add_argument('--muon-tp-mode', type=str, default='blockwise', choices=['blockwise', 'duplicated', 'distributed'], help='How to perform NS calculation for tensor model parallel weights') + group.add_argument('--muon-use-syrk', action='store_true', + help='Use the Triton SYRK kernel for the Gram matrix ' + 'in Newton-Schulz iteration.') group.add_argument('--muon-extra-scale-factor', type=float, default=1.0, help='Additional scale factor for the muon update') group.add_argument('--muon-scalar-optimizer', type=str, default='adam',