diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 70f757f2889..a88dfe539e3 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -4,7 +4,7 @@ import warnings from collections import defaultdict from dataclasses import astuple -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import torch from torch.optim import SGD as CPUSGD @@ -722,6 +722,136 @@ def check_config_overrides_consistency( return True +def _tag_muon_split_qkv_parameters( + config: OptimizerConfig, model_chunks: list[MegatronModule] +) -> None: + """Tag standard QKV and MLA up-projection weights with Muon split metadata. + + Note that MLA down-projection weights aren't split with the current implementation. + + The applied tags are: + + - `is_qkv` (bool): whether the layer is a fused/packed collection of standard Multi-Head + Attention Q, K, and V projections. The layer contains a single weight, which includes all of + the Q, K, and V projections. + - `qkv_split_shapes` (tuple[int, ...] | None): per-group row counts for the projection slices + packed in this weight, if applicable. The gradients for the tagged parameter will be split + along the row dimension. The splits will be orthogonalized independently. + + For example, without an Attention output gate, we have the following for the fused/packed + standard Multi-Head Attention QKV projection when not using tensor model parallelism: + + ``` + # Remember that PyTorch uses transposed weights. + linear_qkv.weight.shape == ( + ( + num_attention_heads * kv_channels # <- Qs + + kv_channels * num_query_groups # <- Ks + + kv_channels * num_query_groups # <- Vs + ), + hidden_size, + ) + + linear_qkv.qkv_split_shapes = ( + num_attention_heads // num_query_groups * kv_channels, # <- Q + kv_channels, # <- K + kv_channels, # <- V + ) + ``` + + In the optimizer step, we transform the gradient like the following: + ``` + # The initial shape is like `linear_qkv.weight.shape` before. + linear_qkv.grad.shape == linear_qkv.weight.shape + + qs_grad, ks_grad, vs_grad = muon_split_grads(linear_qkv.grad, linear_qkv.qkv_split_shapes) + + qs_grad.shape == ( + num_query_groups, + num_attention_heads // num_query_groups * kv_channels, + hidden_size, + ) + ks_grad.shape == vs_grad.shape == ( + num_query_groups, + kv_channels, + hidden_size, + ) + + # Contract first two dimensions. + qs_grad = qs_grad.reshape(num_attention_heads * kv_channels, hidden_size) + qs_grad = orthogonalize(qs_grad) + + ks_grad = ks_grad.reshape(num_query_groups * kv_channels, hidden_size) + ks_grad = orthogonalize(ks_grad) + + vs_grad = vs_grad.reshape(num_query_groups * kv_channels, hidden_size) + vs_grad = orthogonalize(vs_grad) + + result_grad = muon_concat_grads(qs_grad, ks_grad, vs_grad) + result_grad.shape == linear_qkv.grad.shape + ``` + """ + split_qkv = getattr(config, 'muon_split_qkv', True) + + def _is_muon_split_metadata_managed_param( + name: str, attn_variant: Literal["mha", "mla"] + ) -> bool: + """Return whether the parameter with the given name should be managed by Muon split tags. + + E.g., standard QKV or MLA projection weights. + """ + if 'linear_qkv.weight' in name: + return True + if attn_variant != "mla": + return False + return any( + f'{projection}.weight' in name + for projection in ( + 'linear_q_proj', + 'linear_q_up_proj', + 'linear_kv_up_proj', + 'linear_q_down_proj', + 'linear_kv_down_proj', + 'linear_qkv_down_proj', + ) + ) + + for model_chunk in model_chunks: + model_cfg = get_model_config(model_chunk) + attn_variant = "mla" if getattr(model_cfg, 'multi_latent_attention', False) else "mha" + standard_qkv_split_shapes = None + mla_q_split_shapes = None + mla_kv_split_shapes = None + if attn_variant == "mla": + mla_q_split_shapes = (model_cfg.qk_head_dim, model_cfg.qk_pos_emb_head_dim) + mla_kv_split_shapes = (model_cfg.qk_head_dim, model_cfg.v_head_dim) + + for name, param in model_chunk.named_parameters(): + if not param.requires_grad: + continue + + is_managed_param = _is_muon_split_metadata_managed_param(name, attn_variant) + if is_managed_param: + # Remove existing Muon split tags from a parameter. + for attr in ("is_qkv", "qkv_split_shapes"): + if hasattr(param, attr): + delattr(param, attr) + if not is_managed_param or not split_qkv or len(param.shape) != 2: + continue + + if 'linear_qkv.weight' in name: + if standard_qkv_split_shapes is None: + standard_qkv_split_shapes = tuple(_get_qkv_split_shapes(model_cfg)) + param.qkv_split_shapes = standard_qkv_split_shapes + param.is_qkv = True + elif mla_q_split_shapes is not None and ( + 'linear_q_up_proj.weight' in name or 'linear_q_proj.weight' in name + ): + param.qkv_split_shapes = mla_q_split_shapes + elif mla_kv_split_shapes is not None and 'linear_kv_up_proj.weight' in name: + param.qkv_split_shapes = mla_kv_split_shapes + + def _get_megatron_emerging_optimizer( config: OptimizerConfig, model_chunks: List[MegatronModule], @@ -774,28 +904,15 @@ def _get_megatron_emerging_optimizer( log_single_rank(logger, logging.INFO, f'Setting up emerging optimizer with config {config}') - # Tag parameters with optimizer-specific attributes (expert_tp, is_qkv). + # Tag parameters with optimizer-specific attributes. for model_chunk in model_chunks: - qkv_split_shapes = None for name, param in model_chunk.named_parameters(): if not param.requires_grad: continue if 'experts' in name and 'shared' not in name: param.expert_tp = True - # TODO(deyuf): support MLA - if 'linear_qkv.weight' in name and len(param.shape) == 2: - if qkv_split_shapes is None: - qkv_split_shapes = _get_qkv_split_shapes(model_chunk.config) - if param.shape[0] % sum(qkv_split_shapes) == 0: - param.is_qkv = True - param.qkv_split_shapes = qkv_split_shapes - else: - log_single_rank( - logger, - logging.DEBUG, - f"Emerging optimizer QKV split skipped for {name}: " - f"shape={tuple(param.shape)}, split_shapes={qkv_split_shapes}", - ) + if eopt_name in ('muon', 'adaptive_muon'): + _tag_muon_split_qkv_parameters(config, model_chunks) # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam). # For Muon-family optimizers, the scalar optimizer that handles non-linear/embedding diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 53ac956b35c..4c5e9f29a30 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -266,6 +266,58 @@ 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 _get_muon_split_shapes(self, p: torch.Tensor) -> tuple[int, ...] | None: + """Return per-parameter Muon split shapes, if this parameter should be split.""" + split_shapes: tuple[int, ...] | None = getattr(p, "qkv_split_shapes", None) + if split_shapes is None: + if self.is_qkv_fn is None or not self.is_qkv_fn(p): + return None + split_shapes = self.qkv_split_shapes + if split_shapes is None: + raise ValueError( + "Muon QKV split was requested for a parameter, but `qkv_split_shapes` is not " + "set." + ) + + if not split_shapes or any(shape <= 0 for shape in split_shapes): + raise ValueError(f"Muon split shapes must be positive integers, got {split_shapes}.") + return split_shapes + + def _orthogonalize_split_grad( + self, + param: torch.Tensor, + grad: torch.Tensor, + split_shapes: tuple[int, ...], + tp_group: torch.distributed.ProcessGroup | None, + partition_dim: int | None, + ) -> torch.Tensor: + """Orthogonalize a fused projection gradient by splitting its row layout first.""" + grad_shape = grad.shape + split_size = sum(split_shapes) + if grad_shape[0] % split_size != 0: + raise ValueError( + f"Muon split parameter has incompatible grad shape `{tuple(grad_shape)}` " + f"for split shapes `{split_shapes}`: `grad.shape[0]` must be divisible by " + f"`sum(split_shapes)={split_size}`." + ) + + log_single_rank( + logger, + logging.DEBUG, + f'muon split grad shape `{grad_shape}`, split shapes `{split_shapes}`', + ) + num_groups = grad_shape[0] // split_size + split_grads = torch.split(grad.view(num_groups, split_size, -1), split_shapes, dim=1) + split_grads = [g.reshape(-1, grad_shape[-1]) for g in split_grads] + + split_grads = [ + self.scaled_orthogonalize_fn_with_gtp_remat(param, g, tp_group, partition_dim).view( + num_groups, -1, grad_shape[-1] + ) + for g in split_grads + ] + return torch.cat(split_grads, dim=1).view(grad_shape) + def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: """Orthogonalize the momentum. @@ -291,38 +343,15 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t if partition_dim == -1: partition_dim = None - if self.split_qkv and self.is_qkv_fn(p): # type: ignore[misc] - grad_shape = grad.shape - qkv_split_shapes = getattr(p, "qkv_split_shapes", None) - if qkv_split_shapes is None: - qkv_split_shapes = self.qkv_split_shapes - if qkv_split_shapes is None: - 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}" + grad_was_split = False + if self.split_qkv: + split_shapes = self._get_muon_split_shapes(p) + if split_shapes is not None: + grad = self._orthogonalize_split_grad( + p, grad, split_shapes, tp_group, partition_dim ) - log_single_rank( - logger, - logging.DEBUG, - f'qkv split grad shape {grad_shape}, split shapes {qkv_split_shapes}', - ) - num_query_groups = grad_shape[0] // qkv_split_dim - qkv_grads = torch.split( - grad.view(num_query_groups, qkv_split_dim, -1), qkv_split_shapes, dim=1 - ) - qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads] - - qkv_grads = [ - self.scaled_orthogonalize_fn_with_gtp_remat(p, g, tp_group, partition_dim).view( - num_query_groups, -1, grad_shape[-1] - ) - for g in qkv_grads - ] - grad = torch.cat(qkv_grads, dim=1).view(grad_shape) - else: + grad_was_split = True + if not grad_was_split: grad = self.scaled_orthogonalize_fn_with_gtp_remat(p, grad, tp_group, partition_dim) return grad diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index e3b9f666fb2..3b6ed003734 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -720,6 +720,101 @@ def test_muon_optimizer_qkv_split(): ), "Weights should be different between split_qkv=True and split_qkv=False" +def test_muon_optimizer_qkv_split_uses_legacy_is_qkv_metadata(): + """Test that legacy is_qkv + qkv_split_shapes still uses the split path.""" + split_shapes = (3, 2, 1) + num_groups = 4 + hidden_size = 5 + param = torch.nn.Parameter( + torch.zeros(num_groups * sum(split_shapes), hidden_size, device='cuda') + ) + param.is_qkv = True + optimizer = TensorParallelMuon( + params=[param], + lr=0.01, + split_qkv=True, + is_qkv_fn=lambda p: getattr(p, 'is_qkv', False), + qkv_split_shapes=split_shapes, + num_ns_steps=1, + pg_collection=None, + tp_mode="duplicated", + ) + + calls = [] + + def fake_orthogonalize(grad, tp_group, partition_dim): + calls.append(tuple(grad.shape)) + return torch.full_like(grad, len(calls)) + + optimizer.scaled_orthogonalize_fn = fake_orthogonalize + result = optimizer.orthogonalize(param, torch.zeros_like(param)) + + assert calls == [ + (num_groups * split_shapes[0], hidden_size), + (num_groups * split_shapes[1], hidden_size), + (num_groups * split_shapes[2], hidden_size), + ] + result = result.view(num_groups, sum(split_shapes), hidden_size) + assert torch.all(result[:, : split_shapes[0]] == 1).item() + assert torch.all(result[:, split_shapes[0] : split_shapes[0] + split_shapes[1]] == 2).item() + assert torch.all(result[:, -split_shapes[2] :] == 3).item() + + +def test_muon_optimizer_mla_kv_split_uses_parameter_metadata(): + """Test that MLA-style 2-way split metadata uses the split path.""" + split_shapes = (3, 5) + num_heads = 4 + hidden_size = 7 + param = torch.nn.Parameter( + torch.zeros(num_heads * sum(split_shapes), hidden_size, device='cuda') + ) + param.qkv_split_shapes = split_shapes + optimizer = TensorParallelMuon( + params=[param], + lr=0.01, + split_qkv=True, + num_ns_steps=1, + pg_collection=None, + tp_mode="duplicated", + ) + + calls = [] + + def fake_orthogonalize(grad, tp_group, partition_dim): + calls.append(tuple(grad.shape)) + return torch.full_like(grad, len(calls)) + + optimizer.scaled_orthogonalize_fn = fake_orthogonalize + result = optimizer.orthogonalize(param, torch.zeros_like(param)) + + assert calls == [ + (num_heads * split_shapes[0], hidden_size), + (num_heads * split_shapes[1], hidden_size), + ] + result = result.view(num_heads, sum(split_shapes), hidden_size) + assert torch.all(result[:, : split_shapes[0]] == 1).item() + assert torch.all(result[:, split_shapes[0] :] == 2).item() + + +def test_muon_optimizer_mla_split_metadata_validates_grad_shape(): + """Test that invalid MLA split metadata reports an incompatible layout.""" + split_shapes = (3, 5) + hidden_size = 7 + param = torch.nn.Parameter(torch.zeros(sum(split_shapes) + 1, hidden_size, device='cuda')) + param.qkv_split_shapes = split_shapes + optimizer = TensorParallelMuon( + params=[param], + lr=0.01, + split_qkv=True, + num_ns_steps=1, + pg_collection=None, + tp_mode="duplicated", + ) + + with pytest.raises(ValueError, match=r"grad\.shape\[0\].*sum\(split_shapes\)=8"): + optimizer.orthogonalize(param, torch.zeros_like(param)) + + 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') diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index d822f9bbf3f..097aa4d0cd7 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -22,6 +23,7 @@ ParamKey, ParamPredicate, _get_param_groups, + _tag_muon_split_qkv_parameters, check_config_overrides_consistency, get_megatron_optimizer, get_standard_config_overrides, @@ -30,6 +32,7 @@ from megatron.core.optimizer.optimizer import copy_optimizer_param_metadata from megatron.core.optimizer_param_scheduler import ParamGroupOverride from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.layers import copy_tensor_model_parallel_attributes from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig from megatron.core.transformer.enums import AttnMaskType @@ -91,6 +94,125 @@ def test_copy_optimizer_param_metadata_preserves_allreduce(): assert destination.allreduce is False +def _test_muon_split_model_config(multi_latent_attention=False): + return SimpleNamespace( + num_attention_heads=8, + num_query_groups=2, + kv_channels=16, + multi_latent_attention=multi_latent_attention, + qk_head_dim=32, + qk_pos_emb_head_dim=8, + v_head_dim=24, + ) + + +class MuonSplitTaggingModel(nn.Module): + def __init__(self, config, include_standard_qkv=True, include_mla=True): + super().__init__() + self.config = config + self.self_attention = nn.Module() + hidden_size = 4 + + if include_standard_qkv: + qkv_split_shapes = ( + config.num_attention_heads // config.num_query_groups * config.kv_channels, + config.kv_channels, + config.kv_channels, + ) + qkv_size = config.num_query_groups * sum(qkv_split_shapes) + self.self_attention.linear_qkv = nn.Linear(hidden_size, qkv_size, bias=False) + + if include_mla: + q_out = config.num_attention_heads * (config.qk_head_dim + config.qk_pos_emb_head_dim) + kv_out = config.num_attention_heads * (config.qk_head_dim + config.v_head_dim) + self.self_attention.linear_q_proj = nn.Linear(hidden_size, q_out, bias=False) + self.self_attention.linear_q_up_proj = nn.Linear(hidden_size, q_out, bias=False) + self.self_attention.linear_kv_up_proj = nn.Linear(hidden_size, kv_out, bias=False) + self.self_attention.linear_q_down_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.self_attention.linear_kv_down_proj = nn.Linear( + hidden_size, hidden_size, bias=False + ) + self.self_attention.linear_qkv_down_proj = nn.Linear( + hidden_size, hidden_size, bias=False + ) + + +def test_tag_muon_split_qkv_parameters_tags_standard_qkv(): + model = MuonSplitTaggingModel( + _test_muon_split_model_config(multi_latent_attention=False), + include_standard_qkv=True, + include_mla=False, + ) + config = OptimizerConfig(optimizer='muon', muon_split_qkv=True) + + _tag_muon_split_qkv_parameters(config, [model]) + + params = dict(model.named_parameters()) + qkv_weight = params['self_attention.linear_qkv.weight'] + assert qkv_weight.is_qkv + assert qkv_weight.qkv_split_shapes == (64, 16, 16) + + +def test_tag_muon_split_qkv_parameters_tags_mla_up_projections(): + model = MuonSplitTaggingModel( + _test_muon_split_model_config(multi_latent_attention=True), + include_standard_qkv=False, + include_mla=True, + ) + config = OptimizerConfig(optimizer='muon', muon_split_qkv=True) + + _tag_muon_split_qkv_parameters(config, [model]) + + params = dict(model.named_parameters()) + assert params['self_attention.linear_q_proj.weight'].qkv_split_shapes == (32, 8) + assert params['self_attention.linear_q_up_proj.weight'].qkv_split_shapes == (32, 8) + assert params['self_attention.linear_kv_up_proj.weight'].qkv_split_shapes == (32, 24) + + +def test_tag_muon_split_qkv_parameters_skips_mla_down_projections(): + model = MuonSplitTaggingModel( + _test_muon_split_model_config(multi_latent_attention=True), + include_standard_qkv=False, + include_mla=True, + ) + config = OptimizerConfig(optimizer='muon', muon_split_qkv=True) + + _tag_muon_split_qkv_parameters(config, [model]) + + params = dict(model.named_parameters()) + assert not hasattr(params['self_attention.linear_q_down_proj.weight'], 'qkv_split_shapes') + assert not hasattr(params['self_attention.linear_kv_down_proj.weight'], 'qkv_split_shapes') + assert not hasattr(params['self_attention.linear_qkv_down_proj.weight'], 'qkv_split_shapes') + + +def test_tag_muon_split_qkv_parameters_respects_muon_no_split_qkv(): + model = MuonSplitTaggingModel( + _test_muon_split_model_config(multi_latent_attention=True), + include_standard_qkv=True, + include_mla=True, + ) + for param in model.parameters(): + param.is_qkv = True + param.qkv_split_shapes = (1,) + config = OptimizerConfig(optimizer='muon', muon_split_qkv=False) + + _tag_muon_split_qkv_parameters(config, [model]) + + for param in model.parameters(): + assert not hasattr(param, 'is_qkv') + assert not hasattr(param, 'qkv_split_shapes') + + +def test_muon_split_metadata_copied_to_bf16_main_param_clone(): + model_param = torch.nn.Parameter(torch.zeros(4, 4)) + model_param.qkv_split_shapes = (32, 24) + main_param = model_param.detach().clone().float() + + copy_tensor_model_parallel_attributes(main_param, model_param) + + assert main_param.qkv_split_shapes == (32, 24) + + @patch('torch.distributed.get_world_size', return_value=1) @patch( 'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj)