From 2fb747fe81ee06d4aecbbde14698b445e31bef5e Mon Sep 17 00:00:00 2001 From: mkhona Date: Thu, 6 Aug 2026 10:48:13 -0700 Subject: [PATCH 1/8] Add per-head Muon QKV orthogonalization Signed-off-by: mkhona --- megatron/core/optimizer/__init__.py | 62 +++- .../core/optimizer/emerging_optimizers.py | 282 ++++++++++++++++-- megatron/core/optimizer/optimizer_config.py | 5 + megatron/core/tensor_parallel/layers.py | 3 + megatron/training/arguments.py | 3 + .../test_tp_attrs_without_init.py | 8 +- tests/unit_tests/test_emerging_optimizers.py | 180 +++++++++++ 7 files changed, 504 insertions(+), 39 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index f8f5a813b38..9faab7ddf07 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -66,6 +66,8 @@ HAVE_EMERGING_OPTIMIZERS, _create_emerging_optimizer, _get_qkv_split_shapes, + _localize_qkv_split_shapes, + _qkv_split_groups_are_complete, ) from .fully_sharded_optimizer import FullyShardedOptimizer from .grad_scaler import ConstantGradScaler, DynamicGradScaler @@ -771,6 +773,8 @@ def _get_megatron_emerging_optimizer( raise ValueError(f"Unsupported emerging optimizer: {eopt_name}") if config.fp16: raise ValueError('emerging optimizer with fp16 is not supported.') + if config.muon_split_qkv_per_head and not config.muon_split_qkv: + raise ValueError("muon_split_qkv_per_head requires muon_split_qkv=True") if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -788,16 +792,56 @@ def _get_megatron_emerging_optimizer( # TODO(deyuf): support MLA 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: - param.is_qkv = True - param.qkv_split_shapes = qkv_split_shapes + qkv_split_shapes = _get_qkv_split_shapes( + model_chunk.config, split_qkv_per_head=config.muon_split_qkv_per_head + ) + param.is_qkv = True + global_split_shapes = ( + qkv_split_shapes + if config.muon_split_qkv_per_head + else qkv_split_shapes * model_chunk.config.num_query_groups + ) + param.qkv_split_shapes_global = global_split_shapes + + tp_group = ( + pg_collection.expt_tp + if getattr(param, 'expert_tp', False) + else pg_collection.tp + ) + tp_size = get_pg_size(tp_group) + tp_rank = get_pg_rank(tp_group) + gtp_remat_group = ( + pg_collection.expt_gtp_remat + if getattr(param, 'expert_tp', False) + else pg_collection.gtp_remat + ) + if getattr(param, 'is_gtp_weight_remat', False): + gtp_size = get_pg_size(gtp_remat_group) + gtp_rank = get_pg_rank(gtp_remat_group) 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}", + gtp_size = 1 + gtp_rank = 0 + + tp_local_rows = param.shape[0] * gtp_size + expected_global_rows = tp_local_rows * tp_size + if expected_global_rows != sum(global_split_shapes): + raise RuntimeError( + f"Muon QKV layout mismatch for {name}: " + f"global_rows={sum(global_split_shapes)}, " + f"local_rows={param.shape[0]}, tp_size={tp_size}, " + f"gtp_remat_size={gtp_size}" + ) + local_start = tp_rank * tp_local_rows + gtp_rank * param.shape[0] + if config.muon_split_qkv_per_head: + param.qkv_split_shapes, param.qkv_split_heads_are_complete = ( + _localize_qkv_split_shapes( + qkv_split_shapes, local_start=local_start, local_rows=param.shape[0] + ) + ) + else: + param.qkv_split_shapes = qkv_split_shapes + param.qkv_split_groups_are_complete = _qkv_split_groups_are_complete( + qkv_split_shapes, local_start=local_start, local_rows=param.shape[0] ) # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam). diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 53ac956b35c..d315ad1e376 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -130,11 +130,24 @@ 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]: - """Compute QKV split shapes from model config.""" +def _get_qkv_split_shapes(model_cfg, split_qkv_per_head: bool = False) -> list[int]: + """Compute fused QKV split shapes from a transformer model config. + + Args: + model_cfg: Transformer model configuration. + split_qkv_per_head: Return one split size per physical attention head. When false, + return the per-query-group Q, gate (if present), K, and V projection widths. + """ query_projection_size = ( model_cfg.num_attention_heads // model_cfg.num_query_groups * model_cfg.kv_channels ) + if split_qkv_per_head: + num_query_heads_per_group = model_cfg.num_attention_heads // model_cfg.num_query_groups + per_group_shapes = [model_cfg.kv_channels] * num_query_heads_per_group + if getattr(model_cfg, 'attention_output_gate', False): + per_group_shapes += [model_cfg.kv_channels] * num_query_heads_per_group + per_group_shapes += [model_cfg.kv_channels, model_cfg.kv_channels] + return per_group_shapes * model_cfg.num_query_groups if getattr(model_cfg, 'attention_output_gate', False): return [ query_projection_size, @@ -145,6 +158,46 @@ def _get_qkv_split_shapes(model_cfg) -> list[int]: return [query_projection_size, model_cfg.kv_channels, model_cfg.kv_channels] +def _localize_qkv_split_shapes( + global_split_shapes: list[int], local_start: int, local_rows: int +) -> tuple[list[int], bool]: + """Intersect global per-head split sizes with a rank-local contiguous row range. + + Returns: + The physical rank-local split sizes and whether every intersected head is complete. + """ + local_stop = local_start + local_rows + local_split_shapes = [] + all_heads_complete = True + head_start = 0 + for head_rows in global_split_shapes: + head_stop = head_start + head_rows + overlap_start = max(head_start, local_start) + overlap_stop = min(head_stop, local_stop) + if overlap_start < overlap_stop: + overlap_rows = overlap_stop - overlap_start + local_split_shapes.append(overlap_rows) + all_heads_complete &= overlap_rows == head_rows + head_start = head_stop + + if sum(local_split_shapes) != local_rows: + raise RuntimeError( + f"Muon per-head QKV local range [{local_start}, {local_stop}) is outside " + f"the global split shape with {sum(global_split_shapes)} rows" + ) + return local_split_shapes, all_heads_complete + + +def _qkv_split_groups_are_complete( + split_shapes: list[int], local_start: int, local_rows: int +) -> bool: + """Return whether a local row range contains only complete fused QKV groups.""" + split_width = sum(split_shapes) + if split_width <= 0: + raise ValueError(f"Muon QKV split shapes must sum to a positive size: {split_shapes}") + return local_start % split_width == 0 and local_rows % split_width == 0 + + # =========================================================================== # Registry – populated below only when emerging_optimizers is installed. # =========================================================================== @@ -169,6 +222,7 @@ def __init__( weight_decay: float = 0.01, use_decoupled_weight_decay: bool = True, split_qkv: bool = False, + split_qkv_per_head: bool = False, is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, qkv_split_shapes: list[int] | None = None, fp32_matmul_prec: str = "medium", @@ -181,6 +235,8 @@ def __init__( ) -> None: if num_ns_steps < 1: raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") + if split_qkv_per_head and not split_qkv: + raise ValueError("split_qkv_per_head requires split_qkv=True") def scaled_orthogonalize_fn( grad: torch.Tensor, @@ -211,6 +267,7 @@ def scaled_orthogonalize_fn( self.pg_collection = pg_collection self.tp_mode = tp_mode self.split_qkv = split_qkv + self.split_qkv_per_head = split_qkv_per_head self.is_qkv_fn = is_qkv_fn self.qkv_split_shapes = qkv_split_shapes @@ -230,6 +287,114 @@ def scaled_orthogonalize_fn( scaled_orthogonalize_fn=scaled_orthogonalize_fn, ) + def _get_gtp_remat_group(self, p): + """Return the GTP-remat process group for a parameter, if configured.""" + is_expert = getattr(p, 'expert_tp', False) + return ( + (self.pg_collection.expt_gtp_remat if is_expert else self.pg_collection.gtp_remat) + if self.pg_collection + else None + ) + + def _gather_qkv_grad(self, p, grad, tp_group, expected_rows, gather_gtp=True): + """Reconstruct a fused QKV gradient and record how to restore its local shard.""" + gathered_grad = grad + gtp_slice = None + gtp_remat_group = self._get_gtp_remat_group(p) + if ( + gather_gtp + and gtp_remat_group is not None + and get_pg_size(gtp_remat_group) > 1 + and getattr(p, 'is_gtp_weight_remat', False) + ): + gtp_size = get_pg_size(gtp_remat_group) + gtp_rank = get_pg_rank(gtp_remat_group) + gtp_local_rows = gathered_grad.shape[0] + shards = [torch.empty_like(gathered_grad) for _ in range(gtp_size)] + torch.distributed.all_gather(shards, gathered_grad, gtp_remat_group) + gathered_grad = torch.cat(shards, dim=0) + gtp_slice = (gtp_rank, gtp_local_rows) + + tp_slice = None + if gathered_grad.shape[0] != expected_rows: + partition_dim = getattr(p, "partition_dim", None) + if partition_dim != 0 or tp_group is None: + raise RuntimeError( + f"Muon QKV split shape mismatch: grad_shape={tuple(gathered_grad.shape)}, " + f"expected_rows={expected_rows}, partition_dim={partition_dim}" + ) + tp_size = get_pg_size(tp_group) + if gathered_grad.shape[0] * tp_size != expected_rows: + raise RuntimeError( + "Muon QKV split cannot reconstruct the global tensor: " + f"local_grad_shape={tuple(gathered_grad.shape)}, tp_size={tp_size}, " + f"expected_rows={expected_rows}" + ) + tp_rank = get_pg_rank(tp_group) + tp_local_rows = gathered_grad.shape[0] + shards = [torch.empty_like(gathered_grad) for _ in range(tp_size)] + torch.distributed.all_gather(shards, gathered_grad, tp_group) + gathered_grad = torch.cat(shards, dim=0) + tp_slice = (tp_rank, tp_local_rows) + + if gathered_grad.shape[0] != expected_rows: + raise RuntimeError( + "Muon QKV split shape mismatch after gathering: " + f"grad_shape={tuple(gathered_grad.shape)}, expected_rows={expected_rows}" + ) + return gathered_grad, tp_slice, gtp_slice + + @staticmethod + def _restore_local_qkv_grad(gathered_grad, tp_slice, gtp_slice): + """Restore the TP and GTP-remat shards recorded by ``_gather_qkv_grad``.""" + if tp_slice is not None: + tp_rank, tp_local_rows = tp_slice + gathered_grad = gathered_grad[tp_rank * tp_local_rows : (tp_rank + 1) * tp_local_rows] + if gtp_slice is not None: + gtp_rank, gtp_local_rows = gtp_slice + gathered_grad = gathered_grad[ + gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows + ] + return gathered_grad.contiguous() + + def _orthogonalize_split_qkv(self, grad, split_shapes, orthogonalize_fn): + """Split and reconstruct Megatron's interleaved fused QKV update.""" + if grad.ndim != 2: + raise RuntimeError(f"Muon QKV gradient must be 2D, got {grad.ndim}D") + if not split_shapes or any(size <= 0 for size in split_shapes): + raise RuntimeError(f"Muon QKV split shapes must be positive: {split_shapes}") + split_width = sum(split_shapes) + if self.split_qkv_per_head: + if grad.shape[0] != split_width: + raise RuntimeError( + f"Muon per-head QKV split shape mismatch: grad_shape={tuple(grad.shape)}, " + f"split_shapes={split_shapes}" + ) + if len(set(split_shapes)) == 1: + # A 3D input selects Emerging-Optimizers' batched Newton-Schulz path. + head_rows = split_shapes[0] + return orthogonalize_fn(grad.view(len(split_shapes), head_rows, -1)).view_as(grad) + return torch.cat( + [orthogonalize_fn(head) for head in torch.split(grad, split_shapes, dim=0)], dim=0 + ) + + if grad.shape[0] % split_width != 0: + raise RuntimeError( + f"Muon QKV split shape mismatch: grad_shape={tuple(grad.shape)}, " + f"split_shapes={split_shapes}" + ) + num_query_groups = grad.shape[0] // split_width + grouped_grad = grad.view(num_query_groups, split_width, -1) + projection_grads = torch.split(grouped_grad, split_shapes, dim=1) + projection_grads = [ + projection.reshape(-1, grad.shape[-1]) for projection in projection_grads + ] + projection_grads = [ + orthogonalize_fn(projection).view(num_query_groups, -1, grad.shape[-1]) + for projection in projection_grads + ] + return torch.cat(projection_grads, dim=1).view_as(grad) + 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. @@ -239,12 +404,7 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di When GTP_remat is inactive this is a plain passthrough to scaled_orthogonalize_fn. """ # 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 - ) + gtp_remat_group = self._get_gtp_remat_group(p) if gtp_remat_group is None or get_pg_size(gtp_remat_group) <= 1: return self.scaled_orthogonalize_fn(grad, tp_group, partition_dim) @@ -266,6 +426,70 @@ def scaled_orthogonalize_fn_with_gtp_remat(self, p, grad, tp_group, partition_di shard_size = gathered_grad.shape[0] // gtp_remat_size return gathered_grad[gtp_rank * shard_size : (gtp_rank + 1) * shard_size].contiguous() + def _orthogonalize_qkv_per_head(self, p, grad, tp_group): + """Orthogonalize every Q, gate, K, and V head independently. + + Split sizes may describe complete heads in the local tensor or the global fused + QKV tensor. For a global layout, reconstruct GTP-remat and TP dimension 0 before + splitting so heads crossing rank boundaries remain complete. + """ + local_split_shapes = getattr(p, "qkv_split_shapes", None) + heads_are_complete = getattr(p, "qkv_split_heads_are_complete", None) + use_local_layout = heads_are_complete is True or ( + heads_are_complete is None + and local_split_shapes is not None + and sum(local_split_shapes) == grad.shape[0] + ) + if use_local_layout: + qkv_split_shapes = local_split_shapes + else: + qkv_split_shapes = getattr(p, "qkv_split_shapes_global", None) + if qkv_split_shapes is None: + qkv_split_shapes = self.qkv_split_shapes + if qkv_split_shapes is None: + raise RuntimeError("Muon per-head QKV split requested but qkv_split_shapes is not set") + if not qkv_split_shapes or any(size <= 0 for size in qkv_split_shapes): + raise RuntimeError( + f"Muon per-head QKV split shapes must be positive: {qkv_split_shapes}" + ) + + expected_rows = sum(qkv_split_shapes) + gathered_grad, tp_slice, gtp_slice = self._gather_qkv_grad( + p, grad, tp_group, expected_rows, gather_gtp=not use_local_layout + ) + + gathered_grad = self._orthogonalize_split_qkv( + gathered_grad, + qkv_split_shapes, + lambda head_grad: self.scaled_orthogonalize_fn( + head_grad, tp_group=None, partition_dim=None + ), + ) + + return self._restore_local_qkv_grad(gathered_grad, tp_slice, gtp_slice) + + def _orthogonalize_fragmented_qkv(self, p, grad, tp_group, split_shapes): + """Orthogonalize projections after reconstructing fragmented query-group blocks.""" + global_split_shapes = getattr(p, "qkv_split_shapes_global", None) + if global_split_shapes is None: + raise RuntimeError("Muon fragmented QKV split requires global split shapes") + expected_rows = sum(global_split_shapes) + if expected_rows % sum(split_shapes) != 0: + raise RuntimeError( + f"Muon global QKV layout does not contain complete query groups: " + f"global_split_shapes={global_split_shapes}, split_shapes={split_shapes}" + ) + + gathered_grad, tp_slice, gtp_slice = self._gather_qkv_grad(p, grad, tp_group, expected_rows) + gathered_grad = self._orthogonalize_split_qkv( + gathered_grad, + split_shapes, + lambda projection_grad: self.scaled_orthogonalize_fn( + projection_grad, tp_group=None, partition_dim=None + ), + ) + return self._restore_local_qkv_grad(gathered_grad, tp_slice, gtp_slice) + def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: """Orthogonalize the momentum. @@ -292,36 +516,31 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t partition_dim = None if self.split_qkv and self.is_qkv_fn(p): # type: ignore[misc] - grad_shape = grad.shape + if self.split_qkv_per_head: + return self._orthogonalize_qkv_per_head(p, grad, tp_group) + 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}" - ) + if ( + getattr(p, "qkv_split_groups_are_complete", None) is False + and getattr(p, "qkv_split_shapes_global", None) is not None + ): + return self._orthogonalize_fragmented_qkv(p, grad, tp_group, qkv_split_shapes) log_single_rank( logger, logging.DEBUG, - f'qkv split grad shape {grad_shape}, split shapes {qkv_split_shapes}', + 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 + grad = self._orthogonalize_split_qkv( + grad, + qkv_split_shapes, + lambda projection_grad: self.scaled_orthogonalize_fn_with_gtp_remat( + p, projection_grad, tp_group, partition_dim + ), ) - 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) return grad @@ -344,6 +563,7 @@ class TensorParallelAdaptiveMuon(TensorParallelMuon, AdaptiveMuon): weight_decay: Weight decay coefficient. use_decoupled_weight_decay: Whether to use decoupled weight decay. split_qkv: Whether to split QKV weights for orthogonalization. + split_qkv_per_head: Whether to orthogonalize individual Q, gate, K, and V heads. is_qkv_fn: Function to determine if a tensor is a QKV weight. qkv_split_shapes: Shapes for splitting QKV weights. fp32_matmul_prec: Precision for FP32 matrix multiplication. @@ -367,6 +587,7 @@ def __init__( weight_decay: float = 0.01, use_decoupled_weight_decay: bool = True, split_qkv: bool = False, + split_qkv_per_head: bool = False, is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, qkv_split_shapes: list[int] | None = None, fp32_matmul_prec: str = "medium", @@ -389,6 +610,7 @@ def __init__( weight_decay=weight_decay, use_decoupled_weight_decay=use_decoupled_weight_decay, split_qkv=split_qkv, + split_qkv_per_head=split_qkv_per_head, is_qkv_fn=is_qkv_fn, qkv_split_shapes=qkv_split_shapes, fp32_matmul_prec=fp32_matmul_prec, @@ -436,7 +658,9 @@ def _muon_config_to_kwargs(config, model_chunks, pg_collection) -> Dict[str, Any """Convert OptimizerConfig to TensorParallelMuon constructor kwargs.""" kwargs = _kwargs_from_config(TensorParallelMuon, "muon", config) kwargs["is_qkv_fn"] = lambda p: getattr(p, "is_qkv", False) - kwargs["qkv_split_shapes"] = _get_qkv_split_shapes(model_chunks[0].config) + kwargs["qkv_split_shapes"] = _get_qkv_split_shapes( + model_chunks[0].config, split_qkv_per_head=kwargs.get("split_qkv_per_head", False) + ) kwargs["pg_collection"] = pg_collection return kwargs diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 24f9a032c47..81a3eec0ac7 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -263,6 +263,11 @@ class OptimizerConfig: muon_split_qkv: bool = True """Whether to split QKV parameters for Muon optimizer.""" + muon_split_qkv_per_head: bool = False + """Whether to orthogonalize each Q, gate, K, and V head independently. Requires + ``muon_split_qkv``. By default, Muon orthogonalizes the Q, gate, K, and V + projection matrices separately.""" + muon_nesterov: bool = False """Whether to use Nesterov-style momentum in the internal SGD.""" diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 0f248bcf399..598e0f12951 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -64,6 +64,9 @@ "expert_tp": False, "is_qkv": False, "qkv_split_shapes": None, + "qkv_split_shapes_global": None, + "qkv_split_groups_are_complete": False, + "qkv_split_heads_are_complete": False, "tensor_model_parallel": False, "partition_dim": -1, "partition_stride": 1, diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d4f9eb9c0de..937417edaad 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2517,6 +2517,9 @@ def _add_regularization_args(parser): group.add_argument('--muon-no-split-qkv', action='store_false', default=True, dest='muon_split_qkv', help='Whether to split QKV parameters for Muon optimizer') + group.add_argument('--muon-split-qkv-per-head', action='store_true', + help='Orthogonalize each Q, gate, K, and V head independently. ' + 'By default, Q, gate, K, and V projections are orthogonalized separately') group.add_argument('--muon-nesterov', action='store_true', help='Whether to use Nesterov-style momentum in the internal SGD') group.add_argument('--muon-scale-mode', type=str, default='spectral', diff --git a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py index a76746d7674..983de4dcd6e 100644 --- a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py +++ b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py @@ -95,12 +95,18 @@ 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] + source.qkv_split_shapes = [2, 2] + source.qkv_split_shapes_global = [2] * 4 + source.qkv_split_groups_are_complete = True + source.qkv_split_heads_are_complete = True copy_tensor_model_parallel_attributes(destination, source) assert destination.is_qkv is True assert destination.qkv_split_shapes == source.qkv_split_shapes + assert destination.qkv_split_shapes_global == source.qkv_split_shapes_global + assert destination.qkv_split_groups_are_complete is True + assert destination.qkv_split_heads_are_complete is True def test_non_allreduce_param_uses_expert_tp_group_for_duplicate_filter(): diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index e3b9f666fb2..7e7ce593f74 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -16,6 +16,8 @@ TensorParallelAdaptiveMuon, TensorParallelMuon, _get_qkv_split_shapes, + _localize_qkv_split_shapes, + _qkv_split_groups_are_complete, get_supported_coefficient_types, validate_coefficient_type, ) @@ -80,6 +82,40 @@ def test_muon_qkv_split_shapes(): assert _get_qkv_split_shapes(config) == [128, 64, 64] assert _get_qkv_split_shapes(gated_config) == [128, 128, 64, 64] + assert _get_qkv_split_shapes(config, split_qkv_per_head=True) == [64] * 32 + assert _get_qkv_split_shapes(gated_config, split_qkv_per_head=True) == [64] * 48 + + +def test_muon_local_qkv_head_split_shapes_can_differ_by_tp_rank(): + """Rank-local per-head layouts report complete and fragmented heads.""" + global_split_shapes = [64] * 20 + + rank_0_shapes, rank_0_complete = _localize_qkv_split_shapes( + global_split_shapes, local_start=0, local_rows=160 + ) + rank_1_shapes, rank_1_complete = _localize_qkv_split_shapes( + global_split_shapes, local_start=160, local_rows=160 + ) + aligned_shapes, aligned_complete = _localize_qkv_split_shapes( + global_split_shapes, local_start=0, local_rows=640 + ) + + assert rank_0_shapes == [64, 64, 32] + assert rank_1_shapes == [32, 64, 64] + assert not rank_0_complete + assert not rank_1_complete + assert aligned_shapes == [64] * 10 + assert aligned_complete + + +def test_muon_qkv_query_group_layout_localization(): + """Projection splitting detects query groups fragmented by TP row ranges.""" + split_shapes = [256, 64, 64] + + assert _qkv_split_groups_are_complete(split_shapes, local_start=0, local_rows=384) + assert _qkv_split_groups_are_complete(split_shapes, local_start=384, local_rows=768) + assert not _qkv_split_groups_are_complete(split_shapes, local_start=0, local_rows=192) + assert not _qkv_split_groups_are_complete(split_shapes, local_start=192, local_rows=192) def test_muon_optimizer_smoke(): @@ -516,6 +552,95 @@ def test_muon_optimizer_blockwise_mode_different_result(self): model.weight.data, original_weight ), "Weight should be updated with mode=blockwise" + def test_muon_optimizer_per_head_split_gathers_fragmented_heads(self): + """Per-head splitting reconstructs heads that cross TP rank boundaries.""" + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + tp_group = pg_collection.tp + tp_rank = tp_group.rank() + local_grad = torch.arange(12, dtype=torch.float32, device='cuda').view(3, 4) + local_grad = local_grad + tp_rank * local_grad.numel() + param = torch.nn.Parameter(torch.zeros_like(local_grad)) + param.partition_dim = 0 + param.is_qkv = True + param.qkv_split_shapes = [2, 2] + param.qkv_split_shapes_global = [2, 2, 2] + param.qkv_split_heads_are_complete = False + + optimizer = TensorParallelMuon( + params=[param], + split_qkv=True, + split_qkv_per_head=True, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=[2, 2, 2], + pg_collection=pg_collection, + tp_mode="blockwise", + ) + + def center_rows(x, tp_group=None, partition_dim=None): + del tp_group, partition_dim + return x - x.mean(dim=-2, keepdim=True) + + optimizer.scaled_orthogonalize_fn = center_rows + actual = optimizer.orthogonalize(param, local_grad) + + shards = [torch.empty_like(local_grad) for _ in range(tp_group.size())] + torch.distributed.all_gather(shards, local_grad, tp_group) + global_grad = torch.cat(shards, dim=0) + expected_global = torch.cat( + [center_rows(head) for head in torch.split(global_grad, [2, 2, 2], dim=0)], dim=0 + ) + expected = expected_global[tp_rank * 3 : (tp_rank + 1) * 3] + torch.testing.assert_close(actual, expected) + + @pytest.mark.parametrize("split_shapes", ([4, 2, 2], [4, 4, 2, 2])) + def test_muon_optimizer_projection_split_gathers_fragmented_query_group(self, split_shapes): + """Projection splitting reconstructs a query group split over TP ranks.""" + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + tp_group = pg_collection.tp + tp_rank = tp_group.rank() + global_rows = sum(split_shapes) + assert global_rows % tp_group.size() == 0 + local_rows = global_rows // tp_group.size() + local_grad = torch.arange(local_rows * 4, dtype=torch.float32, device='cuda').view( + local_rows, 4 + ) + local_grad = local_grad + tp_rank * local_grad.numel() + param = torch.nn.Parameter(torch.zeros_like(local_grad)) + param.partition_dim = 0 + param.is_qkv = True + param.qkv_split_shapes = split_shapes + param.qkv_split_shapes_global = split_shapes + param.qkv_split_groups_are_complete = False + + optimizer = TensorParallelMuon( + params=[param], + split_qkv=True, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=split_shapes, + pg_collection=pg_collection, + tp_mode="blockwise", + ) + + def center_rows(x, tp_group=None, partition_dim=None): + del tp_group, partition_dim + return x - x.mean(dim=-2, keepdim=True) + + optimizer.scaled_orthogonalize_fn = center_rows + actual = optimizer.orthogonalize(param, local_grad) + + shards = [torch.empty_like(local_grad) for _ in range(tp_group.size())] + torch.distributed.all_gather(shards, local_grad, tp_group) + global_grad = torch.cat(shards, dim=0) + expected_global = torch.cat( + [ + center_rows(projection) + for projection in torch.split(global_grad, split_shapes, dim=0) + ], + dim=0, + ) + expected = expected_global[tp_rank * local_rows : (tp_rank + 1) * local_rows] + torch.testing.assert_close(actual, expected) + # All non-custom coefficient types supported by emerging_optimizers. _TESTABLE_COEFFICIENT_TYPES = ( @@ -720,6 +845,61 @@ def test_muon_optimizer_qkv_split(): ), "Weights should be different between split_qkv=True and split_qkv=False" +def test_muon_optimizer_qkv_split_per_head_is_opt_in(): + """Per-head splitting is guarded and differs from projection splitting.""" + grad = torch.arange(48, dtype=torch.float32, device='cuda').view(16, 3) + projection_param = torch.nn.Parameter(torch.zeros_like(grad)) + projection_param.is_qkv = True + projection_param.qkv_split_shapes = [4, 2, 2] + head_param = torch.nn.Parameter(torch.zeros_like(grad)) + head_param.is_qkv = True + head_param.qkv_split_shapes = [2] * 8 + orthogonalize_call_shapes = [] + + def center_rows(x, tp_group=None, partition_dim=None): + del tp_group, partition_dim + orthogonalize_call_shapes.append(tuple(x.shape)) + return x - x.mean(dim=-2, keepdim=True) + + projection_optimizer = TensorParallelMuon( + params=[projection_param], + split_qkv=True, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=[4, 2, 2], + pg_collection=None, + ) + projection_optimizer.scaled_orthogonalize_fn = center_rows + projection_out = projection_optimizer.orthogonalize(projection_param, grad) + orthogonalize_call_shapes.clear() + + head_optimizer = TensorParallelMuon( + params=[head_param], + split_qkv=True, + split_qkv_per_head=True, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=[2] * 8, + pg_collection=None, + ) + head_optimizer.scaled_orthogonalize_fn = center_rows + head_out = head_optimizer.orthogonalize(head_param, grad) + assert orthogonalize_call_shapes == [(8, 2, 3)] + + expected_head_out = torch.cat( + [center_rows(head) for head in torch.split(grad, [2] * 8, dim=0)], dim=0 + ) + torch.testing.assert_close(head_out, expected_head_out) + assert not torch.equal(projection_out, head_out) + + +def test_muon_optimizer_qkv_split_per_head_requires_split_qkv(): + """The per-head switch cannot enable QKV splitting by itself.""" + param = torch.nn.Parameter(torch.zeros(4, 4, dtype=torch.float32, device='cuda')) + with pytest.raises(ValueError, match="split_qkv_per_head requires split_qkv=True"): + TensorParallelMuon( + params=[param], split_qkv=False, split_qkv_per_head=True, pg_collection=None + ) + + def test_muon_optimizer_extra_scale_factor(): """Test TensorParallelMuon optimizer with different extra_scale_factor values.""" model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda') From 33bc3221ead788e7d8f7d68c24bff40364d7b1f4 Mon Sep 17 00:00:00 2001 From: mkhona Date: Fri, 7 Aug 2026 09:49:08 -0700 Subject: [PATCH 2/8] Handle GTP padding in Muon QKV splits Signed-off-by: mkhona --- megatron/core/optimizer/__init__.py | 41 +++++-- .../core/optimizer/emerging_optimizers.py | 36 ++++-- megatron/core/tensor_parallel/layers.py | 1 + .../test_tp_gtp.py | 104 ++++++++++++++++++ .../test_tp_attrs_without_init.py | 2 + 5 files changed, 166 insertions(+), 18 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 9faab7ddf07..395a8e970f6 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -822,26 +822,47 @@ def _get_megatron_emerging_optimizer( gtp_size = 1 gtp_rank = 0 - tp_local_rows = param.shape[0] * gtp_size - expected_global_rows = tp_local_rows * tp_size + qkv_gtp_pad_length = ( + int(getattr(param, 'pad_length', 0)) + if getattr(param, 'is_gtp_weight_remat', False) + else 0 + ) + physical_tp_local_rows = param.shape[0] * gtp_size + if not 0 <= qkv_gtp_pad_length < physical_tp_local_rows: + raise RuntimeError( + f"Invalid Muon QKV GTP padding for {name}: " + f"pad_length={qkv_gtp_pad_length}, " + f"physical_tp_local_rows={physical_tp_local_rows}" + ) + param.qkv_gtp_pad_length = qkv_gtp_pad_length + logical_tp_local_rows = physical_tp_local_rows - qkv_gtp_pad_length + expected_global_rows = logical_tp_local_rows * tp_size if expected_global_rows != sum(global_split_shapes): raise RuntimeError( f"Muon QKV layout mismatch for {name}: " f"global_rows={sum(global_split_shapes)}, " f"local_rows={param.shape[0]}, tp_size={tp_size}, " - f"gtp_remat_size={gtp_size}" + f"gtp_remat_size={gtp_size}, " + f"gtp_pad_length={qkv_gtp_pad_length}" ) - local_start = tp_rank * tp_local_rows + gtp_rank * param.shape[0] + local_start = tp_rank * logical_tp_local_rows + gtp_rank * param.shape[0] if config.muon_split_qkv_per_head: - param.qkv_split_shapes, param.qkv_split_heads_are_complete = ( - _localize_qkv_split_shapes( - qkv_split_shapes, local_start=local_start, local_rows=param.shape[0] + if qkv_gtp_pad_length > 0: + param.qkv_split_shapes = qkv_split_shapes + param.qkv_split_heads_are_complete = False + else: + param.qkv_split_shapes, param.qkv_split_heads_are_complete = ( + _localize_qkv_split_shapes( + qkv_split_shapes, local_start=local_start, local_rows=param.shape[0] + ) ) - ) else: param.qkv_split_shapes = qkv_split_shapes - param.qkv_split_groups_are_complete = _qkv_split_groups_are_complete( - qkv_split_shapes, local_start=local_start, local_rows=param.shape[0] + param.qkv_split_groups_are_complete = ( + qkv_gtp_pad_length == 0 + and _qkv_split_groups_are_complete( + qkv_split_shapes, local_start=local_start, local_rows=param.shape[0] + ) ) # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam). diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index d315ad1e376..3c319acb034 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -301,6 +301,9 @@ def _gather_qkv_grad(self, p, grad, tp_group, expected_rows, gather_gtp=True): gathered_grad = grad gtp_slice = None gtp_remat_group = self._get_gtp_remat_group(p) + gtp_pad_length = int(getattr(p, "qkv_gtp_pad_length", 0)) + if gtp_pad_length < 0: + raise RuntimeError(f"Muon QKV GTP padding must be non-negative: {gtp_pad_length}") if ( gather_gtp and gtp_remat_group is not None @@ -313,7 +316,18 @@ def _gather_qkv_grad(self, p, grad, tp_group, expected_rows, gather_gtp=True): shards = [torch.empty_like(gathered_grad) for _ in range(gtp_size)] torch.distributed.all_gather(shards, gathered_grad, gtp_remat_group) gathered_grad = torch.cat(shards, dim=0) - gtp_slice = (gtp_rank, gtp_local_rows) + if gtp_pad_length >= gathered_grad.shape[0]: + raise RuntimeError( + "Invalid Muon QKV GTP padding after gathering: " + f"pad_length={gtp_pad_length}, gathered_rows={gathered_grad.shape[0]}" + ) + if gtp_pad_length > 0: + gathered_grad = gathered_grad[:-gtp_pad_length] + gtp_slice = (gtp_rank, gtp_local_rows, gtp_pad_length) + elif gtp_pad_length > 0: + raise RuntimeError( + "Muon QKV has GTP padding but its GTP-remat shards were not gathered" + ) tp_slice = None if gathered_grad.shape[0] != expected_rows: @@ -351,7 +365,9 @@ def _restore_local_qkv_grad(gathered_grad, tp_slice, gtp_slice): tp_rank, tp_local_rows = tp_slice gathered_grad = gathered_grad[tp_rank * tp_local_rows : (tp_rank + 1) * tp_local_rows] if gtp_slice is not None: - gtp_rank, gtp_local_rows = gtp_slice + gtp_rank, gtp_local_rows, gtp_pad_length = gtp_slice + if gtp_pad_length > 0: + gathered_grad = torch.nn.functional.pad(gathered_grad, (0, 0, 0, gtp_pad_length)) gathered_grad = gathered_grad[ gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows ] @@ -435,10 +451,14 @@ def _orthogonalize_qkv_per_head(self, p, grad, tp_group): """ local_split_shapes = getattr(p, "qkv_split_shapes", None) heads_are_complete = getattr(p, "qkv_split_heads_are_complete", None) - use_local_layout = heads_are_complete is True or ( - heads_are_complete is None - and local_split_shapes is not None - and sum(local_split_shapes) == grad.shape[0] + has_gtp_padding = int(getattr(p, "qkv_gtp_pad_length", 0)) > 0 + use_local_layout = not has_gtp_padding and ( + heads_are_complete is True + or ( + heads_are_complete is None + and local_split_shapes is not None + and sum(local_split_shapes) == grad.shape[0] + ) ) if use_local_layout: qkv_split_shapes = local_split_shapes @@ -526,8 +546,8 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t raise RuntimeError("Muon QKV split requested but qkv_split_shapes is not set") if ( getattr(p, "qkv_split_groups_are_complete", None) is False - and getattr(p, "qkv_split_shapes_global", None) is not None - ): + or int(getattr(p, "qkv_gtp_pad_length", 0)) > 0 + ) and getattr(p, "qkv_split_shapes_global", None) is not None: return self._orthogonalize_fragmented_qkv(p, grad, tp_group, qkv_split_shapes) log_single_rank( logger, diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 598e0f12951..85ea5503b11 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -65,6 +65,7 @@ "is_qkv": False, "qkv_split_shapes": None, "qkv_split_shapes_global": None, + "qkv_gtp_pad_length": 0, "qkv_split_groups_are_complete": False, "qkv_split_heads_are_complete": False, "tensor_model_parallel": False, diff --git a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py index 961ab070556..0399f153c0b 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py @@ -15,6 +15,8 @@ 2. TestTPGTPColumnParallelLinear - column-parallel Linear: fwd/bwd correctness (weight shape verified inline) 3. TestTPGTPRowParallelLinear - row-parallel Linear: fwd/bwd smoke test + numerical correctness 4. TestTPGTPLayerNormLinear - LayerNormLinear column-parallel smoke test +5. TestTPGTPPaddingAlignment - alignment padding is applied independently per TP slice +6. TestTPGTPMuonQKVPadding - Muon excludes GTP padding from QKV orthogonalization Tests use (tp_size, gtp_remat_size) = (2, 2) → world_size = 4 (runs on 4-GPU machines). @@ -27,7 +29,10 @@ import pytest import torch import torch.distributed as dist +import torch.nn.functional as F +from megatron.core.optimizer.emerging_optimizers import HAVE_EMERGING_OPTIMIZERS, TensorParallelMuon +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.gtp_api import HAVE_GTP if not HAVE_GTP: @@ -395,3 +400,102 @@ def test_pre_init_pads_per_tp_slice(self, tp_size, gtp_remat_size): world_size = tp_size * gtp_remat_size _requires_multi_gpu(world_size) _run_distributed(_worker_pre_init_tp_padding, world_size, tp_size, gtp_remat_size) + + +def _worker_muon_qkv_padding( + rank, world_size, port, tp_size, gtp_remat_size, split_per_head, split_shapes +): + """Muon must exclude per-TP GTP padding from QKV orthogonalization.""" + del port + tp_group, gtp_remat_group, tp_rank, gtp_rank = _build_groups( + rank, world_size, tp_size, gtp_remat_size + ) + + logical_rows = sum(split_shapes) + logical_tp_rows = logical_rows // tp_size + pad_length = 3 + physical_tp_rows = logical_tp_rows + pad_length + gtp_local_rows = physical_tp_rows // gtp_remat_size + hidden_size = 4 + + global_grad = torch.arange(logical_rows * hidden_size, dtype=torch.float32, device="cuda").view( + logical_rows, hidden_size + ) + tp_grad = global_grad[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows] + # Use a sentinel so the test detects padding entering the orthogonalization. + physical_tp_grad = F.pad(tp_grad, (0, 0, 0, pad_length), value=10_000.0) + local_grad = physical_tp_grad[ + gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows + ].clone() + + param = torch.nn.Parameter(torch.zeros_like(local_grad)) + param.partition_dim = 0 + param.is_qkv = True + param.is_gtp_weight_remat = True + param.qkv_gtp_pad_length = pad_length + param.qkv_split_shapes = split_shapes + param.qkv_split_shapes_global = split_shapes + param.qkv_split_heads_are_complete = False + param.qkv_split_groups_are_complete = False + + optimizer = TensorParallelMuon( + params=[param], + split_qkv=True, + split_qkv_per_head=split_per_head, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=split_shapes, + pg_collection=ProcessGroupCollection(tp=tp_group, gtp_remat=gtp_remat_group), + tp_mode="blockwise", + ) + + def center_rows(x, tp_group=None, partition_dim=None): + del tp_group, partition_dim + return x - x.mean(dim=-2, keepdim=True) + + optimizer.scaled_orthogonalize_fn = center_rows + actual = optimizer.orthogonalize(param, local_grad) + + if split_per_head: + expected_global = torch.cat( + [center_rows(head) for head in torch.split(global_grad, split_shapes, dim=0)], dim=0 + ) + else: + grouped_grad = global_grad.view(1, logical_rows, hidden_size) + projections = torch.split(grouped_grad, split_shapes, dim=1) + expected_global = torch.cat( + [ + center_rows(projection.reshape(-1, hidden_size)).view_as(projection) + for projection in projections + ], + dim=1, + ).view_as(global_grad) + + expected_tp = expected_global[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows] + expected_tp = F.pad(expected_tp, (0, 0, 0, pad_length)) + expected = expected_tp[gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows] + + torch.testing.assert_close(actual, expected) + assert actual.shape == local_grad.shape + if gtp_rank == gtp_remat_size - 1: + assert torch.count_nonzero(actual[-pad_length:]) == 0 + + +class TestTPGTPMuonQKVPadding: + @pytest.mark.parametrize( + "split_per_head,split_shapes", [(True, [2, 2, 2, 2, 2]), (False, [6, 2, 2])] + ) + def test_padding_is_excluded_from_qkv_orthogonalization(self, split_per_head, split_shapes): + if not HAVE_EMERGING_OPTIMIZERS: + pytest.skip("emerging_optimizers package is not installed") + tp_size = 2 + gtp_remat_size = 2 + world_size = tp_size * gtp_remat_size + _requires_multi_gpu(world_size) + _run_distributed( + _worker_muon_qkv_padding, + world_size, + tp_size, + gtp_remat_size, + split_per_head, + split_shapes, + ) diff --git a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py index 983de4dcd6e..17ac0bbbe99 100644 --- a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py +++ b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py @@ -97,6 +97,7 @@ def test_copy_tensor_model_parallel_attributes_preserves_qkv_split_shapes(): source.is_qkv = True source.qkv_split_shapes = [2, 2] source.qkv_split_shapes_global = [2] * 4 + source.qkv_gtp_pad_length = 3 source.qkv_split_groups_are_complete = True source.qkv_split_heads_are_complete = True @@ -105,6 +106,7 @@ def test_copy_tensor_model_parallel_attributes_preserves_qkv_split_shapes(): assert destination.is_qkv is True assert destination.qkv_split_shapes == source.qkv_split_shapes assert destination.qkv_split_shapes_global == source.qkv_split_shapes_global + assert destination.qkv_gtp_pad_length == 3 assert destination.qkv_split_groups_are_complete is True assert destination.qkv_split_heads_are_complete is True From eeeb515257b6779cc14410be10c437a382d98426 Mon Sep 17 00:00:00 2001 From: mkhona Date: Mon, 17 Aug 2026 15:51:31 -0700 Subject: [PATCH 3/8] Clarify padded per-head QKV metadata Signed-off-by: mkhona --- megatron/core/optimizer/__init__.py | 4 +++- tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 395a8e970f6..657c1688158 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -848,7 +848,9 @@ def _get_megatron_emerging_optimizer( local_start = tp_rank * logical_tp_local_rows + gtp_rank * param.shape[0] if config.muon_split_qkv_per_head: if qkv_gtp_pad_length > 0: - param.qkv_split_shapes = qkv_split_shapes + # A padded GTP shard does not have a purely logical local layout. + # Force the per-head path to use qkv_split_shapes_global instead. + param.qkv_split_shapes = None param.qkv_split_heads_are_complete = False else: param.qkv_split_shapes, param.qkv_split_heads_are_complete = ( diff --git a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py index 0399f153c0b..1b5ddf58388 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py @@ -433,7 +433,8 @@ def _worker_muon_qkv_padding( param.is_qkv = True param.is_gtp_weight_remat = True param.qkv_gtp_pad_length = pad_length - param.qkv_split_shapes = split_shapes + # Padded per-head shards intentionally have no rank-local logical split metadata. + param.qkv_split_shapes = None if split_per_head else split_shapes param.qkv_split_shapes_global = split_shapes param.qkv_split_heads_are_complete = False param.qkv_split_groups_are_complete = False From 8356b9cdadef215dc60ffd3f44525acd60b637d6 Mon Sep 17 00:00:00 2001 From: mkhona Date: Mon, 17 Aug 2026 16:47:21 -0700 Subject: [PATCH 4/8] Guard batched Muon NS by optimizer version Signed-off-by: mkhona --- .../core/optimizer/emerging_optimizers.py | 23 ++++++ .../test_tp_gtp.py | 10 ++- tests/unit_tests/test_emerging_optimizers.py | 78 ++++++++++++++++++- 3 files changed, 108 insertions(+), 3 deletions(-) diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 3c319acb034..2a736003c04 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -11,9 +11,12 @@ import inspect import logging from dataclasses import dataclass, field +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version from typing import Any, Callable, Dict, Literal, Optional, get_args import torch +from packaging.version import Version from torch.optim.optimizer import ParamsT from megatron.core.process_groups_config import ProcessGroupCollection @@ -43,6 +46,23 @@ logger = logging.getLogger(__name__) +try: + EMERGING_OPTIMIZERS_VERSION = Version(package_version("emerging-optimizers")) +except PackageNotFoundError: + EMERGING_OPTIMIZERS_VERSION = Version("0") + +_BATCHED_NEWTON_SCHULZ_MIN_VERSION = Version("0.3.0") + + +def _require_batched_newton_schulz_support() -> None: + """Ensure the installed Emerging-Optimizers supports batched Newton-Schulz.""" + if EMERGING_OPTIMIZERS_VERSION < _BATCHED_NEWTON_SCHULZ_MIN_VERSION: + raise RuntimeError( + "Batched Newton-Schulz requires emerging-optimizers>=0.3.0; " + f"found {EMERGING_OPTIMIZERS_VERSION}. Disable muon_split_qkv_per_head or " + "upgrade emerging-optimizers." + ) + def get_supported_coefficient_types() -> tuple[str, ...]: """Return the coefficient types supported by the installed emerging_optimizers. @@ -237,6 +257,8 @@ def __init__( raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") if split_qkv_per_head and not split_qkv: raise ValueError("split_qkv_per_head requires split_qkv=True") + if split_qkv_per_head and qkv_split_shapes and len(set(qkv_split_shapes)) == 1: + _require_batched_newton_schulz_support() def scaled_orthogonalize_fn( grad: torch.Tensor, @@ -388,6 +410,7 @@ def _orthogonalize_split_qkv(self, grad, split_shapes, orthogonalize_fn): ) if len(set(split_shapes)) == 1: # A 3D input selects Emerging-Optimizers' batched Newton-Schulz path. + _require_batched_newton_schulz_support() head_rows = split_shapes[0] return orthogonalize_fn(grad.view(len(split_shapes), head_rows, -1)).view_as(grad) return torch.cat( diff --git a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py index 1b5ddf58388..24b24c48eab 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py @@ -30,6 +30,7 @@ import torch import torch.distributed as dist import torch.nn.functional as F +from packaging.version import Version from megatron.core.optimizer.emerging_optimizers import HAVE_EMERGING_OPTIMIZERS, TensorParallelMuon from megatron.core.process_groups_config import ProcessGroupCollection @@ -485,9 +486,16 @@ class TestTPGTPMuonQKVPadding: @pytest.mark.parametrize( "split_per_head,split_shapes", [(True, [2, 2, 2, 2, 2]), (False, [6, 2, 2])] ) - def test_padding_is_excluded_from_qkv_orthogonalization(self, split_per_head, split_shapes): + def test_padding_is_excluded_from_qkv_orthogonalization( + self, split_per_head, split_shapes, monkeypatch + ): if not HAVE_EMERGING_OPTIMIZERS: pytest.skip("emerging_optimizers package is not installed") + if split_per_head: + monkeypatch.setattr( + "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", + Version("0.3.0"), + ) tp_size = 2 gtp_remat_size = 2 world_size = tp_size * gtp_remat_size diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index 7e7ce593f74..ac8caf925bd 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -552,8 +552,12 @@ def test_muon_optimizer_blockwise_mode_different_result(self): model.weight.data, original_weight ), "Weight should be updated with mode=blockwise" - def test_muon_optimizer_per_head_split_gathers_fragmented_heads(self): + def test_muon_optimizer_per_head_split_gathers_fragmented_heads(self, monkeypatch): """Per-head splitting reconstructs heads that cross TP rank boundaries.""" + monkeypatch.setattr( + "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", + Version("0.3.0"), + ) pg_collection = ProcessGroupCollection.use_mpu_process_groups() tp_group = pg_collection.tp tp_rank = tp_group.rank() @@ -845,8 +849,11 @@ def test_muon_optimizer_qkv_split(): ), "Weights should be different between split_qkv=True and split_qkv=False" -def test_muon_optimizer_qkv_split_per_head_is_opt_in(): +def test_muon_optimizer_qkv_split_per_head_is_opt_in(monkeypatch): """Per-head splitting is guarded and differs from projection splitting.""" + monkeypatch.setattr( + "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.3.0") + ) grad = torch.arange(48, dtype=torch.float32, device='cuda').view(16, 3) projection_param = torch.nn.Parameter(torch.zeros_like(grad)) projection_param.is_qkv = True @@ -900,6 +907,73 @@ def test_muon_optimizer_qkv_split_per_head_requires_split_qkv(): ) +def test_muon_optimizer_batched_ns_requires_emerging_optimizers_0_3(monkeypatch): + """Uniform per-head splits require Emerging-Optimizers batched NS support.""" + monkeypatch.setattr( + "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.2.0") + ) + param = torch.nn.Parameter(torch.zeros(4, 4, dtype=torch.float32, device='cuda')) + + with pytest.raises(RuntimeError, match="Batched Newton-Schulz requires.*>=0.3.0"): + TensorParallelMuon( + params=[param], + split_qkv=True, + split_qkv_per_head=True, + qkv_split_shapes=[2, 2], + pg_collection=None, + ) + + +def test_muon_optimizer_runtime_batched_ns_version_guard(monkeypatch): + """Per-parameter uniform splits are guarded when constructor shapes are unavailable.""" + monkeypatch.setattr( + "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.2.0") + ) + grad = torch.zeros(4, 4, dtype=torch.float32, device='cuda') + param = torch.nn.Parameter(torch.zeros_like(grad)) + param.is_qkv = True + param.qkv_split_shapes = [2, 2] + optimizer = TensorParallelMuon( + params=[param], + split_qkv=True, + split_qkv_per_head=True, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=None, + pg_collection=None, + ) + + with pytest.raises(RuntimeError, match="Batched Newton-Schulz requires.*>=0.3.0"): + optimizer.orthogonalize(param, grad) + + +def test_muon_optimizer_nonuniform_per_head_splits_do_not_require_batched_ns(monkeypatch): + """Nonuniform per-head splits keep using the unbatched compatibility path.""" + monkeypatch.setattr( + "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.2.0") + ) + grad = torch.arange(12, dtype=torch.float32, device='cuda').view(3, 4) + param = torch.nn.Parameter(torch.zeros_like(grad)) + param.is_qkv = True + param.qkv_split_shapes = [2, 1] + optimizer = TensorParallelMuon( + params=[param], + split_qkv=True, + split_qkv_per_head=True, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=[2, 1], + pg_collection=None, + ) + + def center_rows(x, tp_group=None, partition_dim=None): + del tp_group, partition_dim + return x - x.mean(dim=-2, keepdim=True) + + optimizer.scaled_orthogonalize_fn = center_rows + actual = optimizer.orthogonalize(param, grad) + expected = torch.cat([center_rows(head) for head in torch.split(grad, [2, 1])]) + torch.testing.assert_close(actual, expected) + + def test_muon_optimizer_extra_scale_factor(): """Test TensorParallelMuon optimizer with different extra_scale_factor values.""" model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda') From 0f327f5d7d09c268f8956eee3b79d4272c39cc8b Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 18 Aug 2026 16:51:19 -0700 Subject: [PATCH 5/8] Address Muon QKV review feedback Signed-off-by: mkhona --- megatron/core/optimizer/__init__.py | 22 ++- .../core/optimizer/emerging_optimizers.py | 16 +-- megatron/core/optimizer/optimizer_config.py | 5 +- megatron/training/arguments.py | 4 +- .../test_tp_gtp.py | 10 +- tests/unit_tests/test_emerging_optimizers.py | 125 ++++++++++++++---- 6 files changed, 129 insertions(+), 53 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 657c1688158..c41c32952a7 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -795,13 +795,11 @@ def _get_megatron_emerging_optimizer( qkv_split_shapes = _get_qkv_split_shapes( model_chunk.config, split_qkv_per_head=config.muon_split_qkv_per_head ) - param.is_qkv = True global_split_shapes = ( qkv_split_shapes if config.muon_split_qkv_per_head else qkv_split_shapes * model_chunk.config.num_query_groups ) - param.qkv_split_shapes_global = global_split_shapes tp_group = ( pg_collection.expt_tp @@ -834,17 +832,29 @@ def _get_megatron_emerging_optimizer( f"pad_length={qkv_gtp_pad_length}, " f"physical_tp_local_rows={physical_tp_local_rows}" ) - param.qkv_gtp_pad_length = qkv_gtp_pad_length logical_tp_local_rows = physical_tp_local_rows - qkv_gtp_pad_length expected_global_rows = logical_tp_local_rows * tp_size if expected_global_rows != sum(global_split_shapes): - raise RuntimeError( - f"Muon QKV layout mismatch for {name}: " + log_single_rank( + logger, + logging.DEBUG, + f"Emerging optimizer QKV split skipped for {name}: " f"global_rows={sum(global_split_shapes)}, " f"local_rows={param.shape[0]}, tp_size={tp_size}, " f"gtp_remat_size={gtp_size}, " - f"gtp_pad_length={qkv_gtp_pad_length}" + f"gtp_pad_length={qkv_gtp_pad_length}", ) + param.is_qkv = False + param.qkv_split_shapes = None + param.qkv_split_shapes_global = None + param.qkv_gtp_pad_length = 0 + param.qkv_split_groups_are_complete = False + param.qkv_split_heads_are_complete = False + continue + + param.is_qkv = True + param.qkv_split_shapes_global = global_split_shapes + param.qkv_gtp_pad_length = qkv_gtp_pad_length local_start = tp_rank * logical_tp_local_rows + gtp_rank * param.shape[0] if config.muon_split_qkv_per_head: if qkv_gtp_pad_length > 0: diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 2a736003c04..248059a058f 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -54,14 +54,9 @@ _BATCHED_NEWTON_SCHULZ_MIN_VERSION = Version("0.3.0") -def _require_batched_newton_schulz_support() -> None: - """Ensure the installed Emerging-Optimizers supports batched Newton-Schulz.""" - if EMERGING_OPTIMIZERS_VERSION < _BATCHED_NEWTON_SCHULZ_MIN_VERSION: - raise RuntimeError( - "Batched Newton-Schulz requires emerging-optimizers>=0.3.0; " - f"found {EMERGING_OPTIMIZERS_VERSION}. Disable muon_split_qkv_per_head or " - "upgrade emerging-optimizers." - ) +def _supports_batched_newton_schulz() -> bool: + """Return whether Emerging-Optimizers supports batched Newton-Schulz.""" + return EMERGING_OPTIMIZERS_VERSION >= _BATCHED_NEWTON_SCHULZ_MIN_VERSION def get_supported_coefficient_types() -> tuple[str, ...]: @@ -257,8 +252,6 @@ def __init__( raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") if split_qkv_per_head and not split_qkv: raise ValueError("split_qkv_per_head requires split_qkv=True") - if split_qkv_per_head and qkv_split_shapes and len(set(qkv_split_shapes)) == 1: - _require_batched_newton_schulz_support() def scaled_orthogonalize_fn( grad: torch.Tensor, @@ -408,9 +401,8 @@ def _orthogonalize_split_qkv(self, grad, split_shapes, orthogonalize_fn): f"Muon per-head QKV split shape mismatch: grad_shape={tuple(grad.shape)}, " f"split_shapes={split_shapes}" ) - if len(set(split_shapes)) == 1: + if len(set(split_shapes)) == 1 and _supports_batched_newton_schulz(): # A 3D input selects Emerging-Optimizers' batched Newton-Schulz path. - _require_batched_newton_schulz_support() head_rows = split_shapes[0] return orthogonalize_fn(grad.view(len(split_shapes), head_rows, -1)).view_as(grad) return torch.cat( diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 81a3eec0ac7..44bcdd74e85 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -265,8 +265,9 @@ class OptimizerConfig: muon_split_qkv_per_head: bool = False """Whether to orthogonalize each Q, gate, K, and V head independently. Requires - ``muon_split_qkv``. By default, Muon orthogonalizes the Q, gate, K, and V - projection matrices separately.""" + ``muon_split_qkv``. Batched execution requires Emerging-Optimizers 0.3.0 or newer; + older versions process heads individually. By default, Muon orthogonalizes the Q, + gate, K, and V projection matrices separately.""" muon_nesterov: bool = False """Whether to use Nesterov-style momentum in the internal SGD.""" diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 937417edaad..31d361cc022 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2519,7 +2519,9 @@ def _add_regularization_args(parser): help='Whether to split QKV parameters for Muon optimizer') group.add_argument('--muon-split-qkv-per-head', action='store_true', help='Orthogonalize each Q, gate, K, and V head independently. ' - 'By default, Q, gate, K, and V projections are orthogonalized separately') + 'Batched execution requires emerging-optimizers>=0.3.0; older versions ' + 'process heads individually. By default, Q, gate, K, and V projections ' + 'are orthogonalized separately') group.add_argument('--muon-nesterov', action='store_true', help='Whether to use Nesterov-style momentum in the internal SGD') group.add_argument('--muon-scale-mode', type=str, default='spectral', diff --git a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py index 24b24c48eab..1b5ddf58388 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py @@ -30,7 +30,6 @@ import torch import torch.distributed as dist import torch.nn.functional as F -from packaging.version import Version from megatron.core.optimizer.emerging_optimizers import HAVE_EMERGING_OPTIMIZERS, TensorParallelMuon from megatron.core.process_groups_config import ProcessGroupCollection @@ -486,16 +485,9 @@ class TestTPGTPMuonQKVPadding: @pytest.mark.parametrize( "split_per_head,split_shapes", [(True, [2, 2, 2, 2, 2]), (False, [6, 2, 2])] ) - def test_padding_is_excluded_from_qkv_orthogonalization( - self, split_per_head, split_shapes, monkeypatch - ): + def test_padding_is_excluded_from_qkv_orthogonalization(self, split_per_head, split_shapes): if not HAVE_EMERGING_OPTIMIZERS: pytest.skip("emerging_optimizers package is not installed") - if split_per_head: - monkeypatch.setattr( - "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", - Version("0.3.0"), - ) tp_size = 2 gtp_remat_size = 2 world_size = tp_size * gtp_remat_size diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index ac8caf925bd..5e3c9038497 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -10,6 +10,8 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec +from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer from megatron.core.optimizer.emerging_optimizers import ( HAVE_EMERGING_OPTIMIZERS, @@ -23,6 +25,7 @@ ) from megatron.core.optimizer.muon import get_megatron_muon_optimizer from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -511,6 +514,88 @@ def create_tp_model_and_optimizer(self, mode): return model, optimizer + def test_optimizer_factory_tags_qkv_when_tp_exceeds_query_groups(self): + """The public optimizer factory tags query groups fragmented across TP ranks.""" + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + tp_size = pg_collection.tp.size() + assert tp_size == 2 + model_parallel_cuda_manual_seed(123) + transformer_config = TransformerConfig( + num_layers=1, + hidden_size=8, + num_attention_heads=2, + num_query_groups=1, + kv_channels=4, + tensor_model_parallel_size=tp_size, + use_cpu_initialization=False, + add_bias_linear=False, + ) + model = GPTModel( + config=transformer_config, + transformer_layer_spec=get_gpt_layer_local_spec(), + vocab_size=32, + max_sequence_length=8, + pre_process=False, + post_process=False, + pg_collection=pg_collection, + ) + optimizer_config = OptimizerConfig( + optimizer='muon', + lr=0.01, + use_distributed_optimizer=False, + muon_split_qkv=True, + muon_tp_mode="blockwise", + ) + + optimizer = get_megatron_optimizer( + config=optimizer_config, + model_chunks=[model], + use_gloo_process_groups=False, + pg_collection=pg_collection, + ) + + qkv_weight = model.decoder.layers[0].self_attention.linear_qkv.weight + assert optimizer is not None + assert qkv_weight.shape[0] == 8 + assert qkv_weight.is_qkv + assert qkv_weight.qkv_split_shapes == [8, 4, 4] + assert qkv_weight.qkv_split_shapes_global == [8, 4, 4] + assert not qkv_weight.qkv_split_groups_are_complete + + def test_optimizer_factory_skips_mismatched_qkv_layout(self): + """A QKV layout mismatch falls back to whole-matrix Muon orthogonalization.""" + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + tp_size = pg_collection.tp.size() + transformer_config = TransformerConfig( + num_layers=1, + hidden_size=8, + num_attention_heads=2, + num_query_groups=1, + kv_channels=4, + tensor_model_parallel_size=tp_size, + ) + model = torch.nn.Module() + model.config = transformer_config + model.linear_qkv = torch.nn.Linear(8, 7, bias=False, dtype=torch.float32, device='cuda') + model.linear_qkv.weight.tensor_model_parallel = True + model.linear_qkv.weight.partition_dim = 0 + optimizer_config = OptimizerConfig( + optimizer='muon', lr=0.01, use_distributed_optimizer=False, muon_split_qkv=True + ) + + optimizer = get_megatron_optimizer( + config=optimizer_config, + model_chunks=[model], + use_gloo_process_groups=False, + pg_collection=pg_collection, + ) + + qkv_weight = model.linear_qkv.weight + assert optimizer is not None + assert not qkv_weight.is_qkv + assert qkv_weight.qkv_split_shapes is None + assert qkv_weight.qkv_split_shapes_global is None + @pytest.mark.parametrize("mode", ["duplicated", "distributed"]) def test_muon_optimizer_modes_multirank_same_result(self, mode): """Test that duplicated and distributed modes produce same results with TP > 1.""" @@ -907,29 +992,12 @@ def test_muon_optimizer_qkv_split_per_head_requires_split_qkv(): ) -def test_muon_optimizer_batched_ns_requires_emerging_optimizers_0_3(monkeypatch): - """Uniform per-head splits require Emerging-Optimizers batched NS support.""" +def test_muon_optimizer_uniform_per_head_splits_fall_back_without_batched_ns(monkeypatch): + """Uniform per-head splits use individual 2D calls before Emerging-Optimizers 0.3.""" monkeypatch.setattr( "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.2.0") ) - param = torch.nn.Parameter(torch.zeros(4, 4, dtype=torch.float32, device='cuda')) - - with pytest.raises(RuntimeError, match="Batched Newton-Schulz requires.*>=0.3.0"): - TensorParallelMuon( - params=[param], - split_qkv=True, - split_qkv_per_head=True, - qkv_split_shapes=[2, 2], - pg_collection=None, - ) - - -def test_muon_optimizer_runtime_batched_ns_version_guard(monkeypatch): - """Per-parameter uniform splits are guarded when constructor shapes are unavailable.""" - monkeypatch.setattr( - "megatron.core.optimizer.emerging_optimizers.EMERGING_OPTIMIZERS_VERSION", Version("0.2.0") - ) - grad = torch.zeros(4, 4, dtype=torch.float32, device='cuda') + grad = torch.arange(16, dtype=torch.float32, device='cuda').view(4, 4) param = torch.nn.Parameter(torch.zeros_like(grad)) param.is_qkv = True param.qkv_split_shapes = [2, 2] @@ -938,12 +1006,23 @@ def test_muon_optimizer_runtime_batched_ns_version_guard(monkeypatch): split_qkv=True, split_qkv_per_head=True, is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), - qkv_split_shapes=None, + qkv_split_shapes=[2, 2], pg_collection=None, ) + call_shapes = [] - with pytest.raises(RuntimeError, match="Batched Newton-Schulz requires.*>=0.3.0"): - optimizer.orthogonalize(param, grad) + def center_rows(x, tp_group=None, partition_dim=None): + del tp_group, partition_dim + call_shapes.append(tuple(x.shape)) + return x - x.mean(dim=-2, keepdim=True) + + optimizer.scaled_orthogonalize_fn = center_rows + actual = optimizer.orthogonalize(param, grad) + assert call_shapes == [(2, 4), (2, 4)] + expected = torch.cat( + [head - head.mean(dim=-2, keepdim=True) for head in torch.split(grad, [2, 2])] + ) + torch.testing.assert_close(actual, expected) def test_muon_optimizer_nonuniform_per_head_splits_do_not_require_batched_ns(monkeypatch): From 3fc9694cbd61158db8e3e63679857979199db0f5 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 18 Aug 2026 17:17:49 -0700 Subject: [PATCH 6/8] Use layer QKV layouts for Muon Signed-off-by: mkhona --- megatron/core/optimizer/__init__.py | 15 +-- .../core/optimizer/emerging_optimizers.py | 6 +- megatron/core/tensor_parallel/layers.py | 1 + megatron/core/transformer/attention.py | 16 +++ .../test_tp_attrs_without_init.py | 10 +- tests/unit_tests/test_emerging_optimizers.py | 98 +++++++++++++++++++ 6 files changed, 137 insertions(+), 9 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index c41c32952a7..329f92acdf1 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -783,7 +783,6 @@ 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 @@ -791,14 +790,18 @@ def _get_megatron_emerging_optimizer( param.expert_tp = True # TODO(deyuf): support MLA 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, split_qkv_per_head=config.muon_split_qkv_per_head - ) + qkv_layout = getattr(param, 'qkv_layout', None) + if qkv_layout is None: + # Backward compatibility for custom QKV modules that do not annotate + # their weight with the owning attention layer's logical layout. + qkv_layout = model_chunk.config + qkv_split_shapes = _get_qkv_split_shapes( + qkv_layout, split_qkv_per_head=config.muon_split_qkv_per_head + ) global_split_shapes = ( qkv_split_shapes if config.muon_split_qkv_per_head - else qkv_split_shapes * model_chunk.config.num_query_groups + else qkv_split_shapes * qkv_layout.num_query_groups ) tp_group = ( diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 248059a058f..ae4ebec030a 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -146,10 +146,12 @@ def _is_nonlinear_or_embedding(param): def _get_qkv_split_shapes(model_cfg, split_qkv_per_head: bool = False) -> list[int]: - """Compute fused QKV split shapes from a transformer model config. + """Compute fused QKV split shapes from logical attention layout metadata. Args: - model_cfg: Transformer model configuration. + model_cfg: Object exposing ``num_attention_heads``, ``num_query_groups``, + ``kv_channels``, and ``attention_output_gate``. This can be a transformer + config or the owning attention layer's parameter metadata. split_qkv_per_head: Return one split size per physical attention head. When false, return the per-query-group Q, gate (if present), K, and V projection widths. """ diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 85ea5503b11..8e5e76cbc11 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -63,6 +63,7 @@ _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = { "expert_tp": False, "is_qkv": False, + "qkv_layout": None, "qkv_split_shapes": None, "qkv_split_shapes_global": None, "qkv_gtp_pad_length": 0, diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 682b75fb701..1b82028c9fa 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -136,6 +136,16 @@ HAVE_FUSED_QKV_ROPE = False +@dataclass(frozen=True) +class QKVLayout: + """Logical layout metadata for a fused QKV projection weight.""" + + num_attention_heads: int + num_query_groups: int + kv_channels: int + attention_output_gate: bool + + class LinearQkvInterface(Protocol): """Interface for linear_qkv modules.""" @@ -1693,6 +1703,12 @@ def __init__( pg_collection=self.pg_collection, name=(name + ".linear_qkv") if name is not None else None, ) + self.linear_qkv.weight.qkv_layout = QKVLayout( + num_attention_heads=self.config.num_attention_heads, + num_query_groups=self.config.num_query_groups, + kv_channels=self.config.kv_channels, + attention_output_gate=self.config.attention_output_gate, + ) # Resolve which norm class to use for Q and K. # Config selects the default norm class; spec overrides if set. diff --git a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py index 17ac0bbbe99..983c11c202c 100644 --- a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py +++ b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py @@ -10,6 +10,7 @@ copy_tensor_model_parallel_attributes, param_is_not_tensor_parallel_duplicate, ) +from megatron.core.transformer.attention import QKVLayout from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -91,10 +92,16 @@ def test_row_parallel_linear_tp_attrs_no_init(self, use_cpu_init): assert hasattr(w, "partition_stride") and w.partition_stride == 1 -def test_copy_tensor_model_parallel_attributes_preserves_qkv_split_shapes(): +def test_copy_tensor_model_parallel_attributes_preserves_qkv_metadata(): source = torch.empty(4, 4) destination = torch.empty_like(source) source.is_qkv = True + source.qkv_layout = QKVLayout( + num_attention_heads=8, + num_query_groups=2, + kv_channels=64, + attention_output_gate=False, + ) source.qkv_split_shapes = [2, 2] source.qkv_split_shapes_global = [2] * 4 source.qkv_gtp_pad_length = 3 @@ -104,6 +111,7 @@ def test_copy_tensor_model_parallel_attributes_preserves_qkv_split_shapes(): copy_tensor_model_parallel_attributes(destination, source) assert destination.is_qkv is True + assert destination.qkv_layout == source.qkv_layout assert destination.qkv_split_shapes == source.qkv_split_shapes assert destination.qkv_split_shapes_global == source.qkv_split_shapes_global assert destination.qkv_gtp_pad_length == 3 diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index 5e3c9038497..646c2d1cab1 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import json import os import pytest @@ -12,6 +13,9 @@ from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import ( + get_gpt_heterogeneous_layer_spec, +) from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer from megatron.core.optimizer.emerging_optimizers import ( HAVE_EMERGING_OPTIMIZERS, @@ -27,6 +31,9 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.heterogeneous.heterogeneous_config import ( + HeterogeneousTransformerConfig, +) from tests.unit_tests.test_utilities import Utils if HAVE_EMERGING_OPTIMIZERS: @@ -562,6 +569,97 @@ def test_optimizer_factory_tags_qkv_when_tp_exceeds_query_groups(self): assert qkv_weight.qkv_split_shapes_global == [8, 4, 4] assert not qkv_weight.qkv_split_groups_are_complete + @pytest.mark.parametrize("split_per_head", [False, True]) + def test_optimizer_factory_uses_heterogeneous_layer_qkv_layout(self, split_per_head): + """Each heterogeneous attention layer supplies its own logical QKV layout.""" + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + tp_size = pg_collection.tp.size() + assert tp_size == 2 + model_parallel_cuda_manual_seed(123) + block_configs = { + "block_configs": [ + { + "attention": { + "no_op": False, + "replace_with_linear": False, + "num_query_groups": 2, + }, + "mlp": { + "no_op": False, + "replace_with_linear": False, + "ffn_hidden_size": 16, + }, + }, + { + "attention": { + "no_op": False, + "replace_with_linear": False, + "num_query_groups": 1, + }, + "mlp": { + "no_op": False, + "replace_with_linear": False, + "ffn_hidden_size": 16, + }, + }, + ] + } + transformer_config = HeterogeneousTransformerConfig( + num_layers=2, + hidden_size=8, + num_attention_heads=2, + kv_channels=4, + tensor_model_parallel_size=tp_size, + use_cpu_initialization=False, + add_bias_linear=False, + heterogeneous_layers_config_encoded_json=json.dumps(block_configs), + ) + model = GPTModel( + config=transformer_config, + transformer_layer_spec=get_gpt_heterogeneous_layer_spec(transformer_config), + vocab_size=32, + max_sequence_length=8, + pre_process=False, + post_process=False, + pg_collection=pg_collection, + ) + optimizer_config = OptimizerConfig( + optimizer='muon', + lr=0.01, + use_distributed_optimizer=False, + muon_split_qkv=True, + muon_split_qkv_per_head=split_per_head, + muon_tp_mode="blockwise", + ) + + optimizer = get_megatron_optimizer( + config=optimizer_config, + model_chunks=[model], + use_gloo_process_groups=False, + pg_collection=pg_collection, + ) + + first_qkv = model.decoder.layers[0].self_attention.linear_qkv.weight + second_qkv = model.decoder.layers[1].self_attention.linear_qkv.weight + assert optimizer is not None + assert transformer_config.num_query_groups == 2 + assert first_qkv.qkv_layout.num_query_groups == 2 + assert second_qkv.qkv_layout.num_query_groups == 1 + assert first_qkv.shape[0] == 12 + assert second_qkv.shape[0] == 8 + assert first_qkv.is_qkv + assert second_qkv.is_qkv + if split_per_head: + assert first_qkv.qkv_split_shapes_global == [4] * 6 + assert second_qkv.qkv_split_shapes_global == [4] * 4 + assert first_qkv.qkv_split_heads_are_complete + assert second_qkv.qkv_split_heads_are_complete + else: + assert first_qkv.qkv_split_shapes_global == [4, 4, 4] * 2 + assert second_qkv.qkv_split_shapes_global == [8, 4, 4] + assert first_qkv.qkv_split_groups_are_complete + assert not second_qkv.qkv_split_groups_are_complete + def test_optimizer_factory_skips_mismatched_qkv_layout(self): """A QKV layout mismatch falls back to whole-matrix Muon orthogonalization.""" pg_collection = ProcessGroupCollection.use_mpu_process_groups() From af711504cb1027b66f2d7b61d5d8a1cba9b76ed0 Mon Sep 17 00:00:00 2001 From: mkhona Date: Tue, 18 Aug 2026 17:28:08 -0700 Subject: [PATCH 7/8] Unify Muon QKV and MLA split layouts Signed-off-by: mkhona --- megatron/core/optimizer/__init__.py | 33 ++++--- .../core/optimizer/emerging_optimizers.py | 9 +- megatron/core/transformer/attention.py | 50 +++++++--- .../absorbed_mla.py | 11 ++- .../transformer/multi_latent_attention.py | 19 +++- .../test_tp_attrs_without_init.py | 7 +- tests/unit_tests/test_emerging_optimizers.py | 97 ++++++++++++++++--- .../test_absorbed_mla.py | 5 + .../test_multi_latent_attention.py | 13 +++ 9 files changed, 197 insertions(+), 47 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 329f92acdf1..f4b33447a33 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -788,21 +788,28 @@ def _get_megatron_emerging_optimizer( 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: - qkv_layout = getattr(param, 'qkv_layout', None) - if qkv_layout is None: + qkv_layout = getattr(param, 'qkv_layout', None) + if (qkv_layout is not None or 'linear_qkv.weight' in name) and len(param.shape) == 2: + if qkv_layout is not None: + qkv_split_shapes = _get_qkv_split_shapes( + qkv_layout, split_qkv_per_head=config.muon_split_qkv_per_head + ) + global_split_shapes = ( + qkv_split_shapes + if config.muon_split_qkv_per_head + else qkv_split_shapes * qkv_layout.num_groups + ) + else: # Backward compatibility for custom QKV modules that do not annotate # their weight with the owning attention layer's logical layout. - qkv_layout = model_chunk.config - qkv_split_shapes = _get_qkv_split_shapes( - qkv_layout, split_qkv_per_head=config.muon_split_qkv_per_head - ) - global_split_shapes = ( - qkv_split_shapes - if config.muon_split_qkv_per_head - else qkv_split_shapes * qkv_layout.num_query_groups - ) + qkv_split_shapes = _get_qkv_split_shapes( + model_chunk.config, split_qkv_per_head=config.muon_split_qkv_per_head + ) + global_split_shapes = ( + qkv_split_shapes + if config.muon_split_qkv_per_head + else qkv_split_shapes * model_chunk.config.num_query_groups + ) tp_group = ( pg_collection.expt_tp diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index ae4ebec030a..8bce97892b4 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -149,12 +149,15 @@ def _get_qkv_split_shapes(model_cfg, split_qkv_per_head: bool = False) -> list[i """Compute fused QKV split shapes from logical attention layout metadata. Args: - model_cfg: Object exposing ``num_attention_heads``, ``num_query_groups``, - ``kv_channels``, and ``attention_output_gate``. This can be a transformer - config or the owning attention layer's parameter metadata. + model_cfg: Transformer config or an owning attention layer's ``QKVLayout`` metadata. split_qkv_per_head: Return one split size per physical attention head. When false, return the per-query-group Q, gate (if present), K, and V projection widths. """ + if hasattr(model_cfg, 'projection_split_shapes'): + if split_qkv_per_head: + return list(model_cfg.per_head_split_shapes) * model_cfg.num_groups + return list(model_cfg.projection_split_shapes) + query_projection_size = ( model_cfg.num_attention_heads // model_cfg.num_query_groups * model_cfg.kv_channels ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 1b82028c9fa..abc1bd78b16 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -138,12 +138,45 @@ @dataclass(frozen=True) class QKVLayout: - """Logical layout metadata for a fused QKV projection weight.""" + """Logical row layout for a packed attention projection weight. - num_attention_heads: int - num_query_groups: int - kv_channels: int - attention_output_gate: bool + ``projection_split_shapes`` describes the projection slices repeated in every group. + ``per_head_split_shapes`` describes the independently orthogonalizable head slices in the + same group. Standard fused QKV has one group per query group, while MLA up-projections have + one group per attention head. + """ + + num_groups: int + projection_split_shapes: tuple[int, ...] + per_head_split_shapes: tuple[int, ...] + + @classmethod + def from_standard_attention_config(cls, config: TransformerConfig) -> 'QKVLayout': + """Build the fused QKV row layout described by a transformer config.""" + assert config.num_query_groups is not None + assert config.kv_channels is not None + num_query_heads_per_group = config.num_attention_heads // config.num_query_groups + projection_split_shapes = [num_query_heads_per_group * config.kv_channels] + per_head_split_shapes = [config.kv_channels] * num_query_heads_per_group + if config.attention_output_gate: + projection_split_shapes.append(num_query_heads_per_group * config.kv_channels) + per_head_split_shapes += [config.kv_channels] * num_query_heads_per_group + projection_split_shapes += [config.kv_channels, config.kv_channels] + per_head_split_shapes += [config.kv_channels, config.kv_channels] + return cls( + num_groups=config.num_query_groups, + projection_split_shapes=tuple(projection_split_shapes), + per_head_split_shapes=tuple(per_head_split_shapes), + ) + + @classmethod + def from_repeated_splits(cls, num_groups: int, split_shapes: tuple[int, ...]) -> 'QKVLayout': + """Build a layout whose projection slices are repeated once per attention head.""" + return cls( + num_groups=num_groups, + projection_split_shapes=split_shapes, + per_head_split_shapes=split_shapes, + ) class LinearQkvInterface(Protocol): @@ -1703,12 +1736,7 @@ def __init__( pg_collection=self.pg_collection, name=(name + ".linear_qkv") if name is not None else None, ) - self.linear_qkv.weight.qkv_layout = QKVLayout( - num_attention_heads=self.config.num_attention_heads, - num_query_groups=self.config.num_query_groups, - kv_channels=self.config.kv_channels, - attention_output_gate=self.config.attention_output_gate, - ) + self.linear_qkv.weight.qkv_layout = QKVLayout.from_standard_attention_config(self.config) # Resolve which norm class to use for Q and K. # Config selects the default norm class; spec overrides if set. diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index e0b6af7aa7f..d24a1da6a04 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -33,7 +33,7 @@ gather_from_tensor_model_parallel_region, scatter_to_sequence_parallel_region, ) -from megatron.core.transformer.attention import Attention +from megatron.core.transformer.attention import Attention, QKVLayout from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -372,6 +372,15 @@ def __init__( name=(name + ".linear_kv_up_proj") if name is not None else None, ) + q_up_proj = self.linear_q_proj if self.config.q_lora_rank is None else self.linear_q_up_proj + q_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.config.num_attention_heads, + (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim), + ) + self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim) + ) + if self.config.q_lora_rank is not None: self.q_layernorm = build_module( layer_classes["q_layernorm"], diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index aa21e78ce86..20eb5a27d3b 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -34,7 +34,7 @@ gather_from_tensor_model_parallel_region, scatter_to_sequence_parallel_region, ) -from megatron.core.transformer.attention import Attention, LinearProjBuilder +from megatron.core.transformer.attention import Attention, LinearProjBuilder, QKVLayout from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.mla_qk_norm_config import QKNormConfigResolver from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -616,6 +616,15 @@ def __init__( name=(name + ".linear_kv_up_proj") if name is not None else None, ) + q_up_proj = self.linear_q_proj if self.config.q_lora_rank is None else self.linear_q_up_proj + q_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.config.num_attention_heads, + (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim), + ) + self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim) + ) + if self.config.q_lora_rank is not None: self.q_layernorm = layer_classes["q_layernorm"]( hidden_size=self.config.q_lora_rank, @@ -1329,6 +1338,14 @@ def __init__( name=(name + ".linear_kv_up_proj") if name is not None else None, ) + self.linear_q_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.config.num_attention_heads, + (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim), + ) + self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim) + ) + self.q_layernorm = layer_classes["q_layernorm"]( hidden_size=self.config.q_lora_rank, config=self.config, diff --git a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py index 983c11c202c..09a592110d4 100644 --- a/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py +++ b/tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py @@ -97,10 +97,9 @@ def test_copy_tensor_model_parallel_attributes_preserves_qkv_metadata(): destination = torch.empty_like(source) source.is_qkv = True source.qkv_layout = QKVLayout( - num_attention_heads=8, - num_query_groups=2, - kv_channels=64, - attention_output_gate=False, + num_groups=2, + projection_split_shapes=(256, 64, 64), + per_head_split_shapes=(64, 64, 64, 64, 64, 64), ) source.qkv_split_shapes = [2, 2] source.qkv_split_shapes_global = [2] * 4 diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index 646c2d1cab1..97f9d253251 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -30,7 +30,8 @@ from megatron.core.optimizer.muon import get_megatron_muon_optimizer from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer import TransformerConfig +from megatron.core.transformer import MLATransformerConfig, TransformerConfig +from megatron.core.transformer.attention import QKVLayout from megatron.core.transformer.heterogeneous.heterogeneous_config import ( HeterogeneousTransformerConfig, ) @@ -95,6 +96,10 @@ def test_muon_qkv_split_shapes(): assert _get_qkv_split_shapes(config, split_qkv_per_head=True) == [64] * 32 assert _get_qkv_split_shapes(gated_config, split_qkv_per_head=True) == [64] * 48 + mla_layout = QKVLayout.from_repeated_splits(4, (128, 64)) + assert _get_qkv_split_shapes(mla_layout) == [128, 64] + assert _get_qkv_split_shapes(mla_layout, split_qkv_per_head=True) == [128, 64] * 4 + def test_muon_local_qkv_head_split_shapes_can_differ_by_tp_rank(): """Rank-local per-head layouts report complete and fragmented heads.""" @@ -584,11 +589,7 @@ def test_optimizer_factory_uses_heterogeneous_layer_qkv_layout(self, split_per_h "replace_with_linear": False, "num_query_groups": 2, }, - "mlp": { - "no_op": False, - "replace_with_linear": False, - "ffn_hidden_size": 16, - }, + "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 16}, }, { "attention": { @@ -596,11 +597,7 @@ def test_optimizer_factory_uses_heterogeneous_layer_qkv_layout(self, split_per_h "replace_with_linear": False, "num_query_groups": 1, }, - "mlp": { - "no_op": False, - "replace_with_linear": False, - "ffn_hidden_size": 16, - }, + "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 16}, }, ] } @@ -643,8 +640,8 @@ def test_optimizer_factory_uses_heterogeneous_layer_qkv_layout(self, split_per_h second_qkv = model.decoder.layers[1].self_attention.linear_qkv.weight assert optimizer is not None assert transformer_config.num_query_groups == 2 - assert first_qkv.qkv_layout.num_query_groups == 2 - assert second_qkv.qkv_layout.num_query_groups == 1 + assert first_qkv.qkv_layout.num_groups == 2 + assert second_qkv.qkv_layout.num_groups == 1 assert first_qkv.shape[0] == 12 assert second_qkv.shape[0] == 8 assert first_qkv.is_qkv @@ -660,6 +657,78 @@ def test_optimizer_factory_uses_heterogeneous_layer_qkv_layout(self, split_per_h assert first_qkv.qkv_split_groups_are_complete assert not second_qkv.qkv_split_groups_are_complete + @pytest.mark.parametrize("split_per_head", [False, True]) + @pytest.mark.parametrize("q_lora_rank", [None, 4]) + def test_optimizer_factory_uses_mla_up_projection_layouts(self, split_per_head, q_lora_rank): + """MLA up-projections use module-owned layouts with TP-aware Muon splitting.""" + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + tp_size = pg_collection.tp.size() + assert tp_size == 2 + model_parallel_cuda_manual_seed(123) + transformer_config = MLATransformerConfig( + num_layers=1, + hidden_size=8, + num_attention_heads=2, + q_lora_rank=q_lora_rank, + kv_lora_rank=4, + qk_head_dim=4, + qk_pos_emb_head_dim=2, + v_head_dim=3, + tensor_model_parallel_size=tp_size, + use_cpu_initialization=False, + add_bias_linear=False, + multi_latent_attention=True, + rope_type="rope", + rotary_base=10000, + original_max_position_embeddings=8, + ) + model = GPTModel( + config=transformer_config, + transformer_layer_spec=get_gpt_layer_local_spec(multi_latent_attention=True), + vocab_size=32, + max_sequence_length=8, + pre_process=False, + post_process=False, + pg_collection=pg_collection, + ) + optimizer_config = OptimizerConfig( + optimizer='muon', + lr=0.01, + use_distributed_optimizer=False, + muon_split_qkv=True, + muon_split_qkv_per_head=split_per_head, + muon_tp_mode="blockwise", + ) + + optimizer = get_megatron_optimizer( + config=optimizer_config, + model_chunks=[model], + use_gloo_process_groups=False, + pg_collection=pg_collection, + ) + + attention = model.decoder.layers[0].self_attention + q_up_weight = ( + attention.linear_q_proj.weight + if q_lora_rank is None + else attention.linear_q_up_proj.weight + ) + kv_up_weight = attention.linear_kv_up_proj.weight + assert optimizer is not None + assert q_up_weight.is_qkv + assert kv_up_weight.is_qkv + assert q_up_weight.qkv_split_shapes_global == [4, 2] * 2 + assert kv_up_weight.qkv_split_shapes_global == [4, 3] * 2 + if split_per_head: + assert q_up_weight.qkv_split_heads_are_complete + assert kv_up_weight.qkv_split_heads_are_complete + else: + assert q_up_weight.qkv_split_groups_are_complete + assert kv_up_weight.qkv_split_groups_are_complete + if q_lora_rank is not None: + assert not getattr(attention.linear_q_down_proj.weight, 'is_qkv', False) + assert not getattr(attention.linear_kv_down_proj.weight, 'is_qkv', False) + def test_optimizer_factory_skips_mismatched_qkv_layout(self): """A QKV layout mismatch falls back to whole-matrix Muon orthogonalization.""" pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -779,7 +848,7 @@ def center_rows(x, tp_group=None, partition_dim=None): expected = expected_global[tp_rank * 3 : (tp_rank + 1) * 3] torch.testing.assert_close(actual, expected) - @pytest.mark.parametrize("split_shapes", ([4, 2, 2], [4, 4, 2, 2])) + @pytest.mark.parametrize("split_shapes", ([4, 2, 2], [4, 4, 2, 2], [3, 5])) def test_muon_optimizer_projection_split_gathers_fragmented_query_group(self, split_shapes): """Projection splitting reconstructs a query group split over TP ranks.""" pg_collection = ProcessGroupCollection.use_mpu_process_groups() diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py index fc1778f649f..7f2c2afa16f 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_absorbed_mla.py @@ -430,6 +430,11 @@ def test_functionality(tp_cp: List[int], qkv_format: str, down_proj_use_column_p cp_comm_type="all_gather" if cp_size > 1 else None, pg_collection=None, ).cuda() + assert absorbed_mla.linear_q_up_proj.weight.qkv_layout.num_groups == config.num_attention_heads + assert absorbed_mla.linear_kv_up_proj.weight.qkv_layout.projection_split_shapes == ( + config.qk_head_dim, + config.v_head_dim, + ) state_dict = standard_mla.state_dict() absorbed_mla.load_state_dict(state_dict) diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index 95517f7176d..5cc21bbfd26 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -219,6 +219,16 @@ def test_dynamic_inference_forwards_decode_only_to_flash_attention(self, is_deco def test_constructor(self): assert isinstance(self.parallel_attention, MLASelfAttention) assert self.parallel_attention.layer_number == 1 + assert self.parallel_attention.linear_q_up_proj.weight.qkv_layout.num_groups == 4 + assert ( + self.parallel_attention.linear_q_up_proj.weight.qkv_layout.projection_split_shapes + == (128, 64) + ) + assert self.parallel_attention.linear_kv_up_proj.weight.qkv_layout.num_groups == 4 + assert ( + self.parallel_attention.linear_kv_up_proj.weight.qkv_layout.projection_split_shapes + == (128, 128) + ) num_weights = sum([p.numel() for p in self.parallel_attention.parameters()]) assert num_weights == 65036 @@ -1700,6 +1710,9 @@ def test_constructor(self): assert isinstance(self.fused_attention, MLASelfAttention) assert self.fused_attention.layer_number == 1 assert hasattr(self.fused_attention, 'linear_qkv_down_proj') + assert self.fused_attention.linear_q_up_proj.weight.qkv_layout.num_groups == 4 + assert self.fused_attention.linear_kv_up_proj.weight.qkv_layout.num_groups == 4 + assert getattr(self.fused_attention.linear_qkv_down_proj.weight, 'qkv_layout', None) is None def test_fused_weight_shape(self): config = self.transformer_config From febf2f185b664ccb0ce22fdceff27f0aab740b46 Mon Sep 17 00:00:00 2001 From: mkhona Date: Wed, 19 Aug 2026 11:22:51 -0700 Subject: [PATCH 8/8] Address QKV layout review feedback Signed-off-by: mkhona --- megatron/core/optimizer/__init__.py | 30 +++++++++---------- megatron/core/transformer/attention.py | 6 ++-- .../absorbed_mla.py | 4 +-- .../transformer/multi_latent_attention.py | 8 ++--- .../test_tp_gtp.py | 21 ++++++------- tests/unit_tests/test_emerging_optimizers.py | 2 +- 6 files changed, 36 insertions(+), 35 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index f4b33447a33..ca98033b0ca 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -794,7 +794,7 @@ def _get_megatron_emerging_optimizer( qkv_split_shapes = _get_qkv_split_shapes( qkv_layout, split_qkv_per_head=config.muon_split_qkv_per_head ) - global_split_shapes = ( + logical_split_shapes = ( qkv_split_shapes if config.muon_split_qkv_per_head else qkv_split_shapes * qkv_layout.num_groups @@ -805,7 +805,7 @@ def _get_megatron_emerging_optimizer( qkv_split_shapes = _get_qkv_split_shapes( model_chunk.config, split_qkv_per_head=config.muon_split_qkv_per_head ) - global_split_shapes = ( + logical_split_shapes = ( qkv_split_shapes if config.muon_split_qkv_per_head else qkv_split_shapes * model_chunk.config.num_query_groups @@ -819,16 +819,16 @@ def _get_megatron_emerging_optimizer( tp_size = get_pg_size(tp_group) tp_rank = get_pg_rank(tp_group) gtp_remat_group = ( - pg_collection.expt_gtp_remat - if getattr(param, 'expert_tp', False) - else pg_collection.gtp_remat + ( + pg_collection.expt_gtp_remat + if getattr(param, 'expert_tp', False) + else pg_collection.gtp_remat + ) + if getattr(param, 'is_gtp_weight_remat', False) + else None ) - if getattr(param, 'is_gtp_weight_remat', False): - gtp_size = get_pg_size(gtp_remat_group) - gtp_rank = get_pg_rank(gtp_remat_group) - else: - gtp_size = 1 - gtp_rank = 0 + gtp_size = get_pg_size(gtp_remat_group) + gtp_rank = get_pg_rank(gtp_remat_group) qkv_gtp_pad_length = ( int(getattr(param, 'pad_length', 0)) @@ -843,13 +843,13 @@ def _get_megatron_emerging_optimizer( f"physical_tp_local_rows={physical_tp_local_rows}" ) logical_tp_local_rows = physical_tp_local_rows - qkv_gtp_pad_length - expected_global_rows = logical_tp_local_rows * tp_size - if expected_global_rows != sum(global_split_shapes): + expected_logical_rows = logical_tp_local_rows * tp_size + if expected_logical_rows != sum(logical_split_shapes): log_single_rank( logger, logging.DEBUG, f"Emerging optimizer QKV split skipped for {name}: " - f"global_rows={sum(global_split_shapes)}, " + f"logical_rows={sum(logical_split_shapes)}, " f"local_rows={param.shape[0]}, tp_size={tp_size}, " f"gtp_remat_size={gtp_size}, " f"gtp_pad_length={qkv_gtp_pad_length}", @@ -863,7 +863,7 @@ def _get_megatron_emerging_optimizer( continue param.is_qkv = True - param.qkv_split_shapes_global = global_split_shapes + param.qkv_split_shapes_global = logical_split_shapes param.qkv_gtp_pad_length = qkv_gtp_pad_length local_start = tp_rank * logical_tp_local_rows + gtp_rank * param.shape[0] if config.muon_split_qkv_per_head: diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index abc1bd78b16..750d20c4e4c 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -151,7 +151,7 @@ class QKVLayout: per_head_split_shapes: tuple[int, ...] @classmethod - def from_standard_attention_config(cls, config: TransformerConfig) -> 'QKVLayout': + def from_transformer_config(cls, config: TransformerConfig) -> 'QKVLayout': """Build the fused QKV row layout described by a transformer config.""" assert config.num_query_groups is not None assert config.kv_channels is not None @@ -170,7 +170,7 @@ def from_standard_attention_config(cls, config: TransformerConfig) -> 'QKVLayout ) @classmethod - def from_repeated_splits(cls, num_groups: int, split_shapes: tuple[int, ...]) -> 'QKVLayout': + def from_splits(cls, num_groups: int, split_shapes: tuple[int, ...]) -> 'QKVLayout': """Build a layout whose projection slices are repeated once per attention head.""" return cls( num_groups=num_groups, @@ -1736,7 +1736,7 @@ def __init__( pg_collection=self.pg_collection, name=(name + ".linear_qkv") if name is not None else None, ) - self.linear_qkv.weight.qkv_layout = QKVLayout.from_standard_attention_config(self.config) + self.linear_qkv.weight.qkv_layout = QKVLayout.from_transformer_config(self.config) # Resolve which norm class to use for Q and K. # Config selects the default norm class; spec overrides if set. diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index d24a1da6a04..f776c102b97 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -373,11 +373,11 @@ def __init__( ) q_up_proj = self.linear_q_proj if self.config.q_lora_rank is None else self.linear_q_up_proj - q_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + q_up_proj.weight.qkv_layout = QKVLayout.from_splits( self.config.num_attention_heads, (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim), ) - self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_splits( self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim) ) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 20eb5a27d3b..aff59ff7c74 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -617,11 +617,11 @@ def __init__( ) q_up_proj = self.linear_q_proj if self.config.q_lora_rank is None else self.linear_q_up_proj - q_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + q_up_proj.weight.qkv_layout = QKVLayout.from_splits( self.config.num_attention_heads, (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim), ) - self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_splits( self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim) ) @@ -1338,11 +1338,11 @@ def __init__( name=(name + ".linear_kv_up_proj") if name is not None else None, ) - self.linear_q_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.linear_q_up_proj.weight.qkv_layout = QKVLayout.from_splits( self.config.num_attention_heads, (self.config.qk_head_dim, self.config.qk_pos_emb_head_dim), ) - self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_repeated_splits( + self.linear_kv_up_proj.weight.qkv_layout = QKVLayout.from_splits( self.config.num_attention_heads, (self.config.qk_head_dim, self.config.v_head_dim) ) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py index 1b5ddf58388..e8858bfbcb2 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_tp_gtp.py @@ -418,10 +418,11 @@ def _worker_muon_qkv_padding( gtp_local_rows = physical_tp_rows // gtp_remat_size hidden_size = 4 - global_grad = torch.arange(logical_rows * hidden_size, dtype=torch.float32, device="cuda").view( - logical_rows, hidden_size - ) - tp_grad = global_grad[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows] + # Full logical gradient before TP/GTP sharding; physical GTP padding is not included. + logical_grad = torch.arange( + logical_rows * hidden_size, dtype=torch.float32, device="cuda" + ).view(logical_rows, hidden_size) + tp_grad = logical_grad[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows] # Use a sentinel so the test detects padding entering the orthogonalization. physical_tp_grad = F.pad(tp_grad, (0, 0, 0, pad_length), value=10_000.0) local_grad = physical_tp_grad[ @@ -457,21 +458,21 @@ def center_rows(x, tp_group=None, partition_dim=None): actual = optimizer.orthogonalize(param, local_grad) if split_per_head: - expected_global = torch.cat( - [center_rows(head) for head in torch.split(global_grad, split_shapes, dim=0)], dim=0 + expected_logical_grad = torch.cat( + [center_rows(head) for head in torch.split(logical_grad, split_shapes, dim=0)], dim=0 ) else: - grouped_grad = global_grad.view(1, logical_rows, hidden_size) + grouped_grad = logical_grad.view(1, logical_rows, hidden_size) projections = torch.split(grouped_grad, split_shapes, dim=1) - expected_global = torch.cat( + expected_logical_grad = torch.cat( [ center_rows(projection.reshape(-1, hidden_size)).view_as(projection) for projection in projections ], dim=1, - ).view_as(global_grad) + ).view_as(logical_grad) - expected_tp = expected_global[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows] + expected_tp = expected_logical_grad[tp_rank * logical_tp_rows : (tp_rank + 1) * logical_tp_rows] expected_tp = F.pad(expected_tp, (0, 0, 0, pad_length)) expected = expected_tp[gtp_rank * gtp_local_rows : (gtp_rank + 1) * gtp_local_rows] diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index 97f9d253251..e9688528f89 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -96,7 +96,7 @@ def test_muon_qkv_split_shapes(): assert _get_qkv_split_shapes(config, split_qkv_per_head=True) == [64] * 32 assert _get_qkv_split_shapes(gated_config, split_qkv_per_head=True) == [64] * 48 - mla_layout = QKVLayout.from_repeated_splits(4, (128, 64)) + mla_layout = QKVLayout.from_splits(4, (128, 64)) assert _get_qkv_split_shapes(mla_layout) == [128, 64] assert _get_qkv_split_shapes(mla_layout, split_qkv_per_head=True) == [128, 64] * 4