fix(muon): split the fully-gathered QKV tensor before orthogonalizing under GTP

docs/api-guide/core/generalized_tensor_parallel.md

@@ -190,6 +190,10 @@ GTP_remat runs under both the standard **Adam** `DistributedOptimizer` and **Muo
 
 - **Adam** shards optimizer state over the gtp_remat/egtp_remat-excluded replicate group, like any GTP_remat run (§3.2).
 - **Muon** keeps matrix params *whole* (Newton–Schulz needs the full 2D weight). A GTP_remat-replicated whole param (e.g. MoE router, latent-proj MLPs by default) then lands on one checkpoint key shared by all GTP_remat peers, so the LayerWise optimizer folds `gtp_rank` into its `replica_id` — exactly one peer writes (the optimizer-state analog of the model-side fold in §3.3).
+- **Split-QKV is layout-invariant.** Muon orthogonalizes a fused `linear_qkv` weight as separate **q**, **k** and **v** blocks, and *whether* it does must not depend on the parallel layout. Splitting and not splitting are **two different update rules**, not small perturbations of one: q is typically far larger than k and v, so orthogonalizing the fused matrix lets q dominate the spectrum the k and v rows are scaled by. Newton–Schulz is a nonlinear iteration, so the gap compounds over steps rather than averaging out — the same weight would train differently at TP1 than under GTP_remat.
+  - **Decide on the across-shards row count** — `local × gtp_remat_size − pad` — **never on this rank's shard.** GTP_remat row-shards dim 0, so a shard-local divisibility test flips to `False` the moment the GTP_remat degree stops dividing the query-group count: the split then vanishes on every GTP_remat rank while TP1, holding all the rows, keeps it.
+  - **Split *after* the all-gather**, never before — a row shard cuts q/k/v mid-boundary. The split is therefore reachable on the **`duplicated`** Newton–Schulz mode only (`--muon-tp-mode`, default `duplicated`).
+  - **`blockwise` and `distributed` fall back** to whole-matrix Newton–Schulz and **warn once.** Both keep the rows sharded through NS, so no rank holds a q/k/v boundary. That is what those modes already did under GTP_remat, so **existing runs on them are unchanged**; `duplicated`, the default, is the path that gets the split.
 - **Native-FP8 optimizer-state matching (Muon path).** The save-side dequantize (§3.3) hands DCP a *fresh* BF16 tensor, which breaks the id-based optimizer-param → model-`ShardedTensor` match for every native-FP8 GTP_remat weight. The dequantized copy carries a `_gtp_dequant_src` backlink to the live FP8 param, and `_backfill_gtp_sharded_param_map` reuses the model's **own** entry (backlink first, tagged-name second) — preserving its full offsets (expert axes included) and `replica_id`. Only truly-unmatched params (the SSM `in_proj` weights, gathered+split factories) take the per-shard rebuild, which refuses expert-parallel params rather than emit EP-colliding shards.
 
 Neither path adds a GTP_remat-specific checkpoint format or call site.
