Fix Muon QKV split for gated attention

megatron/core/optimizer/__init__.py

@@ -65,6 +65,7 @@
     _EMERGING_OPTIMIZERS,
     HAVE_EMERGING_OPTIMIZERS,
     _create_emerging_optimizer,
+    _get_qkv_split_shapes,
 )
 from .grad_scaler import ConstantGradScaler, DynamicGradScaler
 from .layer_wise_optimizer import LayerWiseDistributedOptimizer
@@ -769,14 +770,26 @@ def _get_megatron_emerging_optimizer(
 
     # Tag parameters with optimizer-specific attributes (expert_tp, is_qkv).
     for model_chunk in model_chunks:
+        qkv_split_shapes = None
         for name, param in model_chunk.named_parameters():
             if not param.requires_grad:
                 continue
             if 'experts' in name and 'shared' not in name:
                 param.expert_tp = True
             # TODO(deyuf): support MLA
             if 'linear_qkv.weight' in name and len(param.shape) == 2:
-                param.is_qkv = True
+                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:
+                    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}",
+                    )
 
     # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam).
     config_overrides.update(_EMERGING_OPTIMIZERS[eopt_name].default_param_overrides)

megatron/core/optimizer/emerging_optimizers.py

@@ -11,7 +11,7 @@
 import inspect
 import logging
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Literal, Optional, get_args
+from typing import Any, Callable, Dict, Literal, Optional, get_args
 
 import torch
 from torch.optim.optimizer import ParamsT
@@ -130,13 +130,19 @@ def _is_nonlinear_or_embedding(param):
     return getattr(param, 'is_embedding_or_output_parameter', False) or len(param.shape) != 2
 
 
-def _get_qkv_split_shapes(model_cfg) -> List[int]:
+def _get_qkv_split_shapes(model_cfg) -> list[int]:
     """Compute QKV split shapes from model config."""
-    return [
-        model_cfg.num_attention_heads // model_cfg.num_query_groups * model_cfg.kv_channels,
-        model_cfg.kv_channels,
-        model_cfg.kv_channels,
-    ]
+    query_projection_size = (
+        model_cfg.num_attention_heads // model_cfg.num_query_groups * model_cfg.kv_channels
+    )
+    if getattr(model_cfg, 'attention_output_gate', False):
+        return [
+            query_projection_size,
+            query_projection_size,
+            model_cfg.kv_channels,
+            model_cfg.kv_channels,
+        ]
+    return [query_projection_size, model_cfg.kv_channels, model_cfg.kv_channels]
 
 
 # ===========================================================================
@@ -164,7 +170,7 @@ def __init__(
         use_decoupled_weight_decay: bool = True,
         split_qkv: bool = False,
         is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
-        qkv_split_shapes: tuple[int, int, int] | None = None,
+        qkv_split_shapes: list[int] | None = None,
         fp32_matmul_prec: str = "medium",
         coefficient_type: str = "quintic",
         num_ns_steps: int = 5,
@@ -251,16 +257,25 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t
 
         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:
+                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}, ' f'split shapes {self.qkv_split_shapes}',
+                f'qkv split grad shape {grad_shape}, split shapes {qkv_split_shapes}',
             )
-            num_query_groups = grad_shape[0] // sum(self.qkv_split_shapes)
+            num_query_groups = grad_shape[0] // qkv_split_dim
             qkv_grads = torch.split(
-                grad.view(num_query_groups, sum(self.qkv_split_shapes), -1),
-                self.qkv_split_shapes,
-                dim=1,
+                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]
 
@@ -317,7 +332,7 @@ def __init__(
         use_decoupled_weight_decay: bool = True,
         split_qkv: bool = False,
         is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
-        qkv_split_shapes: tuple[int, int, int] | None = None,
+        qkv_split_shapes: list[int] | None = None,
         fp32_matmul_prec: str = "medium",
         coefficient_type: str = "quintic",
         num_ns_steps: int = 5,

megatron/core/tensor_parallel/layers.py

@@ -60,6 +60,7 @@
 _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = {
     "expert_tp": False,
     "is_qkv": False,
+    "qkv_split_shapes": None,
     "tensor_model_parallel": False,
     "partition_dim": -1,
     "partition_stride": 1,

tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py

@@ -7,6 +7,7 @@
     ColumnParallelLinear,
     RowParallelLinear,
     VocabParallelEmbedding,
+    copy_tensor_model_parallel_attributes,
 )
 from megatron.core.transformer.transformer_config import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
@@ -87,3 +88,15 @@ def test_row_parallel_linear_tp_attrs_no_init(self, use_cpu_init):
         assert hasattr(w, "tensor_model_parallel") and w.tensor_model_parallel is True
         assert hasattr(w, "partition_dim") and w.partition_dim == 1
         assert hasattr(w, "partition_stride") and w.partition_stride == 1
+
+
+def test_copy_tensor_model_parallel_attributes_preserves_qkv_split_shapes():
+    source = torch.empty(4, 4)
+    destination = torch.empty_like(source)
+    source.is_qkv = True
+    source.qkv_split_shapes = [256, 64, 64]
+
+    copy_tensor_model_parallel_attributes(destination, source)
+
+    assert destination.is_qkv is True
+    assert destination.qkv_split_shapes == source.qkv_split_shapes

tests/unit_tests/test_emerging_optimizers.py

@@ -15,6 +15,7 @@
     HAVE_EMERGING_OPTIMIZERS,
     TensorParallelAdaptiveMuon,
     TensorParallelMuon,
+    _get_qkv_split_shapes,
     get_supported_coefficient_types,
     validate_coefficient_type,
 )
@@ -65,6 +66,22 @@ def forward(self, x):
 # ===========================================================================
 
 
+def test_muon_qkv_split_shapes():
+    config = TransformerConfig(
+        num_layers=1, hidden_size=1024, num_attention_heads=16, num_query_groups=8
+    )
+    gated_config = TransformerConfig(
+        num_layers=1,
+        hidden_size=1024,
+        num_attention_heads=16,
+        num_query_groups=8,
+        attention_output_gate=True,
+    )
+
+    assert _get_qkv_split_shapes(config) == [128, 64, 64]
+    assert _get_qkv_split_shapes(gated_config) == [128, 128, 64, 64]
+
+
 def test_muon_optimizer_smoke():
     """Smoke test for TensorParallelMuon optimizer."""
     # Create a simple linear model for testing