nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/export/plugins/mcore_gptoss.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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.
+
+"""Custom mapping from GPT-OSS Hugging Face models to Megatron Core models."""
+
+from .mcore_custom import (
+    COL_TP,
+    PACK_EP,
+    REPLICATE,
+    ROW_TP,
+    CustomModuleMapping,
+    NameRemapping,
+    PackNameRemappingGPT,
+    QKVMerging,
+    QKVSlicing,
+    UnpackNameRemappingGPT,
+)
+
+gptoss_causal_lm_export: dict[str, CustomModuleMapping | bool] = {
+    "word_embeddings": NameRemapping("model.embed_tokens."),
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
+    "linear_qkv": QKVSlicing("model.layers.{}.self_attn."),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    "softmax_offset": NameRemapping("model.layers.{}.self_attn.sinks"),
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
+    "use_packed_local_experts": True,
+    "local_experts.linear_fc1": PackNameRemappingGPT(
+        "model.layers.{}.mlp.experts.gate_up_proj",
+        {"layer_type": "linear_fc1"},
+    ),
+    "local_experts.linear_fc2": PackNameRemappingGPT(
+        "model.layers.{}.mlp.experts.down_proj",
+        {"layer_type": "linear_fc2"},
+    ),
+    "router": NameRemapping("model.layers.{}.mlp.router."),
+    "final_layernorm": NameRemapping("model.norm."),
+    "output_layer": NameRemapping("lm_head."),
+}
+
+gptoss_causal_lm_import: dict[str, CustomModuleMapping | bool] = {
+    "word_embeddings": NameRemapping("model.embed_tokens.", COL_TP),
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm.", REPLICATE),
+    "linear_qkv": QKVMerging("model.layers.{}.self_attn.", COL_TP),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj.", ROW_TP),
+    "softmax_offset": NameRemapping("model.layers.{}.self_attn.sinks", COL_TP),
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm.", REPLICATE),
+    "router": NameRemapping("model.layers.{}.mlp.router.", REPLICATE),
+    "use_packed_local_experts": True,
+    "local_experts.linear_fc1_ep": UnpackNameRemappingGPT(
+        "model.layers.{}.mlp.experts.gate_up_proj",
+        PACK_EP | {"layer_type": "linear_fc1"},
+    ),
+    "local_experts.linear_fc2_ep": UnpackNameRemappingGPT(
+        "model.layers.{}.mlp.experts.down_proj",
+        PACK_EP | {"layer_type": "linear_fc2"},
+    ),
+    "final_layernorm": NameRemapping("model.norm.", REPLICATE),
+    "output_layer": NameRemapping("lm_head.", COL_TP),
+}

modelopt/torch/export/plugins/mcore_llama.py

