diff --git a/modelopt/torch/export/plugins/megatron_importer.py b/modelopt/torch/export/plugins/megatron_importer.py index e38bf7e516c..8bb49a77271 100644 --- a/modelopt/torch/export/plugins/megatron_importer.py +++ b/modelopt/torch/export/plugins/megatron_importer.py @@ -51,6 +51,30 @@ has_mcore = True +class _MambaConv1dCompat(torch.nn.Module): + """Expose direct Mamba conv params through the legacy Conv1d state_dict keys.""" + + def __init__(self, mixer): + super().__init__() + self.weight = mixer.conv1d_weight + self.bias = mixer.conv1d_bias + + +def _get_mamba_conv1d(mixer): + conv1d = getattr(mixer, "conv1d", None) + if conv1d is not None: + return conv1d + + if hasattr(mixer, "conv1d_weight") and hasattr(mixer, "conv1d_bias"): + # Megatron-LM PR #4899 / commit 35992ba changed MambaMixer fields from + # `conv1d` to direct `conv1d_weight` and `conv1d_bias` parameters. + return _MambaConv1dCompat(mixer) + + raise AttributeError( + "MambaMixer does not have `conv1d` or `conv1d_weight`/`conv1d_bias` fields." + ) + + class GPTModelImporter: """Megatron Core GPTModel HuggingFace Importer. @@ -561,7 +585,7 @@ def _import_mamba_layer(self, layer, layer_id, layer_pbar): self.rules["A_log"](layer.mixer.A_log, layer_id) self.rules["D"](layer.mixer.D, layer_id) self.rules["dt_bias"](layer.mixer.dt_bias, layer_id) - self.rules["conv1d"](layer.mixer.conv1d, layer_id) + self.rules["conv1d"](_get_mamba_conv1d(layer.mixer), layer_id) self.rules["in_proj"](layer.mixer.in_proj, layer_id) self.rules["out_proj"](layer.mixer.out_proj, layer_id) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index cabb0d77580..070a4478838 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -57,7 +57,7 @@ get_safetensor, save_safetensors_by_layer_index, ) -from .plugins.megatron_importer import GPTModelImporter +from .plugins.megatron_importer import GPTModelImporter, _get_mamba_conv1d from .quant_utils import ( get_activation_scaling_factor, get_kv_cache_dtype, @@ -664,7 +664,7 @@ def _get_mamba_layer_state_dict(self, layer, layer_id): self.rules["D"](layer.mixer.D, layer_id) self.rules["dt_bias"](layer.mixer.dt_bias, layer_id) - self.rules["conv1d"](layer.mixer.conv1d, layer_id) + self.rules["conv1d"](_get_mamba_conv1d(layer.mixer), layer_id) self.rules["in_proj"](layer.mixer.in_proj, layer_id) self.rules["out_proj"](layer.mixer.out_proj, layer_id) diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index 5f49615d638..593c0f4f2e5 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -216,6 +216,55 @@ def create_tiny_nemotron_dir( return nemotron_dir +##### NEMOTRON-H (Mamba + Attention + MoE/MLP hybrid) ##### +def get_tiny_nemotron_h(**config_kwargs) -> PreTrainedModel: + set_seed(SEED) + + # Lazy import — NemotronHConfig only exists in newer transformers builds, and this + # module is imported broadly (including by the min-transformers CI job). + from transformers import NemotronHConfig + + # Tiny NemotronH hybrid. hybrid_override_pattern letters: M=Mamba, E=MoE/FFN, *=Attention. + # "ME*E" matches the NemotronH default and exercises Mamba + MoE + attention layers. + kwargs = { + "dtype": torch.bfloat16, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 4, + "hybrid_override_pattern": "ME*E", + "num_attention_heads": 8, + "num_key_value_heads": 4, + "head_dim": 8, + "mamba_num_heads": 8, + "mamba_head_dim": 16, + "ssm_state_size": 32, + "n_groups": 2, + "conv_kernel": 4, + # MoE + "n_routed_experts": 8, + "num_experts_per_tok": 2, + "moe_intermediate_size": 64, + "n_shared_experts": 1, + "moe_shared_expert_intermediate_size": 32, + "vocab_size": 32, + "max_position_embeddings": 32, + } + kwargs.update(**config_kwargs) + return AutoModelForCausalLM.from_config(NemotronHConfig(**kwargs)) + + +def create_tiny_nemotron_h_dir( + tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs +) -> Path: + nemotron_h_dir = Path(tmp_path) / "tiny_nemotron_h" + if with_tokenizer: + tokenizer = get_tiny_tokenizer() + tokenizer.save_pretrained(nemotron_h_dir) + config_kwargs["vocab_size"] = tokenizer.vocab_size + get_tiny_nemotron_h(**config_kwargs).save_pretrained(nemotron_h_dir) + return nemotron_h_dir + + ##### DeepSeek V3 ##### def get_tiny_deepseek_v3(**config_kwargs) -> PreTrainedModel: set_seed(SEED) diff --git a/tests/gpu_megatron/torch/export/test_megatron_importer.py b/tests/gpu_megatron/torch/export/test_megatron_importer.py new file mode 100644 index 00000000000..8776f50cf11 --- /dev/null +++ b/tests/gpu_megatron/torch/export/test_megatron_importer.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from functools import partial + +import torch +import transformers +from _test_utils.torch.megatron.models import get_mcore_mamba_hybrid_model +from _test_utils.torch.transformers_models import create_tiny_nemotron_h_dir +from safetensors import safe_open + +from modelopt.torch.export import export_mcore_gpt_to_hf, import_mcore_gpt_from_hf +from modelopt.torch.export.plugins.megatron_importer import _get_mamba_conv1d + + +class _Mixer: + pass + + +def test_get_mamba_conv1d_returns_legacy_module(): + mixer = _Mixer() + mixer.conv1d = torch.nn.Conv1d(4, 4, 3) + + assert _get_mamba_conv1d(mixer) is mixer.conv1d + + +def test_get_mamba_conv1d_wraps_direct_params(): + mixer = _Mixer() + mixer.conv1d_weight = torch.nn.Parameter(torch.zeros(4, 1, 3)) + mixer.conv1d_bias = torch.nn.Parameter(torch.zeros(4)) + + conv1d = _get_mamba_conv1d(mixer) + new_weight = torch.ones_like(mixer.conv1d_weight) + new_bias = torch.ones_like(mixer.conv1d_bias) + conv1d.load_state_dict({"weight": new_weight, "bias": new_bias}) + + assert set(conv1d.state_dict()) == {"weight", "bias"} + assert conv1d.weight is mixer.conv1d_weight + assert conv1d.bias is mixer.conv1d_bias + torch.testing.assert_close(mixer.conv1d_weight, new_weight) + torch.testing.assert_close(mixer.conv1d_bias, new_bias) + + +# NemotronH-style MoE + Mamba hybrid: exercise import/export of a model with Mamba +# (conv1d/in_proj/out_proj), attention, MLP, and MoE expert layers. + + +def _build_mcore_nemotron_h(config, size, initialize=True): + return get_mcore_mamba_hybrid_model( + tensor_model_parallel_size=size, + pipeline_model_parallel_size=1, + initialize_megatron=initialize, + num_layers=config.num_hidden_layers, + hybrid_override_pattern=config.hybrid_override_pattern, + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_query_groups=config.num_key_value_heads, + ffn_hidden_size=config.intermediate_size, + max_sequence_length=config.max_position_embeddings, + vocab_size=config.vocab_size, + mamba_state_dim=config.ssm_state_size, + mamba_num_heads=config.mamba_num_heads, + mamba_head_dim=config.mamba_head_dim, + mamba_num_groups=config.n_groups, + num_moe_experts=config.n_routed_experts, + moe_ffn_hidden_size=config.moe_intermediate_size, + moe_shared_expert_intermediate_size=config.moe_shared_expert_intermediate_size, + ).cuda() + + +def _test_nemotron_h_export_import(tmp_path, model_dir, rank, size): + config = transformers.AutoConfig.from_pretrained(model_dir) + + # Export a Mamba + MoE hybrid mcore model to HF safetensors. + model = _build_mcore_nemotron_h(config, size) + export_dir = tmp_path / "export" + export_mcore_gpt_to_hf(model, str(model_dir), dtype=torch.bfloat16, export_dir=str(export_dir)) + + if rank == 0: + keys = [] + for sf in export_dir.glob("*.safetensors"): + with safe_open(str(sf), framework="pt", device="cpu") as f: + keys.extend(f.keys()) + # Mamba mixer (conv1d is the layer fixed by _get_mamba_conv1d), plus MoE experts. + assert any("mixer.conv1d" in k for k in keys), "mamba conv1d weights missing from export" + assert any("mixer.in_proj" in k for k in keys), "mamba in_proj weights missing from export" + assert any("mixer.experts" in k for k in keys), "moe expert weights missing from export" + + # Import the exported checkpoint back into a fresh mcore model + # (megatron is already initialized from the first build). + imported = _build_mcore_nemotron_h(config, size, initialize=False) + import_mcore_gpt_from_hf(imported, str(export_dir)) + + +def test_nemotron_h_export_import(dist_workers_size_1, tmp_path): + model_dir = create_tiny_nemotron_h_dir(tmp_path) + dist_workers_size_1.run(partial(_test_nemotron_h_export_import, tmp_path, model_dir))