Skip to content

Support EP mcore import for TE Spec and Fix mamba moe config #1342

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jenchen13 merged 8 commits into from
Apr 28, 2026
Merged

Support EP mcore import for TE Spec and Fix mamba moe config #1342

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
4fa630a
support EP import for TE and fix mamba moe config
jenchen13 Apr 24, 2026
948e4de
lint
jenchen13 Apr 24, 2026
67e021d
fix hybrid model export
jenchen13 Apr 24, 2026
27a4eef
lint
jenchen13 Apr 27, 2026
67b8848
getter/setter for QuantCfg
jenchen13 Apr 27, 2026
0b3ca94
remove setter
jenchen13 Apr 27, 2026
a7dc300
revert getter
jenchen13 Apr 27, 2026
09d3950
Merge branch 'main' into jennifchen/te_ep_import
jenchen13 Apr 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions modelopt/torch/export/plugins/megatron_importer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
has_mcore = False
with import_plugin("megatron"):
from megatron.core.parallel_state import (
get_expert_model_parallel_rank,
get_expert_tensor_parallel_world_size,
get_tensor_model_parallel_world_size,
)
Expand Down Expand Up @@ -294,9 +295,13 @@ def _grouped_mlp_merging(
assert module.num_gemms == num_local_experts, (
"num_gemms must be equal to num_local_experts in TEGroupedMLP"
)
for expert_id in range(init_expert_id, init_expert_id + num_local_experts):
tensor = self._get_safetensor(prefix.format(expert_id) + ".weight")
state_dict[f"weight{expert_id}"] = tensor
# init_expert_id is the global index of this rank's first local expert.
# TEGroupedMLP stores weights as weight0..weight{num_local-1} locally, so we
# map global expert_id -> local slot (expert_id - init_expert_id).
for local_id in range(num_local_experts):
global_expert_id = init_expert_id + local_id
tensor = self._get_safetensor(prefix.format(global_expert_id) + ".weight")
state_dict[f"weight{local_id}"] = tensor
# TODO handle weight_scale

module.load_state_dict(state_dict)
Expand Down Expand Up @@ -653,10 +658,13 @@ def _import_transformer_layer(self, layer, layer_id, layer_pbar, is_mtp: bool =
layer_pbar.set_description("Importing MoE grouped local experts")
num_local_experts = experts.num_local_experts
num_global_experts = experts.config.num_moe_experts
assert num_local_experts == num_global_experts, (
"num_local_experts must be equal to num_global_experts during MoE import"
assert num_global_experts % num_local_experts == 0, (
"num_global_experts must be divisible by num_local_experts "
"during MoE import"
)
init_index = 0
# Each EP rank owns a contiguous slice of global experts:
# [ep_rank * num_local_experts, (ep_rank + 1) * num_local_experts).
init_index = get_expert_model_parallel_rank() * num_local_experts
Comment on lines +661 to +667

@coderabbitai coderabbitai Bot Apr 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

cat -n modelopt/torch/export/plugins/megatron_importer.py | sed -n '650,680p'

Repository: NVIDIA/Model-Optimizer

Length of output: 2049

🏁 Script executed:

# Check imports and function definitions
cat -n modelopt/torch/export/plugins/megatron_importer.py | head -50

Repository: NVIDIA/Model-Optimizer

Length of output: 2174

🏁 Script executed:

# Search for where these functions are defined or imported
rg "get_expert_model_parallel_rank|get_expert_model_parallel_world_size" modelopt/torch/export/plugins/megatron_importer.py -B 2 -A 2

Repository: NVIDIA/Model-Optimizer

Length of output: 650

Replace assert with explicit EP-topology validation.

Line 662 uses assert for runtime validation, which Python removes when executed with optimization flags (-O). Additionally, divisibility alone (num_global_experts % num_local_experts == 0) does not guarantee correct expert mapping. The code distributes experts by EP rank using the formula [ep_rank * num_local_experts, (ep_rank + 1) * num_local_experts), which requires num_global_experts == num_local_experts * ep_size. Without this constraint, mismatches between expert count and EP topology can silently produce incorrect indexing.

Suggested fix 
-                        assert num_global_experts % num_local_experts == 0, (
-                            "num_global_experts must be divisible by num_local_experts "
-                            "during MoE import"
-                        )
-                        # Each EP rank owns a contiguous slice of global experts:
-                        # [ep_rank * num_local_experts, (ep_rank + 1) * num_local_experts).
-                        init_index = get_expert_model_parallel_rank() * num_local_experts
+                        ep_rank = get_expert_model_parallel_rank()
+                        ep_size = get_expert_model_parallel_world_size()
+                        if num_global_experts != num_local_experts * ep_size:
+                            raise ValueError(
+                                "Expected num_global_experts == num_local_experts * ep_size "
+                                f"for TEGroupedMLP import, got {num_global_experts=}, "
+                                f"{num_local_experts=}, {ep_size=}."
+                            )
+                        # Each EP rank owns a contiguous slice of global experts:
+                        # [ep_rank * num_local_experts, (ep_rank + 1) * num_local_experts).
+                        init_index = ep_rank * num_local_experts
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/export/plugins/megatron_importer.py` around lines 662 - 668,
Replace the fragile assert-based check with an explicit validation that raises a
clear exception: verify both that num_global_experts is divisible by
num_local_experts and that num_global_experts == num_local_experts * ep_size
(where ep_size is the EP topology size used by
get_expert_model_parallel_rank()); if these conditions fail, raise a ValueError
with a descriptive message mentioning num_global_experts, num_local_experts and
ep_size so the expert slice computation (init_index =
get_expert_model_parallel_rank() * num_local_experts) cannot proceed with an
invalid topology.


self.rules["experts.linear_fc1"](
experts.linear_fc1,
Expand Down
7 changes: 6 additions & 1 deletion modelopt/torch/export/unified_export_megatron.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,11 @@
with import_plugin("megatron"):
from megatron.core.models.gpt import GPTModel
from megatron.core.models.mamba import MambaModel

try:
from megatron.core.models.hybrid.hybrid_model import HybridModel
except ImportError:
HybridModel = MambaModel
from megatron.core.models.multimodal.llava_model import LLaVAModel
from megatron.core.parallel_state import (
get_pipeline_model_parallel_rank,
Expand Down Expand Up @@ -121,7 +126,7 @@ def __init__(
moe_router_dtype: str | None = None,
):
"""Create a GPTModel exporter instance."""
if not isinstance(model, (GPTModel, MambaModel, LLaVAModel)):
if not isinstance(model, (GPTModel, MambaModel, HybridModel, LLaVAModel)):
raise ValueError("Input to GPTModelExport must be a megatron.core.models.GPTModel!")

self._state_dict = OrderedDict()
Expand Down
16 changes: 12 additions & 4 deletions modelopt/torch/quantization/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,10 +236,18 @@ def find_quant_cfg_entry_by_path(
_mamba_moe_disabled_quantizer_cfg: list[QuantizerCfgEntry] = [
{"quantizer_name": "*fc1_latent_proj*", "enable": False}, # Skip Latent MOE
{"quantizer_name": "*fc2_latent_proj*", "enable": False}, # Skip Latent MOE
{"quantizer_name": "*q_proj*", "enable": False}, # Skip QKV Linear
{"quantizer_name": "*k_proj*", "enable": False}, # Skip QKV Linear
{"quantizer_name": "*v_proj*", "enable": False}, # Skip QKV Linear
{"quantizer_name": "*o_proj*", "enable": False}, # Skip QKV Output Projection
{"quantizer_name": "*q_proj*", "enable": False}, # Skip QKV Linear (HF naming)
{"quantizer_name": "*k_proj*", "enable": False}, # Skip QKV Linear (HF naming)
{"quantizer_name": "*v_proj*", "enable": False}, # Skip QKV Linear (HF naming)
{"quantizer_name": "*o_proj*", "enable": False}, # Skip QKV Output Projection (HF naming)
{
"quantizer_name": "*self_attention.linear_qkv*",
"enable": False,
}, # Skip QKV Linear (Mcore naming)
{
"quantizer_name": "*self_attention.linear_proj*",
"enable": False,
}, # Skip QKV Output Projection (Mcore naming)
]

INT8_DEFAULT_CFG = {
Expand Down
Loading