@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+
+"""Custom mapping from Llama Hugging Face models to Megatron Core models."""
+
+from .mcore_custom import (
+    COL_TP,
+    PACK_COL_ETP,
+    PACK_EP,
+    PACK_ROW_ETP,
+    REPLICATE,
+    ROW_TP,
+    CustomModuleMapping,
+    GatedMLPMerging,
+    GatedMLPSlicing,
+    NameRemapping,
+    PackNameRemapping,
+    QKVMerging,
+    QKVSlicing,
+    UnpackNameRemapping,
+)
+
+llama_causal_lm_export: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("model.embed_tokens."),
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
+    "linear_qkv": QKVSlicing("model.layers.{}.self_attn."),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
+    "linear_fc1": GatedMLPSlicing("model.layers.{}.mlp."),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
+    "final_layernorm": NameRemapping("model.norm."),
+    "output_layer": NameRemapping("lm_head."),
+}
+
+llama4_causal_lm_export: dict[str, CustomModuleMapping | bool] = {
+    "word_embeddings": NameRemapping("language_model.model.embed_tokens."),
+    "input_layernorm": NameRemapping("language_model.model.layers.{}.input_layernorm."),
+    # self_attn
+    "linear_qkv": QKVSlicing("language_model.model.layers.{}.self_attn."),
+    "linear_proj": NameRemapping("language_model.model.layers.{}.self_attn.o_proj."),
+    # mlp
+    "pre_mlp_layernorm": NameRemapping("language_model.model.layers.{}.post_attention_layernorm."),
+    "shared_experts.linear_fc1": GatedMLPSlicing(
+        "language_model.model.layers.{}.feed_forward.shared_expert.",
+    ),
+    "shared_experts.linear_fc2": NameRemapping(
+        "language_model.model.layers.{}.feed_forward.shared_expert.down_proj.",
+    ),
+    # moe_layer
+    "router": NameRemapping("language_model.model.layers.{}.feed_forward.router."),
+    "use_packed_local_experts": True,
+    "local_experts.linear_fc1": PackNameRemapping(
+        "language_model.model.layers.{}.feed_forward.experts.gate_up_proj",
+        {"layer_type": "linear_fc1"},
+    ),
+    "local_experts.linear_fc2": PackNameRemapping(
+        "language_model.model.layers.{}.feed_forward.experts.down_proj",
+        {"layer_type": "linear_fc2"},
+    ),
+    "final_layernorm": NameRemapping("language_model.model.norm."),
+    "output_layer": NameRemapping("language_model.lm_head."),
+}
+
+medusa_llama_causal_lm_export: dict[str, CustomModuleMapping] = {
+    # MedusaForCausalLM support
+    "lm_head": NameRemapping(
+        "medusa_heads.{}.1."
+    ),  # TODO: lm_head is hardcoded to .1 as currently only support using 1 layer in medusa head
+    # needs a fix
+    "linear": NameRemapping("medusa_heads.{}.{}.linear."),
+}
+
+eagle_llama_causal_lm_export: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("embed_tokens."),
+    "enorm": NameRemapping("enorm."),
+    "hnorm": NameRemapping("hnorm."),
+    "fc": NameRemapping("fc."),
+    "input_layernorm": NameRemapping("layers.{}.input_layernorm."),
+    "linear_qkv": QKVSlicing("layers.{}.self_attn."),
+    "linear_proj": NameRemapping("layers.{}.self_attn.o_proj."),
+    "pre_mlp_layernorm": NameRemapping("layers.{}.post_attention_layernorm."),
+    "linear_fc1": GatedMLPSlicing("layers.{}.mlp."),
+    "linear_fc2": NameRemapping("layers.{}.mlp.down_proj."),
+    "final_layernorm": NameRemapping("norm."),
+    "d2t": NameRemapping("d2t"),
+    "output_layer": NameRemapping("lm_head."),
+}
+
+eagle3_llama_causal_lm_export: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("embed_tokens."),
+    "enorm": NameRemapping("midlayer.input_layernorm."),
+    "fc": NameRemapping("fc."),
+    "input_layernorm": NameRemapping("midlayer.hidden_norm."),
+    "linear_qkv": QKVSlicing("midlayer.self_attn."),
+    "linear_proj": NameRemapping("midlayer.self_attn.o_proj."),
+    "pre_mlp_layernorm": NameRemapping("midlayer.post_attention_layernorm."),
+    "linear_fc1": GatedMLPSlicing("midlayer.mlp."),
+    "linear_fc2": NameRemapping("midlayer.mlp.down_proj."),
+    "final_layernorm": NameRemapping("norm."),
+    "d2t": NameRemapping("d2t"),
+    "output_layer": NameRemapping("lm_head."),
+}
+
+
+llama_causal_lm_import: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("model.embed_tokens.", COL_TP),
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm.", REPLICATE),
+    "linear_qkv": QKVMerging("model.layers.{}.self_attn.", COL_TP),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj.", ROW_TP),
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm.", REPLICATE),
+    "linear_fc1": GatedMLPMerging("model.layers.{}.mlp.", COL_TP),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj.", ROW_TP),
+    "final_layernorm": NameRemapping("model.norm.", REPLICATE),
+    "output_layer": NameRemapping("lm_head.", COL_TP),
+}
+
+llama4_causal_lm_import: dict[str, CustomModuleMapping | bool] = {
+    "word_embeddings": NameRemapping("language_model.model.embed_tokens.", COL_TP),
+    "input_layernorm": NameRemapping("language_model.model.layers.{}.input_layernorm.", REPLICATE),
+    "linear_qkv": QKVMerging("language_model.model.layers.{}.self_attn.", COL_TP),
+    "linear_proj": NameRemapping("language_model.model.layers.{}.self_attn.o_proj.", ROW_TP),
+    "pre_mlp_layernorm": NameRemapping(
+        "language_model.model.layers.{}.post_attention_layernorm.", REPLICATE
+    ),
+    "shared_experts.linear_fc1": GatedMLPMerging(
+        "language_model.model.layers.{}.feed_forward.shared_expert.", COL_TP
+    ),
+    "shared_experts.linear_fc2": NameRemapping(
+        "language_model.model.layers.{}.feed_forward.shared_expert.down_proj.", ROW_TP
+    ),
+    "router": NameRemapping("language_model.model.layers.{}.feed_forward.router.", REPLICATE),
+    "use_packed_local_experts": True,
+    "local_experts.linear_fc1_etp": UnpackNameRemapping(
+        "language_model.model.layers.{}.feed_forward.experts.gate_up_proj",
+        PACK_COL_ETP | {"layer_type": "linear_fc1"},
+    ),
+    "local_experts.linear_fc2_etp": UnpackNameRemapping(
+        "language_model.model.layers.{}.feed_forward.experts.down_proj",
+        PACK_ROW_ETP | {"layer_type": "linear_fc2"},
+    ),
+    "local_experts.linear_fc1_ep": UnpackNameRemapping(
+        "language_model.model.layers.{}.feed_forward.experts.gate_up_proj",
+        PACK_EP | {"layer_type": "linear_fc1"},
+    ),
+    "local_experts.linear_fc2_ep": UnpackNameRemapping(
+        "language_model.model.layers.{}.feed_forward.experts.down_proj",
+        PACK_EP | {"layer_type": "linear_fc2"},
+    ),
+    "final_layernorm": NameRemapping("language_model.model.norm.", REPLICATE),
+    "output_layer": NameRemapping("language_model.lm_head.", COL_TP),
+}