@@ -867,7 +871,9 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall
 | `test_gtp_cudagraph_grad.py` | Capture-step grad-norm guard (§1.2): `_backup_grads_before_capture`/`_restore_grads_after_capture` keep a graph capture from clobbering finalized `main_grad` (own params + cross-graph `next_w`, incl. routed-expert `weight_list`). |
 | `test_gtp_partial_cg.py` | Four-layer partial-CG loss and eager-vs-replay grad-norm parity with two-slot ring reuse across independently replayed graphs (§3.5). |
 | `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. Also the SSM `in_proj` gather+split: the gated-delta-product mixer's factory build/merge at MXFP8 alignment, and a full DCP save→load roundtrip of that mixer. |
+| `test_gtp_muon.py` | Newton–Schulz per-shard parity for all three GTP_remat modes (§1.6): blockwise/duplicated/distributed at TP1, plus row- and column-parallel cases at `tp_size != gtp_remat_size` in both directions. |
 | `test_gtp_muon_dcp.py` | Muon optimizer-state DCP roundtrip (§1.6): `replica_id` fold + native-FP8 backfill matching. |
+| `test_gtp_muon_qkv.py` | Layout-invariant split-QKV (§1.6): the decision uses the across-shards row count (with the production shape that a shard-local test rejected), the split runs after the all-gather so the GTP_remat shard equals TP1's result restricted to this rank's rows, and `blockwise`/`distributed` fall back to whole-matrix Newton–Schulz bitwise, warning once. |
 | `test_gtp_recompute_chain.py` | Recompute-chain buffers (§3.1): adjacent nodes never share a gather buffer, dense and grouped, plus dgrad/wgrad parity vs no-recompute. |
 | `test_gtp_mtp.py` | GTP_remat + MTP shared weights (§3.5), 14 cases over `mtp_use_repeated_layer` × dense/MoE. Both MTP hazards are silent, so each needs its own guard: the async reduce-scatter path is compared numerically against the sync path on an identical model/sharding/batch, and all-gathers issued are tallied against consumes to catch a consume reading a buffer nothing gathered into. |
 | `test_gtp_fp8_param_gather.py` | Native-FP8 GTP_remat (§1.3): fp8-vs-BF16 loss parity (TP1/TP2, MoE), post-save-spike guard. |

megatron/core/optimizer/__init__.py

@@ -53,7 +53,7 @@
     combine_param_group_overrides,
     param_group_override_to_tuple,
 )
-from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.process_groups_config import ProcessGroupCollection, resolve_gtp_remat_group
 from megatron.core.transformer.fsdp_dtensor_checkpoint import get_global_unique_param_name
 
 from ..distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallelV2
@@ -732,6 +732,42 @@ def check_config_overrides_consistency(
     return True
 
 
+def qkv_rows_after_gtp_gather(
+    param: torch.nn.Parameter,
+    qkv_split_shapes: List[int],
+    gtp_remat_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> Tuple[int, int, bool]:
+    """Total rows of a fused QKV weight across its GTP shards, and whether [q|k|v] divides them.
+
+    Not ``param.shape[0]``: once the GTP degree stops dividing the query-group count, a
+    shard-local test switches the split off on GTP ranks while TP1 keeps it, and one
+    weight gets two different Muon rules.
+
+    The count is still TP-local, which is enough to decide because TP splits whole query
+    groups -- but it is not model-global, so do not reuse it for checkpoint shapes.
+
+    Args:
+        param: Fused QKV weight; reads ``is_gtp_weight_remat`` and ``pad_length``.
+        qkv_split_shapes: Per-query-group rows, e.g. ``[q, k, v]``. Must be non-empty.
+        gtp_remat_group: The group the optimizer gathers over -- ``param.group`` can name
+            a different one under a ``MultiModuleProcessGroupCollection``.
+
+    Returns:
+        ``(gathered_rows, gtp_size, splittable)``, padding already subtracted.
+    """
+    if not qkv_split_shapes:
+        raise ValueError("qkv_split_shapes must be non-empty to decide the QKV split")
+    gtp_size = (
+        get_pg_size(gtp_remat_group)
+        if gtp_remat_group is not None and getattr(param, 'is_gtp_weight_remat', False)
+        else 1
+    )
+    # Subtract the pad only when the rows were scaled up.
+    alignment_pad_rows = getattr(param, 'pad_length', 0) if gtp_size > 1 else 0
+    gathered_rows = param.shape[0] * gtp_size - alignment_pad_rows
+    return gathered_rows, gtp_size, gathered_rows % sum(qkv_split_shapes) == 0
+
+
 def _get_megatron_emerging_optimizer(
     config: OptimizerConfig,
     model_chunks: List[MegatronModule],
@@ -800,15 +836,22 @@ def _get_megatron_emerging_optimizer(
             if 'linear_qkv.weight' in name and len(param.shape) == 2:
                 if qkv_split_shapes is None:
                     qkv_split_shapes = _get_qkv_split_shapes(model_chunk.config)
-                if param.shape[0] % sum(qkv_split_shapes) == 0:
+                gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert=False)
+                gathered_rows, gtp_size, splittable = qkv_rows_after_gtp_gather(
+                    param, qkv_split_shapes, gtp_remat_group
+                )
+                if splittable:
                     param.is_qkv = True
                     param.qkv_split_shapes = qkv_split_shapes
                 else:
                     log_single_rank(
                         logger,
-                        logging.DEBUG,
-                        f"Emerging optimizer QKV split skipped for {name}: "
-                        f"shape={tuple(param.shape)}, split_shapes={qkv_split_shapes}",
+                        logging.INFO,
+                        f"Emerging optimizer QKV split disabled for {name}: gathered rows "
+                        f"{gathered_rows} (local {param.shape[0]} x gtp {gtp_size} - pad "
+                        f"{param.shape[0] * gtp_size - gathered_rows}) not divisible by "
+                        f"{sum(qkv_split_shapes)}={qkv_split_shapes}; whole-matrix "
+                        f"Newton-Schulz at every layout.",
                     )
 
     # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam).

megatron/core/optimizer/emerging_optimizers.py

@@ -16,7 +16,7 @@
 import torch
 from torch.optim.optimizer import ParamsT
 
-from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.process_groups_config import ProcessGroupCollection, resolve_gtp_remat_group
 from megatron.core.utils import (
     get_emerging_optimizers_version,
     get_pg_rank,
@@ -321,6 +321,7 @@ def scaled_orthogonalize_fn(
         self.elem_size = 2 if fp32_matmul_prec == "medium" else 4  # bf16 vs tf32/fp32
         self._tp_mode_cache: Dict[tuple, str] = {}
         self._hw_profile = _hardware_profile() if tp_mode == "auto" else None
+        self._warned_qkv_split_disabled = False
 
         weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2"
         # Use explicit class call instead of super() so that subclasses with
@@ -381,7 +382,9 @@ def _resolve_tp_mode(self, m: int, n: int, group_size: int) -> str:
             )
         return self._tp_mode_cache[key]
 
-    def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_dim):
+    def scaled_orthogonalize_fn_with_gtp_remat(
+        self, p, grad, tp_group, partition_dim, qkv_split_shapes=None
+    ):
         """Orthogonalize a (possibly GTP-sharded) momentum, then reshard.
 
         When GTP is inactive this is a plain passthrough to ``scaled_orthogonalize_fn``.
@@ -399,13 +402,20 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di
         after, since every rank holds a uniform, fully-reconstructed tensor by then.
         distributed does not: it stays row-sharded through its own collective, where
         stripping isn't safe (known limitation).
+
+        ``qkv_split_shapes`` (when set) runs Newton-Schulz on q, k and v separately. The
+        split needs the whole matrix, so under GTP it is available on the duplicated path
+        only -- blockwise and distributed fall back to whole-matrix NS and warn once.
         """
+        if qkv_split_shapes is not None and not qkv_split_shapes:
+            raise ValueError(
+                "qkv_split_shapes must be None (no split) or non-empty; got an empty sequence"
+            )
+
         # 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)
         gtp_remat_group = (
-            (self.pg_collection.expt_gtp_remat if is_expert else self.pg_collection.gtp_remat)
-            if self.pg_collection
-            else None
+            resolve_gtp_remat_group(self.pg_collection, is_expert) if self.pg_collection else None
         )
 
         # Parameters with is_gtp_weight_remat=False are not sharded along the
@@ -427,14 +437,63 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di
                 else "duplicated"
             )
 
-        if not gtp_active:
-            return self.scaled_orthogonalize_fn(
-                grad, tp_group, partition_dim, tp_mode_this_group=mode
+        def orthogonalize_whole_matrix(whole_grad):
+            """Orthogonalize `whole_grad`, splitting [q|k|v] first when asked.
+
+            `whole_grad` is always the WHOLE matrix -- GTP is inactive, or the caller
+            has already gathered. A row shard would cut q/k/v mid-boundary.
+            """
+            if qkv_split_shapes is None:
+                return self.scaled_orthogonalize_fn(
+                    whole_grad, tp_group, partition_dim, tp_mode_this_group=mode
+                )
+            qkv_rows = sum(qkv_split_shapes)
+            if whole_grad.size(0) % qkv_rows != 0:
+                raise RuntimeError(
+                    f"Muon QKV split shape mismatch on the whole matrix: "
+                    f"rows={whole_grad.size(0)}, qkv_split_shapes={qkv_split_shapes} "
+                    f"(sum {qkv_rows}). The row count is decided by the optimizer's qkv "
+                    f"tagging loop (qkv_rows_after_gtp_gather), not here."
+                )
+            num_query_groups = whole_grad.size(0) // qkv_rows
+            cols = whole_grad.size(-1)
+            qkv_grads = torch.split(
+                whole_grad.view(num_query_groups, qkv_rows, cols), qkv_split_shapes, dim=1
             )
+            qkv_grads = [
+                self.scaled_orthogonalize_fn(
+                    g.reshape(-1, cols), tp_group, partition_dim, tp_mode_this_group=mode
+                ).view(num_query_groups, -1, cols)
+                for g in qkv_grads
+            ]
+            return torch.cat(qkv_grads, dim=1).view(whole_grad.shape)
+
+        if not gtp_active:
+            return orthogonalize_whole_matrix(grad)
 
         gtp_rank = get_pg_rank(gtp_remat_group)
         pad_length = getattr(p, 'pad_length', 0)
 
+        if qkv_split_shapes is not None and mode != "duplicated":
+            # Fall back, not raise: these configs ran before the qkv fix, so an abort
+            # here would break live jobs.
+            if not self._warned_qkv_split_disabled:
+                self._warned_qkv_split_disabled = True
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    f"Muon: split-QKV is disabled for GTP-sharded qkv weights under "
+                    f"tp_mode='{mode}'"
+                    + (f" (auto-resolved from '{self.tp_mode}')" if self.tp_mode == "auto" else "")
+                    + "; blockwise/distributed keep the rows sharded, so the q/k/v "
+                    "boundaries are not available on any rank. These weights get "
+                    "whole-matrix Newton-Schulz, which is NOT the update rule the same "
+                    "model gets at TP1. Pass --muon-tp-mode duplicated for a "
+                    "layout-invariant qkv update, or --muon-no-split-qkv to make the "
+                    "whole-matrix rule explicit at every layout.",
+                )
+            qkv_split_shapes = None
+
         if 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
@@ -457,12 +516,7 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di
             # holds the same padded tensor), orthogonalize the whole matrix
             # (scaled_orthogonalize_fn handles any TP sharding per tp_mode), reshard dim 0.
             gathered_grad = self._all_gather_tensor(grad, gtp_remat_group, 0)
-            result = self.scaled_orthogonalize_fn(
-                self._strip_pad(gathered_grad, pad_length),
-                tp_group,
-                partition_dim,
-                tp_mode_this_group=mode,
-            )
+            result = orthogonalize_whole_matrix(self._strip_pad(gathered_grad, pad_length))
             result = self._restore_pad(result, pad_length)
             reshard_size = result.size(0) // gtp_remat_size
             return result[gtp_rank * reshard_size : (gtp_rank + 1) * reshard_size].contiguous()
@@ -532,39 +586,18 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t
         if partition_dim == -1:
             partition_dim = None
 
+        # Look up the shapes only. `grad` is still this rank's row shard here, so the
+        # split has to wait until the callee has gathered the whole matrix.
+        qkv_split_shapes = None
         if self.split_qkv and self.is_qkv_fn(p):  # type: ignore[misc]
-            grad_shape = grad.shape
             qkv_split_shapes = getattr(p, "qkv_split_shapes", None)
             if qkv_split_shapes is None:
                 qkv_split_shapes = self.qkv_split_shapes
-            if qkv_split_shapes is None:
+            if not qkv_split_shapes:
                 raise RuntimeError("Muon QKV split requested but qkv_split_shapes is not set")
-            qkv_split_dim = sum(qkv_split_shapes)
-            if grad_shape[0] % qkv_split_dim != 0:
-                raise RuntimeError(
-                    f"Muon QKV split shape mismatch: grad_shape={tuple(grad_shape)}, "
-                    f"split_shapes={qkv_split_shapes}"
-                )
-            log_single_rank(
-                logger,
-                logging.DEBUG,
-                f'qkv split grad shape {grad_shape}, split shapes {qkv_split_shapes}',
-            )
-            num_query_groups = grad_shape[0] // qkv_split_dim
-            qkv_grads = torch.split(
-                grad.view(num_query_groups, qkv_split_dim, -1), qkv_split_shapes, dim=1
-            )
-            qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads]
-
-            qkv_grads = [
-                self.scaled_orthogonalize_fn_with_gtp_remat(p, g, tp_group, partition_dim).view(
-                    num_query_groups, -1, grad_shape[-1]
-                )
-                for g in qkv_grads
-            ]
-            grad = torch.cat(qkv_grads, dim=1).view(grad_shape)
-        else:
-            grad = self.scaled_orthogonalize_fn_with_gtp_remat(p, grad, tp_group, partition_dim)
+        grad = self.scaled_orthogonalize_fn_with_gtp_remat(
+            p, grad, tp_group, partition_dim, qkv_split_shapes=qkv_split_shapes
+        )
         return grad