From 5916f5ce0eec8d8b76c684893cec0330927523a0 Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Sat, 22 Aug 2026 03:14:59 -0700 Subject: [PATCH 1/9] perf(moe): add frozen torch grouped GEMM backend Signed-off-by: Taufeeque (cherry picked from commit 193d5fe2bf3abf9f987015e3633faf331f94cf02) (cherry picked from commit 1f61e2407b7641c7865c2c87094f5f2308da3ada) Signed-off-by: Taufeeque --- .../core/extensions/transformer_engine.py | 106 ++++++++++++++ .../core/transformer/transformer_config.py | 32 ++++ .../moe/test_frozen_torch_grouped_mlp.py | 137 ++++++++++++++++++ 3 files changed, 275 insertions(+) create mode 100644 tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 7f7424e6c02..bbdaf39e6fe 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,107 @@ 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 + last_weight = getattr(self, f'weight{self.num_gemms - 1}') + if ( + grouped_weight.device != first_weight.device + or grouped_weight.dtype != first_weight.dtype + or grouped_weight.shape[1:] != first_weight.shape + or last_weight.device != first_weight.device + ): + return False + + storage_pointer = grouped_weight.untyped_storage().data_ptr() + expert_elements = first_weight.numel() + return ( + first_weight.untyped_storage().data_ptr() == storage_pointer + and first_weight.storage_offset() == 0 + and first_weight.stride() == grouped_weight[0].stride() + and last_weight.untyped_storage().data_ptr() == storage_pointer + and last_weight.storage_offset() == (self.num_gemms - 1) * expert_elements + and last_weight.stride() == grouped_weight[-1].stride() + ) + + @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 +2750,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/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 66cfaded213..6812ef1880a 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,30 @@ 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.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/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..714a657a571 --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -0,0 +1,137 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from copy import deepcopy + +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"), + ], +) +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): + config = deepcopy(_config(backend="torch")) + model = self._model(config) + 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 From ca36c35d4016a79d568baedf035a04ec0f6db9a6 Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Sat, 22 Aug 2026 03:20:58 -0700 Subject: [PATCH 2/9] test(moe): cover zero-token grouped expert splits Signed-off-by: Taufeeque (cherry picked from commit ecd9ca99d2346c01e5b967c41675aef09e85d90c) (cherry picked from commit dcf358be9ba3956a6f9a46a2859b8451f46e933a) Signed-off-by: Taufeeque --- .../moe/test_frozen_torch_grouped_mlp.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) 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 index 714a657a571..4d3ee6b93b8 100644 --- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -135,3 +135,31 @@ def test_no_tokens(self): assert output.shape == (0, 16) assert hidden_states.grad is not None + + 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 + ) From ede5504833fcada6d3eadd46b5a30b23cc0fcef1 Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Sat, 22 Aug 2026 03:23:34 -0700 Subject: [PATCH 3/9] fix(moe): validate every grouped expert weight view Signed-off-by: Taufeeque (cherry picked from commit f2a9b74de1dfd52407d8481bcd0fdb0a05b13503) (cherry picked from commit 356b80d44a91619b69f07b53e6ae278450ede3b3) Signed-off-by: Taufeeque --- .../core/extensions/transformer_engine.py | 23 +++++++++++-------- .../moe/test_frozen_torch_grouped_mlp.py | 15 ++++++++++++ 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index bbdaf39e6fe..80f6b8eb2c7 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2641,25 +2641,28 @@ def _torch_grouped_weight_layout_is_current(self) -> bool: return False first_weight = self.weight0 - last_weight = getattr(self, f'weight{self.num_gemms - 1}') if ( grouped_weight.device != first_weight.device or grouped_weight.dtype != first_weight.dtype or grouped_weight.shape[1:] != first_weight.shape - or last_weight.device != first_weight.device ): return False storage_pointer = grouped_weight.untyped_storage().data_ptr() expert_elements = first_weight.numel() - return ( - first_weight.untyped_storage().data_ptr() == storage_pointer - and first_weight.storage_offset() == 0 - and first_weight.stride() == grouped_weight[0].stride() - and last_weight.untyped_storage().data_ptr() == storage_pointer - and last_weight.storage_offset() == (self.num_gemms - 1) * expert_elements - and last_weight.stride() == grouped_weight[-1].stride() - ) + 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: 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 index 4d3ee6b93b8..c9970dda8ee 100644 --- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -136,6 +136,21 @@ def test_no_tokens(self): 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")) From 4947c1a8c9b5484cd3419c88e19b654bb3634694 Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Mon, 24 Aug 2026 15:43:35 -0700 Subject: [PATCH 4/9] fix(recompute): keep chunk outputs differentiable with a frozen embedding Re-entrant activation checkpointing only attaches a grad_fn when some tensor input requires grad. Under adapter-only training the embedding is frozen, so every checkpointed chunk output carried no grad_fn and the adapters inside the chunks received no gradient. Detach and re-enable grad on the block input before the chunk loop. The rest of the source commits added a MambaStack checkpointed forward, which upstream has since absorbed: MambaStack moved to megatron/core/models/hybrid/hybrid_block.py and routes full-granularity recompute through megatron/core/recompute.py::checkpointed_forward. Signed-off-by: Taufeeque (cherry picked from commit 91c8b4180) (cherry picked from commit abd235c8d) (cherry picked from commit 8d9ad06ba1bcfa4b5cf5358d92d5693b69c32bf4) Signed-off-by: Taufeeque --- megatron/core/recompute.py | 7 +++++++ 1 file changed, 7 insertions(+) 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. From b992f2828df1d71d66b3cf8651e367a57b1f9cb5 Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Tue, 25 Aug 2026 19:47:18 -0700 Subject: [PATCH 5/9] style(moe): format grouped GEMM tests Signed-off-by: Taufeeque --- .../transformer/moe/test_frozen_torch_grouped_mlp.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) 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 index c9970dda8ee..f8cc6ce09bd 100644 --- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -164,17 +164,11 @@ def test_zero_token_groups(self): 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 - ) + 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 - ) + torch.testing.assert_close(torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02) From 97f7d418ff58a9165f6082759ffa85f55d71a8bc Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Tue, 25 Aug 2026 19:52:26 -0700 Subject: [PATCH 6/9] build: record Megatron-LM fork base Signed-off-by: Taufeeque --- .fork-base.json | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 .fork-base.json 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" +} From 96ecb4b6d46a9647a756c03802c3823b7aacd922 Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Tue, 25 Aug 2026 20:05:35 -0700 Subject: [PATCH 7/9] test(recompute): cover frozen-input adapter gradients Signed-off-by: Taufeeque --- tests/unit_tests/ssm/test_hybrid_block.py | 40 +++++++++++++++++++++++ 1 file changed, 40 insertions(+) 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 From dcee1538e7fdd8c3cc93402b3cfd64dc5557823c Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Tue, 25 Aug 2026 20:12:58 -0700 Subject: [PATCH 8/9] test(moe): simplify no-token setup Signed-off-by: Taufeeque --- .../transformer/moe/test_frozen_torch_grouped_mlp.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 index f8cc6ce09bd..a8a3ce91a8a 100644 --- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -1,7 +1,5 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -from copy import deepcopy - import pytest import torch import torch.nn.functional as F @@ -118,8 +116,7 @@ def test_forward_input_gradient_and_checkpoint_parity(self): ) def test_no_tokens(self): - config = deepcopy(_config(backend="torch")) - model = self._model(config) + model = self._model(_config(backend="torch")) for parameter in model.experts.parameters(): parameter.requires_grad = False hidden_states = torch.empty( From 8447c0d3eac0a71c7367a33da0d4e85c0708a46b Mon Sep 17 00:00:00 2001 From: Taufeeque Date: Tue, 25 Aug 2026 20:18:44 -0700 Subject: [PATCH 9/9] fix(moe): reject TE op fuser with torch backend Signed-off-by: Taufeeque --- megatron/core/transformer/transformer_config.py | 5 +++++ .../transformer/moe/test_frozen_torch_grouped_mlp.py | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 6812ef1880a..7248627d102 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1582,6 +1582,11 @@ def __post_init__(self): "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 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 index a8a3ce91a8a..0267d118fde 100644 --- a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -40,6 +40,10 @@ def _config(*, backend: str = "transformer_engine") -> TransformerConfig: ({"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):