modelopt/torch/export/plugins/mcore_nemotron.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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.
+
+
+"""Custom mapping from Nemotron Hugging Face models to Megatron Core models."""
+
+from .mcore_custom import (
+    COL_TP,
+    REPLICATE,
+    ROW_TP,
+    CustomModuleMapping,
+    NameRemapping,
+    QKVMerging,
+    QKVSlicing,
+)
+
+# Example on adding a new CausalLM.
+nemotron_causal_lm_export: dict[str, CustomModuleMapping] = {
+    # NemotronForCausalLM is using square-relu where no gated handle is needed.
+    "word_embeddings": NameRemapping("model.embed_tokens."),
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
+    "linear_qkv": QKVSlicing("model.layers.{}.self_attn."),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
+    # NemotronForCausalLM is using square-relu where no gated handle is needed.
+    "linear_fc1": NameRemapping("model.layers.{}.mlp.up_proj."),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
+    "final_layernorm": NameRemapping("model.norm."),
+    "output_layer": NameRemapping("lm_head."),
+}
+
+
+nemotron_h_causal_lm_import: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("backbone.embeddings.", COL_TP),
+    "final_norm": NameRemapping("backbone.norm_f.", REPLICATE),
+    "output_layer": NameRemapping("lm_head.", COL_TP),
+    # Mamba
+    "norm": NameRemapping("backbone.layers.{}.norm.", REPLICATE),
+    "mixer_norm": NameRemapping("backbone.layers.{}.mixer.norm.", REPLICATE),
+    "A_log": NameRemapping("backbone.layers.{}.mixer.A_log", REPLICATE),
+    "D": NameRemapping("backbone.layers.{}.mixer.D", REPLICATE),
+    "dt_bias": NameRemapping("backbone.layers.{}.mixer.dt_bias", REPLICATE),
+    "conv1d": NameRemapping("backbone.layers.{}.mixer.conv1d.", REPLICATE),
+    "in_proj": NameRemapping("backbone.layers.{}.mixer.in_proj.", COL_TP),
+    "out_proj": NameRemapping("backbone.layers.{}.mixer.out_proj.", ROW_TP),
+    # Attention
+    "input_layernorm": NameRemapping("backbone.layers.{}.norm.", REPLICATE),
+    "linear_qkv": QKVMerging("backbone.layers.{}.mixer.", COL_TP),
+    "linear_proj": NameRemapping("backbone.layers.{}.mixer.o_proj.", ROW_TP),
+    # MLP
+    "pre_mlp_layernorm": NameRemapping("backbone.layers.{}.norm.", REPLICATE),
+    "linear_fc1": NameRemapping("backbone.layers.{}.mixer.up_proj.", COL_TP),
+    "linear_fc2": NameRemapping("backbone.layers.{}.mixer.down_proj.", ROW_TP),
+}
+
+
+nemotron_h_causal_lm_export: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("backbone.embeddings."),
+    "final_norm": NameRemapping("backbone.norm_f."),
+    "output_layer": NameRemapping("lm_head."),
+    # Mamba
+    "norm": NameRemapping("backbone.layers.{}.norm."),
+    "mixer_norm": NameRemapping("backbone.layers.{}.mixer.norm."),
+    "A_log": NameRemapping("backbone.layers.{}.mixer.A_log"),
+    "D": NameRemapping("backbone.layers.{}.mixer.D"),
+    "dt_bias": NameRemapping("backbone.layers.{}.mixer.dt_bias"),
+    "conv1d": NameRemapping("backbone.layers.{}.mixer.conv1d."),
+    "in_proj": NameRemapping("backbone.layers.{}.mixer.in_proj."),
+    "out_proj": NameRemapping("backbone.layers.{}.mixer.out_proj."),
+    # Attention
+    "input_layernorm": NameRemapping("backbone.layers.{}.norm."),
+    "linear_qkv": QKVSlicing("backbone.layers.{}.mixer."),
+    "linear_proj": NameRemapping("backbone.layers.{}.mixer.o_proj."),
+    # MLP
+    "pre_mlp_layernorm": NameRemapping("backbone.layers.{}.norm."),
+    "linear_fc1": NameRemapping("backbone.layers.{}.mixer.up_proj."),
+    "linear_fc2": NameRemapping("backbone.layers.{}.mixer.down_proj."),
+}

