diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 53ac956b35c..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 @@ -230,13 +244,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. - 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. + When GTP is inactive this is a plain passthrough to ``scaled_orthogonalize_fn``. + Otherwise, ``self.tp_mode`` controls how GTP sharding is handled: + + - **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 +281,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. @@ -376,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, @@ -398,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/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 376f9a1f1c0..79d50e4fe6c 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 @@ -312,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): @@ -332,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 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/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, 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', 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 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