[https://nvbugs/6175060][fix] Fix B300 MegaMoE test selection#14362
Conversation
📝 WalkthroughWalkthroughThis PR expands MegaMoEDeepGemm backend GPU support from strictly SM100 to the SM100 family (SM100/SM103) using a new helper function, documents the change, adds validation tests, and improves distributed test reliability by adding automatic port-conflict retry logic for multiprocessing worker spawning. ChangesMegaMoEDeepGemm SM100-family support
Distributed test reliability improvements
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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.
🧹 Nitpick comments (2)
tests/unittest/_torch/modules/moe/test_moe_backend.py (2)
214-245: QA list update is unnecessary for this PR scope.These are unit-test additions only; no integration test definitions were introduced, so QA list file changes are optional/not required here.
As per coding guidelines: “If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional.”
🤖 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 214 - 245, The PR adds only unit tests (test_megamoe_can_implement_accepts_sm100_family and test_megamoe_can_implement_rejects_non_sm100_family in tests/unittest/_torch/modules/moe/test_moe_backend.py) so explicitly state in the PR description or commit message that QA list updates are unnecessary/optional for this change because it only touches unit tests under tests/unittest/ and does not introduce integration or system-level changes; include the two test names to make the rationale clear.
214-229: ⚡ Quick winCover both supported SM-family members in the acceptance test.
The contract is SM100-family support (SM100 and SM103), but the positive path currently asserts only SM103. Parametrizing this test over both values will prevent regressions on SM100.
Proposed change
+@pytest.mark.parametrize("sm_version", [100, 103], ids=["sm100", "sm103"]) -def test_megamoe_can_implement_accepts_sm100_family(monkeypatch): +def test_megamoe_can_implement_accepts_sm100_family(monkeypatch, sm_version): import tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_deepgemm as megamoe_module - monkeypatch.setattr(megamoe_module, "get_sm_version", lambda: 103) + monkeypatch.setattr(megamoe_module, "get_sm_version", lambda: sm_version) monkeypatch.setattr(megamoe_module, "_import_deep_gemm", lambda: object()) can_implement, reason = MegaMoEDeepGemm.can_implement( QuantAlgo.W4A8_MXFP4_MXFP8, dtype_activation=torch.bfloat16, hidden_size=512, intermediate_size=512, ) assert can_implement assert reason is None🤖 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 214 - 229, The test test_megamoe_can_implement_accepts_sm100_family should cover both SM100 and SM103; change the test to iterate or parametrize over SM family values (100 and 103) when monkeypatching megamoe_module.get_sm_version and asserting that MegaMoEDeepGemm.can_implement returns (True, None); keep the existing monkeypatch of _import_deep_gemm and the same can_implement call/assertions but run them for each SM value to prevent regressions for SM100.
🤖 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.
Nitpick comments:
In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 214-245: The PR adds only unit tests
(test_megamoe_can_implement_accepts_sm100_family and
test_megamoe_can_implement_rejects_non_sm100_family in
tests/unittest/_torch/modules/moe/test_moe_backend.py) so explicitly state in
the PR description or commit message that QA list updates are
unnecessary/optional for this change because it only touches unit tests under
tests/unittest/ and does not introduce integration or system-level changes;
include the two test names to make the rationale clear.
- Around line 214-229: The test test_megamoe_can_implement_accepts_sm100_family
should cover both SM100 and SM103; change the test to iterate or parametrize
over SM family values (100 and 103) when monkeypatching
megamoe_module.get_sm_version and asserting that MegaMoEDeepGemm.can_implement
returns (True, None); keep the existing monkeypatch of _import_deep_gemm and the
same can_implement call/assertions but run them for each SM value to prevent
regressions for SM100.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 505337b2-a56d-475a-889e-d1c97b35e4e5
📒 Files selected for processing (6)
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.pytests/integration/test_lists/waives.txttests/unittest/_torch/modules/moe/test_moe_backend.pytests/unittest/_torch/visual_gen/test_flux_pipeline.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
df66828 to
f912557
Compare
xxi-nv
changed the title
Signed-off-by: xxi <xxi@nvidia.com>
f912557 to
21264b1
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #49435 [ run ] triggered by Bot. Commit: |
xxi-nv
changed the title
|
PR_Github #49435 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49503 [ run ] triggered by Bot. Commit: |
|
PR_Github #49503 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49558 [ run ] triggered by Bot. Commit: |
|
Please update the nvbug link in title to 6175060 |
xxi-nv
changed the title
xxi-nv
left a comment
xxi-nv
left a comment
There was a problem hiding this comment.
Thanks for the catch, update the tile and actually 6175060 and 6175059 are duplicated bugs.
|
PR_Github #49558 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49596 [ run ] triggered by Bot. Commit: |
|
PR_Github #49596 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49658 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #49672 [ run ] triggered by Bot. Commit: |
|
PR_Github #49658 [ run ] completed with state |
|
PR_Github #49672 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49769 [ run ] triggered by Bot. Commit: |
|
PR_Github #49769 [ run ] completed with state |
…#14362) Signed-off-by: xxi <xxi@nvidia.com>
…#14362) Signed-off-by: xxi <xxi@nvidia.com>
3 participants
Summary
MegaMoEDeepGemm.can_implement()to accept the SM100 family (SM100/SM103), matching the underlying DeepGEMMfp8_fp4_mega_moearchitecture gating.MEGAMOE_DEEPGEMMwaives. The source branch did not contain a literal6175059waive entry; the matching MoE waives were linked to6175060and are removed here.tests/unittest/_torch/visual_gen/test_flux_pipeline.pyandtests/unittest/_torch/modules/moe/test_moe_backend.pysource changes.Root Cause
The failing post-merge signal for NVBug 6175059 is the B300 MoE selector:
unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "MEGAMOE_DEEPGEMM"The failure was caused by B300/SM103 being filtered out by a Python-side
sm == 100gate inMegaMoEDeepGemm.can_implement(), so pytest selectors forMEGAMOE_DEEPGEMMselected zero tests.DeepGEMM itself handles the SM100 family path: the MegaMoE C++ API dispatches on
arch_major == 10, the JIT chooses SM100 family arch for major 10 devices, and the CUDA kernel guard accepts__CUDA_ARCH__ >= 1000. B300 reports SM103.Validation
python3 -m py_compile tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.pygit diff --checkMEGAMOE_DEEPGEMMentriesdevice=NVIDIA B300 SXM6 AC,capability=(10, 3),get_sm_version=103,is_sm_100f=True,MegaMoEDeepGemm.can_implement=Truetest_configurable_moe_multi_gpu -k MEGAMOE_DEEPGEMM: 4 collectedtest_configurable_moe_single_gpu -k MEGAMOE_DEEPGEMM: 28 collectedtest_configurable_moe_multi_gpu_eplb -k MEGAMOE_DEEPGEMM: 1 collectedOne additional single-GPU smoke failed in the non-DeepGEMM reference path because the temporary container's TRT-LLM custom ops were not compiled for B300. That does not affect the
can_implement()fix or selector validation.