modelopt/torch/export/plugins/mcore_qwen.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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.
+
+"""Custom mapping from Qwen Hugging Face models to Megatron Core models."""
+
+from .mcore_custom import (
+    COL_ETP,
+    COL_TP,
+    REPLICATE,
+    ROW_ETP,
+    ROW_TP,
+    CustomModuleMapping,
+    GatedMLPMerging,
+    GatedMLPSlicing,
+    NameRemapping,
+    QKVMerging,
+    QKVSlicing,
+)
+
+qwen3_causal_lm_import: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("model.embed_tokens.", COL_TP),
+    "final_layernorm": NameRemapping("model.norm.", REPLICATE),
+    "output_layer": NameRemapping("lm_head.", COL_TP),
+    # Attention
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm.", REPLICATE),
+    "linear_qkv": QKVMerging("model.layers.{}.self_attn.", COL_TP),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj.", ROW_TP),
+    "q_layernorm": NameRemapping("model.layers.{}.self_attn.q_norm.", REPLICATE),
+    "k_layernorm": NameRemapping("model.layers.{}.self_attn.k_norm.", REPLICATE),
+    # MLP
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm.", REPLICATE),
+    "linear_fc1": GatedMLPMerging("model.layers.{}.mlp.", COL_TP),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj.", ROW_TP),
+    # MoE
+    "router": NameRemapping("model.layers.{}.mlp.gate.", REPLICATE),
+    "local_experts.linear_fc1": GatedMLPMerging("model.layers.{}.mlp.experts.{}.", COL_ETP),
+    "local_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.experts.{}.down_proj.", ROW_ETP),
+}
+
+
+qwen3_causal_lm_export: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("model.embed_tokens."),
+    "final_layernorm": NameRemapping("model.norm."),
+    "output_layer": NameRemapping("lm_head."),
+    # Attention
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
+    "linear_qkv": QKVSlicing("model.layers.{}.self_attn."),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    "q_layernorm": NameRemapping("model.layers.{}.self_attn.q_norm."),
+    "k_layernorm": NameRemapping("model.layers.{}.self_attn.k_norm."),
+    # MLP
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
+    "linear_fc1": GatedMLPSlicing("model.layers.{}.mlp."),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
+    # MoE
+    "router": NameRemapping("model.layers.{}.mlp.gate."),
+    "local_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.experts.{}."),
+    "local_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.experts.{}.down_proj."),
+}
+
+qwen25_causal_lm_import: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("model.embed_tokens.", COL_TP),
+    "final_layernorm": NameRemapping("model.norm.", REPLICATE),
+    "output_layer": NameRemapping("lm_head.", COL_TP),
+    # Attention
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm.", REPLICATE),
+    "linear_qkv": QKVMerging("model.layers.{}.self_attn.", COL_TP),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj.", ROW_TP),
+    # MLP
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm.", REPLICATE),
+    "linear_fc1": GatedMLPMerging("model.layers.{}.mlp.", COL_TP),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj.", ROW_TP),
+}
+
+qwen25_causal_lm_export: dict[str, CustomModuleMapping] = {
+    "word_embeddings": NameRemapping("model.embed_tokens."),
+    "final_layernorm": NameRemapping("model.norm."),
+    "output_layer": NameRemapping("lm_head."),
+    # Attention
+    "input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
+    "linear_qkv": QKVSlicing("model.layers.{}.self_attn."),
+    "linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
+    # MLP
+    "pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
+    "linear_fc1": GatedMLPSlicing("model.layers.{}.mlp."),
+    "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
+}