[None][chore] Optimize MOE export by tracing with reduced experts and expanding graph - #11504
Conversation
Speed up torch.export for MOE models by temporarily reducing the number of experts during tracing, then programmatically expanding the FX graph to include all expert weights afterward. Expert nn.ModuleLists are discovered via a TorchDispatchMode probe that intercepts torch_moe-family custom ops during a lightweight forward pass, so the optimization does NOT rely on any particular attribute name. Changes: - export.py: Add _MoeExpertProbe / _find_moe_module_lists for name- independent expert discovery. Add _reduce_moe_experts / _restore_moe_experts / _expand_moe_experts_in_graph helpers and wire them into torch_export_to_gm via num_moe_experts_for_export. - export_to_gm.py: Thread num_moe_experts_for_export through the ExportToGMConfig and transform pipeline. - test_export.py: Add SimpleMoEForExport unit tests (parametrized over expert_attr_name='experts' and 'mlp_bank' to prove name independence) and optional Glm4MoeLiteForCausalLM real-model tests. Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
suyoggupta
changed the title
📝 WalkthroughWalkthroughThe changes introduce MOE (Mixture of Experts) export optimization to the auto-deploy system. This includes utilities to trace models with reduced expert counts, expand exported graphs with full expert weights post-export, discover MOE module structures, and comprehensive tests validating the optimization across generic and GLM4-based MOE models. Changes
Sequence DiagramsequenceDiagram
participant User
participant ExportPipeline as Export Pipeline
participant Reducer as MOE Reducer
participant Exporter as Torch Exporter
participant Restorer as MOE Restorer
participant GraphExpander as Graph Expander
participant Model as Model
User->>ExportPipeline: Call torch_export_to_gm with num_moe_experts_for_export
ExportPipeline->>Reducer: _reduce_moe_experts(model, min_num_experts)
Reducer->>Model: Forward pass with probing
Model-->>Reducer: Module list contributions identified
Reducer->>Model: Truncate expert lists to min_num_experts
Reducer-->>ExportPipeline: Metadata (reductions)
ExportPipeline->>Exporter: torch.export.export(reduced_model)
Exporter->>Exporter: Trace with reduced expert count
Exporter-->>ExportPipeline: GraphModule (reduced)
ExportPipeline->>Restorer: _restore_moe_experts(reductions)
Restorer->>Model: Restore original expert lists
Restorer-->>ExportPipeline: Model restored
ExportPipeline->>GraphExpander: _expand_moe_experts_in_graph(gm, model, reductions)
GraphExpander->>GraphExpander: Inject missing expert parameters
GraphExpander->>GraphExpander: Create get_attr nodes for each expert
GraphExpander-->>ExportPipeline: GraphModule (expanded)
ExportPipeline-->>User: Full GraphModule with all experts
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@tensorrt_llm/_torch/auto_deploy/export/export.py`:
- Around line 316-318: The new get_attr node is being assigned a real
nn.Parameter into new_node.meta["val"] while existing get_attr nodes hold
FakeTensor meta values, causing inconsistency; update the insertion logic
(around graph.inserting_before(first_graph_node), new_node, new_target, and
gm.get_parameter) to copy the meta dict from an existing sibling get_attr (e.g.,
weight_list[0].meta) into new_node.meta instead of directly setting meta["val"],
preserving shape/dtype/placeholder fields (or at least replicate meta["val"],
"tensor_meta", dtype, and shape entries) so downstream FX passes see consistent
FakeTensor-like metadata.
- Around line 675-690: The code currently reduces MOE experts via
_reduce_moe_experts when num_moe_experts_for_export is set but calls
run_forward_for_capture before restoring; if run_forward_for_capture throws the
model stays truncated—wrap the run_forward_for_capture call and the subsequent
_expand_moe_experts_in_graph invocation in a try/finally: after computing
moe_reductions, run run_forward_for_capture inside try, and in finally call
_restore_moe_experts(moe_reductions) and _expand_moe_experts_in_graph(egm,
model, moe_reductions) only if moe_reductions is non-empty, then rethrow the
original exception so behavior is unchanged.
🧹 Nitpick comments (6)
tensorrt_llm/_torch/auto_deploy/export/export.py (3)
13-13: Import style: class imported directly from module.Per the coding guidelines, prefer
from package.subpackage import moduleover importing a class directly. HereTorchDispatchModeis imported directly.However, this is the canonical PyTorch import pattern for
TorchDispatchMode, so this is a minor style nit.As per coding guidelines: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class"
104-105: Duplicate MOE op lists in two places.
_MoeExpertProbe._MOE_OP_NAMES(Line 105) and themoe_opsset in_expand_moe_experts_in_graph(Lines 257-259) both enumerate the same MOE ops in different formats. If a new MOE op is added, both need updating independently. Consider defining a single source of truth.Also applies to: 256-260
236-329: Core graph expansion logic looks correct.The flow of finding MOE nodes → inferring naming patterns → registering missing parameters → creating
get_attrnodes → extending argument lists is sound. The insertion beforefirst_graph_nodeensures parameter nodes appear at the graph's beginning, following FX conventions.One note: Line 302 uses
assertfor a condition that depends on runtime data (finding the expert prefix). In production this could be silently removed with-O.Consider replacing assert with a RuntimeError
- assert ep is not None, ( - f"Could not find expert prefix for target '{weight_list[0].target}'" - ) + if ep is None: + raise RuntimeError( + f"Could not find expert prefix for target '{weight_list[0].target}'" + )tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py (3)
249-249: Import style: importing classes directly.Per the coding guidelines, this should use
from tensorrt_llm._torch.auto_deploy import custom_ops. The side-effect import only needs the module, not specific classes.Also, the
# noqa: F401, E402directive is flagged as unused by Ruff (RUF100) since those rules aren't enabled.Suggested fix
-import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401, E402 +import tensorrt_llm._torch.auto_deploy.custom_ops # side-effect: register custom opsAs per coding guidelines: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class"
331-335: Duplicated helper:_count_moe_experts(Line 331) and_count_moe_experts_in_graph(Line 400).These two functions are nearly identical. Consider extracting a single module-level helper and reusing it in both test functions.
Suggested consolidation
+def _count_moe_experts_in_graph(gm: GraphModule) -> int: + """Return the number of experts in the first ``torch_moe`` call in *gm*.""" + for node in gm.graph.nodes: + if node.op == "call_function" and "torch_moe" in str(node.target): + return len(node.args[3]) # w1_weight list length + return 0 + + `@pytest.mark.parametrize`("expert_attr_name", ["experts", "mlp_bank"]) # ... in test_moe_export_with_reduced_experts: - def _count_moe_experts(gm): - for node in gm.graph.nodes: - if node.op == "call_function" and "torch_moe" in str(node.target): - return len(node.args[3]) # w1_weight list length - return 0 - - assert _count_moe_experts(gm_full) == num_experts - assert _count_moe_experts(gm_reduced) == num_experts + assert _count_moe_experts_in_graph(gm_full) == num_experts + assert _count_moe_experts_in_graph(gm_reduced) == num_expertsAlso applies to: 400-405
359-367: Import style: classes imported directly.Per the coding guidelines, this should import the module rather than classes:
- from tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import ( - Glm4MoeLiteConfig, - Glm4MoeLiteForCausalLM, - ) + from tensorrt_llm._torch.auto_deploy.models.custom import modeling_glm4_moe_liteThen reference as
modeling_glm4_moe_lite.Glm4MoeLiteConfig, etc. However, since this is a conditional import within atry/exceptblock in a test file, this is a minor nit.As per coding guidelines: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class"
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
|
/bot run |
|
PR_Github #35920 [ run ] triggered by Bot. Commit: |
|
PR_Github #35920 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
1 similar comment
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #35947 [ run ] triggered by Bot. Commit: |
|
PR_Github #35947 [ run ] completed with state |
- Replace hardcoded magic number 3 with schema-driven introspection for finding List[Tensor] arguments in MOE ops - Rename min_num_experts to num_moe_experts_for_export for consistency - Remove unnecessary try/except guard around GLM4 MoE Lite import in tests Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
… expanding graph (NVIDIA#11504) Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Signed-off-by: peihu-nv <259410613+peihu-nv@users.noreply.github.com>
4 participants
Speed up torch.export for MOE models by temporarily reducing the number of experts during tracing, then programmatically expanding the FX graph to include all expert weights afterward.
Expert nn.ModuleLists are discovered via a TorchDispatchMode probe that intercepts torch_moe-family custom ops during a lightweight forward pass, so the optimization does NOT rely on any particular attribute name.
Changes:
superv3 export time on H100, tp4, fp8: 66s -> 22s
Summary by CodeRabbit
New Features
Tests