[TRTLLM-13052][feat] Enable TRTLLM moe backend for nemotron-h BF16 ckpt - #14944
Conversation
Wanli-Jiang
changed the title
ec5f7b6 to
cb15612
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #52257 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR extends the TRTLLM-Gen BF16 MoE backend to support multiple activation types beyond Swiglu, introduces a new activation set constant, simplifies MoE class fallback logic, converts activations to FlashInfer format, and adds comprehensive unit and integration tests covering new routing methods and activation combinations. ChangesTRTLLM BF16 MoE multi-activation support
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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: 1
🧹 Nitpick comments (3)
tests/unittest/_torch/modules/moe/test_moe_backend.py (3)
788-877: 💤 Low valueConsider expanding the test docstring to document coverage.
The current docstring is minimal. Expanding it to mention the parametrization dimensions (routing methods, activations, sequence lengths, routing modes) would improve discoverability and make the test's purpose clearer for future maintainers.
📝 Suggested docstring
def test_trtllm_bf16_unquantized_moe( routing_kind, activation_type, seq_len, trtllm_use_router_logits ): - """TRTLLM-Gen BF16 (unquantized) MoE accuracy vs the reference impl.""" + """TRTLLM-Gen BF16 (unquantized) MoE accuracy vs the reference impl. + + This test covers: + - Routing methods: DeepSeekV3, Renormalize + - Activations: Relu2, SwiGLU + - Sequence lengths: 8, 256 + - Routing modes: fused (router_logits) and separated (pre-computed indices) + + Validates backend output accuracy against the reference module after autotuning. + """ backend_type = MoeBackendType.TRTLLM🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_moe_backend.py` around lines 788 - 877, Update the docstring of test_trtllm_bf16_unquantized_moe to document what the test covers: mention the parametrization dimensions (routing_kind, activation_type, seq_len, trtllm_use_router_logits), the backend being tested (MoeBackendType.TRTLLM with bfloat16), the unquantized code path (BaseQuantizeUtil / unquantized weights), that it runs autotune and compares output accuracy against the reference fused MoE implementation, and any GPU/mapping setup assumptions; reference the test function name test_trtllm_bf16_unquantized_moe and key helpers like routing_method.apply, backend.quantize_input, run_backend_moe, and ref_fused_moe.check_accuracy so maintainers can quickly see what's exercised.
766-777: ⚡ Quick winAdd a docstring to the helper function.
The factory function branches between two routing methods with different configuration requirements. A docstring would improve maintainability by documenting the purpose, parameters, and return type.
📝 Suggested docstring
def _make_bf16_routing_method(routing_kind: str, top_k: int, num_experts: int, device: str): + """Create a routing method for BF16 unquantized MoE testing. + + Args: + routing_kind: Either "renormalize" or "deepseekv3" + top_k: Number of experts to activate per token + num_experts: Total number of experts + device: Device string for tensor placement + + Returns: + RenormalizeMoeRoutingMethod or DeepSeekV3MoeRoutingMethod instance + """ if routing_kind == "renormalize":🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_moe_backend.py` around lines 766 - 777, Add a clear docstring to the helper function _make_bf16_routing_method describing its purpose (factory that returns a BF16-capable routing method), the parameters (routing_kind: str, top_k: int, num_experts: int, device: str), the return type (RenormalizeMoeRoutingMethod or DeepSeekV3MoeRoutingMethod), and the branching behavior (when routing_kind == "renormalize" it returns RenormalizeMoeRoutingMethod; otherwise it constructs DeepSeekV3MoeRoutingMethod with a per-expert float32 bias created on the given device and the specific kwargs used: top_k, n_group=1, topk_group=1, routed_scaling_factor=2.5, callable_e_score_correction_bias).
780-787: 💤 Low valueConsider adding
idsto all parametrize decorators for consistency.Lines 781 and 785 provide explicit test IDs, but lines 783 and 787 don't. Adding
idsto all decorators improves test output readability and makes pytest collection output more uniform.📝 Suggested improvement
`@pytest.mark.parametrize`( "trtllm_use_router_logits", [True, False], ids=["fused_routing", "separated_routing"] ) -@pytest.mark.parametrize("seq_len", [8, 256]) +@pytest.mark.parametrize("seq_len", [8, 256], ids=["short", "long"]) `@pytest.mark.parametrize`( "activation_type", [ActivationType.Relu2, ActivationType.Swiglu], ids=["relu2", "swiglu"] ) -@pytest.mark.parametrize("routing_kind", ["deepseekv3", "renormalize"]) +@pytest.mark.parametrize("routing_kind", ["deepseekv3", "renormalize"], ids=["deepseekv3", "renormalize"])🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_moe_backend.py` around lines 780 - 787, Add explicit ids to the remaining pytest.mark.parametrize decorators in the test file to match the others: update the "seq_len" parametrize (currently pytest.mark.parametrize("seq_len", [8, 256])) to include ids (e.g., ids=["short", "long"] or ["8", "256"]) and update the "routing_kind" parametrize (currently pytest.mark.parametrize("routing_kind", ["deepseekv3", "renormalize"])) to include ids (e.g., ids=["deepseekv3", "renormalize"]); keep the existing decorators for "trtllm_use_router_logits" and "activation_type" unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 6823-6840: The test test_bf16_trtllm_gen_moe_backend forces
MoeConfig(backend="TRTLLM") but lacks the SM120/SM121 skip other TRTLLM MoE
tests have; add a guard at the top of test_bf16_trtllm_gen_moe_backend that
skips the test when the current GPU SM is 120 or 121 (use the same skip helper
used elsewhere in this file, e.g., skip_if_sm_in or the project’s SM-check
helper), so the test returns/skips before creating the LLM with
MoeConfig(backend="TRTLLM") on unsupported SM120/SM121 targets.
---
Nitpick comments:
In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 788-877: Update the docstring of test_trtllm_bf16_unquantized_moe
to document what the test covers: mention the parametrization dimensions
(routing_kind, activation_type, seq_len, trtllm_use_router_logits), the backend
being tested (MoeBackendType.TRTLLM with bfloat16), the unquantized code path
(BaseQuantizeUtil / unquantized weights), that it runs autotune and compares
output accuracy against the reference fused MoE implementation, and any
GPU/mapping setup assumptions; reference the test function name
test_trtllm_bf16_unquantized_moe and key helpers like routing_method.apply,
backend.quantize_input, run_backend_moe, and ref_fused_moe.check_accuracy so
maintainers can quickly see what's exercised.
- Around line 766-777: Add a clear docstring to the helper function
_make_bf16_routing_method describing its purpose (factory that returns a
BF16-capable routing method), the parameters (routing_kind: str, top_k: int,
num_experts: int, device: str), the return type (RenormalizeMoeRoutingMethod or
DeepSeekV3MoeRoutingMethod), and the branching behavior (when routing_kind ==
"renormalize" it returns RenormalizeMoeRoutingMethod; otherwise it constructs
DeepSeekV3MoeRoutingMethod with a per-expert float32 bias created on the given
device and the specific kwargs used: top_k, n_group=1, topk_group=1,
routed_scaling_factor=2.5, callable_e_score_correction_bias).
- Around line 780-787: Add explicit ids to the remaining pytest.mark.parametrize
decorators in the test file to match the others: update the "seq_len"
parametrize (currently pytest.mark.parametrize("seq_len", [8, 256])) to include
ids (e.g., ids=["short", "long"] or ["8", "256"]) and update the "routing_kind"
parametrize (currently pytest.mark.parametrize("routing_kind", ["deepseekv3",
"renormalize"])) to include ids (e.g., ids=["deepseekv3", "renormalize"]); keep
the existing decorators for "trtllm_use_router_logits" and "activation_type"
unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4cb1dbde-0900-4d7d-b507-b1f23ab1689f
📒 Files selected for processing (9)
tensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytensorrt_llm/_torch/modules/fused_moe/moe_op_backend.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/integration/test_lists/test-db/l0_gb200_multi_gpus.ymltests/unittest/_torch/modules/moe/test_moe_backend.py
b97c4c0 to
1315f2e
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #52323 [ run ] triggered by Bot. Commit: |
|
PR_Github #52257 [ run ] completed with state |
|
PR_Github #52323 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52514 [ run ] triggered by Bot. Commit: |
|
PR_Github #52514 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52634 [ run ] triggered by Bot. Commit: |
|
PR_Github #52634 [ run ] completed with state
|
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
1315f2e to
e146a3c
Compare
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. 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. |
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #52923 [ run ] triggered by Bot. Commit: |
|
PR_Github #52923 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #52990 [ run ] triggered by Bot. Commit: |
|
PR_Github #52990 [ run ] completed with state
|
|
/bot skip --comment "Test passed at 52634 for single GPUs and 52923 for multi GPUs" |
|
PR_Github #53189 [ skip ] triggered by Bot. Commit: |
|
PR_Github #53189 [ skip ] completed with state |
6 participants
Summary by CodeRabbit
New Features
Improvements
Tests
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.