From 4fa630a3155a75e685d0e8d7b05db739522cbab2 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 24 Apr 2026 12:28:33 -0700 Subject: [PATCH 1/7] support EP import for TE and fix mamba moe config Signed-off-by: Jennifer Chen --- .../torch/export/plugins/megatron_importer.py | 21 +++++++++++++------ modelopt/torch/quantization/config.py | 13 +++++++----- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/export/plugins/megatron_importer.py b/modelopt/torch/export/plugins/megatron_importer.py index b1d37c1ad90..fbaf55e29fe 100644 --- a/modelopt/torch/export/plugins/megatron_importer.py +++ b/modelopt/torch/export/plugins/megatron_importer.py @@ -39,6 +39,8 @@ has_mcore = False with import_plugin("megatron"): from megatron.core.parallel_state import ( + get_expert_model_parallel_rank, + get_expert_model_parallel_world_size, get_expert_tensor_parallel_world_size, get_tensor_model_parallel_world_size, ) @@ -294,9 +296,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) @@ -653,10 +659,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 self.rules["experts.linear_fc1"]( experts.linear_fc1, diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 186ff1c7edd..a225ee9061b 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -236,11 +236,14 @@ 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 = { "quant_cfg": [ From 948e4deb212e215006f32940c2c2d8eceb9900ad Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 24 Apr 2026 13:12:09 -0700 Subject: [PATCH 2/7] lint Signed-off-by: Jennifer Chen --- modelopt/torch/export/plugins/megatron_importer.py | 1 - modelopt/torch/quantization/config.py | 13 +++++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/plugins/megatron_importer.py b/modelopt/torch/export/plugins/megatron_importer.py index fbaf55e29fe..e485731b3d8 100644 --- a/modelopt/torch/export/plugins/megatron_importer.py +++ b/modelopt/torch/export/plugins/megatron_importer.py @@ -40,7 +40,6 @@ with import_plugin("megatron"): from megatron.core.parallel_state import ( get_expert_model_parallel_rank, - get_expert_model_parallel_world_size, get_expert_tensor_parallel_world_size, get_tensor_model_parallel_world_size, ) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index a225ee9061b..34e7f692ca0 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -240,10 +240,15 @@ def find_quant_cfg_entry_by_path( {"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) - - ] + { + "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 = { "quant_cfg": [ From 67e021dfbc28b8a10dfab810e6c2f78dc0214c4e Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 24 Apr 2026 13:39:44 -0700 Subject: [PATCH 3/7] fix hybrid model export Signed-off-by: Jennifer Chen --- modelopt/torch/distill/plugins/megatron.py | 2 +- modelopt/torch/export/unified_export_megatron.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index dbfad6fb6bb..a460204f512 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -163,7 +163,7 @@ def setup_distillation_config( def _adjust_layer_index_for_pp(submodule_name, model_cfg): """Adjust any sequence-based layer indices found in a submodule name for Pipeline Parallelism.""" - match = re.search(r"(?<=\.)\d+(?=\.)", submodule_name) + match = re.search(r"(?<=\.)\d+(?=\.|$)", submodule_name) if not match: return submodule_name diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 62053e549c8..0ef7e916013 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -72,6 +72,10 @@ 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, @@ -121,7 +125,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() From 27a4eef5596539515c30cb517fc80d93736ce6c5 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 10:25:53 -0700 Subject: [PATCH 4/7] lint Signed-off-by: Jennifer Chen --- modelopt/torch/export/unified_export_megatron.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 0ef7e916013..24983a24891 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -72,6 +72,7 @@ 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: From 67b8848e427c12f1cd05098d4408d6a21ccc933a Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 12:09:21 -0700 Subject: [PATCH 5/7] getter/setter for QuantCfg Signed-off-by: Jennifer Chen --- modelopt/torch/quantization/config.py | 28 +++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 34e7f692ca0..c5a8df2604f 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -170,6 +170,34 @@ class QuantizerCfgEntry(TypedDict, total=False): enable: bool | None # toggles matched quantizers on/off; independent of cfg +QuantizerCfg = list[QuantizerCfgEntry] + + +def _set_quant_cfg_entry( + quant_cfg: QuantizerCfg, quantizer_name: str, entry_cfg: QuantizerCfgEntry +) -> None: + """Set a QuantizerCfgEntry by quantizer_name in the QuantizerCfg.""" + for entry in quant_cfg: + if entry.get("quantizer_name") == quantizer_name: + # mypy workaround for TypedDict: treat as a plain dict for the merge. + cast("dict[str, Any]", entry).update(entry_cfg) + return + quant_cfg.append( + cast( + "QuantizerCfgEntry", + {"quantizer_name": quantizer_name, **entry_cfg}, + ) + ) + + +def _get_quant_cfg_entry(quant_cfg: QuantizerCfg, quantizer_name: str) -> QuantizerCfgEntry | None: + """Get a QuantizerCfgEntry by quantizer_name in the QuantizerCfg.""" + for entry in quant_cfg: + if entry.get("quantizer_name") == quantizer_name: + return entry + return None + + def find_quant_cfg_entry_by_path( quant_cfg_list: list[QuantizerCfgEntry], quantizer_name: str ) -> QuantizerCfgEntry: From 0b3ca9452ee86a7d02de32883717d577c0177f90 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 12:29:34 -0700 Subject: [PATCH 6/7] remove setter Signed-off-by: Jennifer Chen --- modelopt/torch/quantization/config.py | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index c5a8df2604f..ac3ce7e6fb9 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -173,26 +173,9 @@ class QuantizerCfgEntry(TypedDict, total=False): QuantizerCfg = list[QuantizerCfgEntry] -def _set_quant_cfg_entry( - quant_cfg: QuantizerCfg, quantizer_name: str, entry_cfg: QuantizerCfgEntry -) -> None: - """Set a QuantizerCfgEntry by quantizer_name in the QuantizerCfg.""" - for entry in quant_cfg: - if entry.get("quantizer_name") == quantizer_name: - # mypy workaround for TypedDict: treat as a plain dict for the merge. - cast("dict[str, Any]", entry).update(entry_cfg) - return - quant_cfg.append( - cast( - "QuantizerCfgEntry", - {"quantizer_name": quantizer_name, **entry_cfg}, - ) - ) - - def _get_quant_cfg_entry(quant_cfg: QuantizerCfg, quantizer_name: str) -> QuantizerCfgEntry | None: """Get a QuantizerCfgEntry by quantizer_name in the QuantizerCfg.""" - for entry in quant_cfg: + for entry in reversed(quant_cfg): if entry.get("quantizer_name") == quantizer_name: return entry return None From a7dc3001948d1d5712a1d4315c39e9c3427ccd63 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Mon, 27 Apr 2026 12:47:44 -0700 Subject: [PATCH 7/7] revert getter Signed-off-by: Jennifer Chen --- modelopt/torch/quantization/config.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index ac3ce7e6fb9..34e7f692ca0 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -170,17 +170,6 @@ class QuantizerCfgEntry(TypedDict, total=False): enable: bool | None # toggles matched quantizers on/off; independent of cfg -QuantizerCfg = list[QuantizerCfgEntry] - - -def _get_quant_cfg_entry(quant_cfg: QuantizerCfg, quantizer_name: str) -> QuantizerCfgEntry | None: - """Get a QuantizerCfgEntry by quantizer_name in the QuantizerCfg.""" - for entry in reversed(quant_cfg): - if entry.get("quantizer_name") == quantizer_name: - return entry - return None - - def find_quant_cfg_entry_by_path( quant_cfg_list: list[QuantizerCfgEntry], quantizer_name: str ) -> QuantizerCfgEntry: