[DO NOT MERGE] GTP+Muon rebased onto main for testing

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,13 +207,17 @@ 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,
                 coefficient_type=coefficient_type,
                 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
 

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

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."""
 

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,

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',

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)

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

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