diff --git a/.fork-base.json b/.fork-base.json new file mode 100644 index 00000000000..2c37262545f --- /dev/null +++ b/.fork-base.json @@ -0,0 +1,7 @@ +{ + "schema": 1, + "fork_repo": "https://github.com/AlignmentResearch/Megatron-LM", + "upstream_repo": "https://github.com/NVIDIA/Megatron-LM", + "upstream_branch": "main", + "upstream_base": "d12f6c8c9aff51e166d872fd70151687a8e3f375" +} diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 7f7424e6c02..80f6b8eb2c7 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -11,6 +11,7 @@ import re import warnings from contextlib import contextmanager, nullcontext +from itertools import accumulate from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, cast import torch @@ -2358,6 +2359,7 @@ def __init__( name (str | None): module instance name passed top-down from its paranet module """ self.config = config + self.expert_gemm_backend = config.moe_expert_gemm_backend # TE returns a zero length Tensor when bias=False and # return_bias=True, but we prefer None. So in that case we @@ -2632,6 +2634,110 @@ def _split_grouped_checkpoint_tensor( f"into {self.num_gemms} GEMM shards." ) + def _torch_grouped_weight_layout_is_current(self) -> bool: + """Return whether each expert parameter still views the contiguous buffer.""" + grouped_weight = self._buffers.get('_torch_grouped_weight') + if grouped_weight is None or grouped_weight.shape[0] != self.num_gemms: + return False + + first_weight = self.weight0 + if ( + grouped_weight.device != first_weight.device + or grouped_weight.dtype != first_weight.dtype + or grouped_weight.shape[1:] != first_weight.shape + ): + return False + + storage_pointer = grouped_weight.untyped_storage().data_ptr() + expert_elements = first_weight.numel() + for index in range(self.num_gemms): + weight = getattr(self, f'weight{index}') + grouped_weight_view = grouped_weight[index] + if ( + weight.device != grouped_weight.device + or weight.dtype != grouped_weight.dtype + or weight.shape != grouped_weight_view.shape + or weight.untyped_storage().data_ptr() != storage_pointer + or weight.storage_offset() != index * expert_elements + or weight.stride() != grouped_weight_view.stride() + ): + return False + return True + + @torch.no_grad() + def prepare_torch_grouped_mm(self) -> int: + """Relocate frozen BF16 expert weights into one contiguous allocation. + + The existing per-expert ``Parameter`` objects become views of the new allocation, so + parameter names and checkpoint structure remain unchanged. The backing allocation is a + non-persistent buffer and therefore does not add a checkpoint entry. + + Returns: + Number of bytes in the contiguous backing allocation. + """ + if self.expert_gemm_backend != 'torch': + raise RuntimeError( + "prepare_torch_grouped_mm requires moe_expert_gemm_backend='torch'" + ) + if not hasattr(torch, '_grouped_mm'): + raise RuntimeError("this PyTorch build does not provide torch._grouped_mm") + if self.use_bias: + raise RuntimeError("torch grouped expert GEMM does not support bias") + if getattr(self, 'single_grouped_weight', False): + raise RuntimeError( + "torch grouped expert GEMM expects per-GEMM weight parameters, but " + "moe_single_grouped_weight=True makes TE hold one grouped weight tensor" + ) + + weights = [getattr(self, f'weight{index}') for index in range(self.num_gemms)] + if any(weight.requires_grad for weight in weights): + raise RuntimeError("torch grouped expert GEMM requires frozen base weights") + if any(weight.device.type != 'cuda' for weight in weights): + raise RuntimeError("torch grouped expert GEMM requires CUDA weights") + if any(weight.dtype != torch.bfloat16 for weight in weights): + raise RuntimeError("torch grouped expert GEMM requires BF16 weights") + if any(weight.shape != weights[0].shape for weight in weights[1:]): + raise RuntimeError("torch grouped expert GEMM requires uniform expert shapes") + + if not self._torch_grouped_weight_layout_is_current(): + grouped_weight = torch.stack([weight.detach() for weight in weights]).contiguous() + for index, weight in enumerate(weights): + weight.data = grouped_weight[index] + if '_torch_grouped_weight' in self._buffers: + self._buffers['_torch_grouped_weight'] = grouped_weight + else: + self.register_buffer('_torch_grouped_weight', grouped_weight, persistent=False) + + grouped_weight = self._buffers['_torch_grouped_weight'] + return grouped_weight.numel() * grouped_weight.element_size() + + @property + def torch_grouped_mm_prepared(self) -> bool: + """Return whether the torch grouped GEMM backing allocation is ready.""" + return self._torch_grouped_weight_layout_is_current() + + def _torch_grouped_mm_forward(self, x, m_splits): + """Run the frozen expert base branch with ``torch._grouped_mm``.""" + if not self._torch_grouped_weight_layout_is_current(): + self.prepare_torch_grouped_mm() + + x_2d = x.reshape(-1, x.shape[-1]) + input_rows = sum(m_splits) + if input_rows != x_2d.shape[0]: + raise RuntimeError( + f"expert splits sum to {input_rows}, but input has {x_2d.shape[0]} rows" + ) + + grouped_weight = self._buffers['_torch_grouped_weight'] + if input_rows == 0: + output = x_2d.matmul(grouped_weight[0].transpose(0, 1)) + else: + offsets = torch.tensor( + list(accumulate(m_splits)), device=x.device, dtype=torch.int32 + ) + output = torch._grouped_mm(x_2d, grouped_weight.transpose(1, 2), offs=offsets) + return output.reshape(*x.shape[:-1], output.shape[-1]), None + def finish_init(self, quantization_config: QuantizationConfig): """Post-init of quantization override""" if quantization_config is None: @@ -2647,6 +2753,9 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool: def forward(self, x, m_splits): """Forward.""" + if self.expert_gemm_backend == 'torch': + return self._torch_grouped_mm_forward(x, m_splits) + _is_first_microbatch = ( None if self.disable_parameter_transpose_cache else self.is_first_microbatch ) diff --git a/megatron/core/recompute.py b/megatron/core/recompute.py index bd0d1bcb3b2..fb68354781c 100644 --- a/megatron/core/recompute.py +++ b/megatron/core/recompute.py @@ -154,6 +154,13 @@ def chunk_runner(start: int, end: int, use_checkpoint: bool): if (start + layer_offset) in extract_layer_indices: intermediate_hidden_states.append(hidden_states) + if not hidden_states.requires_grad: + # Re-entrant checkpointing only attaches a grad_fn when some tensor input + # requires grad. With a frozen embedding (adapter-only training) every chunk + # output would otherwise carry no grad_fn and the adapters inside the chunks + # would receive no gradient. + hidden_states = hidden_states.detach().requires_grad_(True) + if self.config.recompute_method == 'uniform': # Uniformly divide the total number of layers and checkpoint # the input activation of each divided chunk. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 66cfaded213..7248627d102 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -843,6 +843,14 @@ class TransformerConfig(ModelParallelConfig): parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and ``add_bias_linear=True``.""" + moe_expert_gemm_backend: Literal['transformer_engine', 'torch'] = 'transformer_engine' + """Backend for grouped expert linear layers during training. + + ``torch`` uses ``torch._grouped_mm`` for frozen BF16, bias-free expert weights. It is intended + for parameter-efficient fine-tuning where only the adapter branch is trainable. The default + ``transformer_engine`` backend supports trainable expert weights and other precisions. + """ + moe_aux_loss_coeff: Union[float, List[float]] = 0.0 """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. If a list of load balancing types is provided for `moe_router_load_balancing_type`, @@ -1551,6 +1559,35 @@ def __post_init__(self): if self.num_moe_experts is not None and self.num_moe_experts <= 0: raise ValueError("num_moe_experts must be non-negative.") + if self.moe_expert_gemm_backend not in ('transformer_engine', 'torch'): + raise ValueError( + "moe_expert_gemm_backend must be 'transformer_engine' or 'torch', " + f"got {self.moe_expert_gemm_backend!r}" + ) + if self.moe_expert_gemm_backend == 'torch': + if self.num_moe_experts is None: + raise ValueError("moe_expert_gemm_backend='torch' requires num_moe_experts") + if not self.moe_grouped_gemm: + raise ValueError("moe_expert_gemm_backend='torch' requires moe_grouped_gemm=True") + if not self.bf16 or self.params_dtype != torch.bfloat16: + raise ValueError("moe_expert_gemm_backend='torch' requires BF16 parameters") + if self.add_bias_linear: + raise ValueError("moe_expert_gemm_backend='torch' does not support expert bias") + if self.fp8 or self.fp4: + raise ValueError( + "moe_expert_gemm_backend='torch' does not support FP8 or FP4 experts" + ) + if self.moe_single_grouped_weight: + raise ValueError( + "moe_expert_gemm_backend='torch' requires per-GEMM expert weight parameters " + "and is incompatible with moe_single_grouped_weight=True" + ) + if self.use_transformer_engine_op_fuser: + raise ValueError( + "moe_expert_gemm_backend='torch' is incompatible with " + "use_transformer_engine_op_fuser=True" + ) + if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None: self.moe_ffn_hidden_size = self.ffn_hidden_size warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.") diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 5d3c33264f4..6b4e98e288d 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -228,6 +228,46 @@ def run(block, hs, am): gb, gr = base_grads[name], rec_grads[name] assert torch.equal(gr, gb), f"Grad should be bitwise matched for {name}" + def test_full_recompute_with_frozen_input_trains_adapter(self): + """Full re-entrant recompute preserves gradients for adapter-only training.""" + block = self.get_hybrid_block( + Symbols.MLP, + add_bias_linear=False, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + ).cuda() + block.requires_grad_(False) + block.train() + + adapter = torch.nn.Sequential( + torch.nn.Linear(block.config.hidden_size, 4, bias=False), + torch.nn.Linear(4, block.config.hidden_size, bias=False), + ).cuda() + block.layers[0].add_module("adapter", adapter) + + def add_adapter(_module, _args, kwargs, output): + hidden_states, context = output + return hidden_states + adapter(kwargs["hidden_states"]), context + + block.layers[0].register_forward_hook(add_adapter, with_kwargs=True) + + sequence_length, micro_batch_size = 4, 1 + hidden_states = torch.randn( + sequence_length, micro_batch_size, block.config.hidden_size, device="cuda" + ) + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool, device="cuda" + ) + + assert not hidden_states.requires_grad + output = block(hidden_states, attention_mask=attention_mask) + output.float().square().mean().backward() + + for parameter in adapter.parameters(): + assert parameter.grad is not None + assert torch.count_nonzero(parameter.grad) > 0 + def test_layer_types(self): """ Make sure that the layer types specified with layer_pattern diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py new file mode 100644 index 00000000000..0267d118fde --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -0,0 +1,175 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_submodules, +) +from megatron.core.transformer.module import Float16Module +from megatron.core.transformer.moe.moe_layer import MoELayer +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import is_te_min_version +from tests.unit_tests.test_utilities import Utils + + +def _config(*, backend: str = "transformer_engine") -> TransformerConfig: + return TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_grouped_gemm=True, + moe_expert_gemm_backend=backend, + add_bias_linear=False, + gated_linear_unit=False, + activation_func=F.gelu, + bias_activation_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="sinkhorn", + moe_router_topk=1, + ) + + +@pytest.mark.parametrize( + "override,error", + [ + ({"num_moe_experts": None}, "requires num_moe_experts"), + ({"moe_grouped_gemm": False}, "requires moe_grouped_gemm=True"), + ({"bf16": False, "params_dtype": torch.float32}, "requires BF16 parameters"), + ({"add_bias_linear": True}, "does not support expert bias"), + ( + {"use_transformer_engine_op_fuser": True}, + "is incompatible with use_transformer_engine_op_fuser=True", + ), + ], +) +def test_torch_grouped_expert_gemm_config_validation(override, error): + kwargs = { + "num_layers": 1, + "hidden_size": 16, + "num_attention_heads": 4, + "num_moe_experts": 2, + "moe_grouped_gemm": True, + "moe_expert_gemm_backend": "torch", + "add_bias_linear": False, + "bf16": True, + "params_dtype": torch.bfloat16, + } + kwargs.update(override) + + with pytest.raises(ValueError, match=error): + TransformerConfig(**kwargs) + + +@pytest.mark.skipif( + not is_te_min_version("1.9.0.dev0") + or not torch.cuda.is_available() + or not hasattr(torch, "_grouped_mm") + or torch.cuda.get_device_capability()[0] < 10, + reason="torch grouped expert GEMM requires TE and a compatible CUDA device", +) +class TestFrozenTorchGroupedMLP: + def setup_method(self): + Utils.initialize_model_parallel(1, 1) + + def teardown_method(self): + Utils.destroy_model_parallel() + + @staticmethod + def _model(config: TransformerConfig) -> MoELayer: + model = MoELayer( + config, + get_gpt_layer_with_transformer_engine_submodules( + config.num_moe_experts, moe_grouped_gemm=True + ).mlp.submodules, + ) + return Float16Module(config, model).module.cuda() + + def test_forward_input_gradient_and_checkpoint_parity(self): + baseline = self._model(_config()) + torch_model = self._model(_config(backend="torch")) + torch_model.load_state_dict(baseline.state_dict()) + torch_linears = (torch_model.experts.linear_fc1, torch_model.experts.linear_fc2) + for linear in torch_linears: + for parameter in linear.parameters(): + parameter.requires_grad = False + state_keys_before = set(torch_model.state_dict()) + + baseline_input = torch.rand( + (32, 2, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + torch_input = baseline_input.detach().clone().requires_grad_(True) + baseline_output, _ = baseline(baseline_input) + torch_output, _ = torch_model(torch_input) + grad_output = torch.randn_like(baseline_output) + baseline_output.backward(grad_output) + torch_output.backward(grad_output) + + torch.testing.assert_close(torch_output, baseline_output, rtol=0.02, atol=0.02) + torch.testing.assert_close(torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02) + assert set(torch_model.state_dict()) == state_keys_before + assert not any("_torch_grouped_weight" in key for key in state_keys_before) + for linear in torch_linears: + assert linear.torch_grouped_mm_prepared + assert ( + linear._buffers["_torch_grouped_weight"].untyped_storage().data_ptr() + == linear.weight0.untyped_storage().data_ptr() + ) + + def test_no_tokens(self): + model = self._model(_config(backend="torch")) + for parameter in model.experts.parameters(): + parameter.requires_grad = False + hidden_states = torch.empty( + (0, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + + output, _ = model.experts( + hidden_states, + tokens_per_expert=torch.zeros(4, dtype=torch.int32, device="cuda"), + permuted_probs=torch.empty(0, dtype=torch.float32, device="cuda"), + ) + output.sum().backward() + + assert output.shape == (0, 16) + assert hidden_states.grad is not None + + def test_reprepares_after_middle_expert_weight_replacement(self): + model = self._model(_config(backend="torch")) + linear = model.experts.linear_fc1 + for parameter in linear.parameters(): + parameter.requires_grad = False + linear.prepare_torch_grouped_mm() + replacement = torch.randn_like(linear.weight1) + + linear.weight1.data = replacement + + assert not linear.torch_grouped_mm_prepared + linear.prepare_torch_grouped_mm() + assert linear.torch_grouped_mm_prepared + torch.testing.assert_close(linear.weight1, replacement) + + def test_zero_token_groups(self): + baseline = self._model(_config()) + torch_model = self._model(_config(backend="torch")) + torch_model.load_state_dict(baseline.state_dict()) + for parameter in torch_model.experts.parameters(): + parameter.requires_grad = False + + baseline_input = torch.rand( + (5, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + torch_input = baseline_input.detach().clone().requires_grad_(True) + tokens_per_expert = torch.tensor([2, 0, 3, 0], device="cuda") + permuted_probs = torch.rand(5, dtype=torch.float32, device="cuda") + baseline_output, _ = baseline.experts(baseline_input, tokens_per_expert, permuted_probs) + torch_output, _ = torch_model.experts(torch_input, tokens_per_expert, permuted_probs) + grad_output = torch.randn_like(baseline_output) + baseline_output.backward(grad_output) + torch_output.backward(grad_output) + + torch.testing.assert_close(torch_output, baseline_output, rtol=0.02, atol=0.02) + torch.testing.assert_close(torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02)