feat:[AutoDeploy] Support Quantized MoE matcher - Step1 by Fridah-nv · Pull Request #68 · nv-auto-deploy/TensorRT-LLM

Fridah-nv · 2025-06-27T21:27:47Z

Description

update moe pattern matcher to handle fp8_linear in addition to linear. Also, add the corresponding torch_fp8_moe.
Update FP4 for the same
Update EP sharding with quantized ops
Add quantize_moe pass for patched torch_moe ops for FP4 and FP8
Add unit tests for custom ops and transformation pass

TODO:

update fuse_moe to map torch_{fp4, fp8}moe to trtllm{fp4, fp8}_moe
Update README Precision Support section

Test Coverage

Tested e2e with nvidia/Mixtral-8x7B-Instruct-v0.1-FP8 and nvidia/Mixtral-8x7B-Instruct-v0.1-FP4 on B100, works with both world_size=1 and world_size=2
Note that the Mixtral models are now in patched moe path after the patch in qwen3.py

nvidia/DeepSeek-R1-FP4 only work with full_state_dict set to true in load_checkpoint_in_model (see #5892 ), but it also suffers from long execution time (I loaded the full model on 1 node and it took 3 hours to went through the pipeline and start loading weights). Verified that 4 layers can run e2e.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]

Launch build/test pipelines. All previously running jobs will be killed.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-[Post-Merge]-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

lucaslie

Looks great!

suyoggupta · 2025-07-10T21:33:31Z

@Fridah-nv re. I loaded the full model on 1 node and it took 3 hours to went through the pipeline and start loading weights: was the time spent on waiting for torch export to complete?

Fridah-nv · 2025-07-10T22:00:43Z

@suyoggupta was the time spent on waiting for torch export to complete?
Yes, torch export time is still long after we patched moe op, also every transformation pass takes quite some time (unless it is skipped). I think it's mainly because the graph is still large.

…nit tests Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

…cc issue and e2e mixtral fp4 has kernel error wo moe matching Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

… exclude_modules in hf_quant_config Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

… refinements Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

suyoggupta · 2025-07-13T19:31:49Z

+            # model.load_state_dict.
+            # This sync step can interfere with load_hooks by mixing raw checkpoint weights and
+            # model-transformed weights,leading to unexpected key mismatches or format issues.
+            load_checkpoint_in_model(model, checkpoint=ckpt_file, full_state_dict=False)


Hi @Fridah-nv , I think this is breaking support for models where some of the weights are aliased. Before this changes, models in which lm_head.weight and embed.weight are aliased worked fine, but they raise the following error after this change:https://gitlab-master.nvidia.com/dl/jet/ci/-/jobs/187865431/raw

* add torch_fp8_moe and fp8 linear support in pattern matcher, update unit tests Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * add torch-fp4-moe and fp4 support in pattern matcher, unit test has acc issue and e2e mixtral fp4 has kernel error wo moe matching Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * add pre-commit hook Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * hacky fix for e2e run of mixtral FP4 and fp4 op unit test Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * EP support for torch_fp4_moe and torch_fp8_moe Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * fix rebase: op rename, shard_load_hook bug in FP4 Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * fix pre-commit Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * fix weight loading-load_hook issue for FP4, update function to handle exclude_modules in hf_quant_config Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * addressing feedback, add moe op template, update op names,other minor refinements Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * move common functionality to utility Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * fix FP4QuantizationImpl register from rebase Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * add quantize_moe pass for patched torch_moe op Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * add transformation unit tests for FP8 and FP4 Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * update should_skip_quantization to fix bmm unit test Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * update BMMDynamicModel and utils to extract weight for dynamic BMM case Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * update BMMDynamicModel to drop linear op Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> * minor Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> --------- Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

Fridah-nv mentioned this pull request Jun 27, 2025

feat:[AutoDeploy] Support Quantized MoE matcher - 1 NVIDIA/TensorRT-LLM#4968

Closed

lucaslie reviewed Jul 1, 2025

View reviewed changes

nzmora-nvidia reviewed Jul 2, 2025

View reviewed changes

Comment thread tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_ad_moe_op.py

Fridah-nv self-assigned this Jul 3, 2025

Fridah-nv commented Jul 3, 2025

View reviewed changes

Comment thread tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

Fridah-nv requested review from meenchen and sugunav14 July 3, 2025 19:17

lucaslie approved these changes Jul 7, 2025

View reviewed changes

Comment thread tensorrt_llm/_torch/auto_deploy/models/hf.py

Comment thread tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

lucaslie reviewed Jul 7, 2025

View reviewed changes

Comment thread tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_ad_moe_op.py

Fridah-nv force-pushed the user/fridah/moe-quant-match branch from d3eceb0 to 3715a45 Compare July 8, 2025 00:27

Fridah-nv changed the base branch from feat/ad-2025-06-24 to feat/ad-2025-07-07 July 8, 2025 00:28

Fridah-nv commented Jul 10, 2025

View reviewed changes

Comment thread tensorrt_llm/_torch/auto_deploy/transformations/library/quantize_moe.py

Fridah-nv commented Jul 10, 2025

View reviewed changes

Comment thread tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py

Fridah-nv requested review from greg-kwasniewski1, nvchenghaoz and suyoggupta July 10, 2025 21:27

lucaslie approved these changes Jul 11, 2025

View reviewed changes

Comment thread tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_ad_moe_op.py

Fridah-nv force-pushed the user/fridah/moe-quant-match branch 2 times, most recently from dcbb1fe to 2bfd86a Compare July 11, 2025 20:05

nvchenghaoz reviewed Jul 11, 2025

View reviewed changes

Fridah-nv added 8 commits July 11, 2025 14:37

add torch_fp8_moe and fp8 linear support in pattern matcher, update u…

ea6428e

…nit tests Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

add torch-fp4-moe and fp4 support in pattern matcher, unit test has a…

2d12d12

…cc issue and e2e mixtral fp4 has kernel error wo moe matching Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

add pre-commit hook

84201be

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

hacky fix for e2e run of mixtral FP4 and fp4 op unit test

bdf304f

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

EP support for torch_fp4_moe and torch_fp8_moe

e4d8457

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

fix rebase: op rename, shard_load_hook bug in FP4

bb83fa8

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

fix pre-commit

bfd783e

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

fix weight loading-load_hook issue for FP4, update function to handle…

71a2114

… exclude_modules in hf_quant_config Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

Fridah-nv added 9 commits July 11, 2025 14:37

addressing feedback, add moe op template, update op names,other minor…

010f3cf

… refinements Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

move common functionality to utility

7e647fa

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

fix FP4QuantizationImpl register from rebase

a9e3ff7

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

add quantize_moe pass for patched torch_moe op

1e6cdcd

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

add transformation unit tests for FP8 and FP4

b857f06

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

update should_skip_quantization to fix bmm unit test

e0982e7

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

update BMMDynamicModel and utils to extract weight for dynamic BMM case

ee7134e

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

update BMMDynamicModel to drop linear op

d2071c2

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

minor

e20dccc

Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>

Fridah-nv force-pushed the user/fridah/moe-quant-match branch from 606ff78 to e20dccc Compare July 11, 2025 21:37

Fridah-nv enabled auto-merge (squash) July 11, 2025 21:38

Fridah-nv merged commit 6535ee6 into feat/ad-2025-07-07 Jul 11, 2025
1 check passed

Fridah-nv mentioned this pull request Jul 11, 2025

feat:[AutoDeploy] Support Quantized MoE matcher NVIDIA/TensorRT-LLM#4361

Closed

5 tasks

suyoggupta reviewed Jul 13, 2025

View reviewed changes

Uh oh!

Conversation

Fridah-nv commented Jun 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

GitHub Bot Help

kill

skip

reuse-pipeline

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lucaslie left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

suyoggupta commented Jul 10, 2025

Uh oh!

Fridah-nv commented Jul 10, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

suyoggupta Jul 13, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Fridah-nv commented Jun 27, 2025 •

edited

Loading