[TRTLLM-13628][test] Optimize MoE comm test execution - #15662
Conversation
|
/bot run --disable-fail-fast |
sunnyqgg
changed the title
|
PR_Github #56061 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis test module now groups MoE communication configs, caches worker communication objects, carries tensor-decoded routing data through dispatch verification, and rewrites combine-reference calculations with vectorized masking and low-precision simulation. ChangesMoE communication test harness
Sequence Diagram(s)sequenceDiagram
participant TestMoEComm
participant run_group as _run_full_test_group
participant mpi_pool_executor
participant get_comm as _get_worker_comm
participant prepare_ref as _prepare_moe_output_for_combine_reference
participant verify_results as _verify_full_test_results
TestMoEComm->>run_group: run CommTestGroup
run_group->>mpi_pool_executor: submit worker futures for grouped configs
mpi_pool_executor->>get_comm: execute worker pipeline with cached comm
get_comm->>prepare_ref: prepare low-precision reference payload
run_group->>verify_results: verify drained worker results
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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: 3
🤖 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/unittest/_torch/modules/moe/test_moe_comm.py`:
- Line 1682: The loop in the boundary case iteration uses an unused variable
that triggers Ruff B007. Update the tuple unpacking in the test loop to either
rename case_id to _case_id or remove it from the boundary_cases payload, keeping
the logic in test_moe_comm unchanged.
- Around line 1808-1814: The helper `_drain_pending_results` is swallowing all
exceptions from `Future.result()` without any visibility, which hides drained
worker failures. Update this loop to avoid silent broad exception handling:
either catch the specific exception types you expect from worker futures or, if
broad catching is required, log the exception with a clear rationale before
continuing. Keep the behavior of draining pending work, but make failures
observable for diagnosis.
- Around line 1827-1838: The low-precision combine tolerance branch in
test_moe_comm.py is too broad because the current else after config.comm_type ==
COMM_NVLINK_ONE_SIDED also captures DeepEPLL cases. Tighten the branching around
verify_combine_results so the NVFP4 atol logic only applies to the
NVLinkTwoSided path, and keep DeepEPLL handled by its own explicit tolerance
path within the same low_precision combine test flow. Use the existing
config.comm_type checks and the verify_combine_results call site to preserve
backend-specific coverage.
🪄 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: 846e6a61-04dd-4615-9530-52bd09190158
📒 Files selected for processing (1)
tests/unittest/_torch/modules/moe/test_moe_comm.py
|
PR_Github #56061 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56285 [ run ] triggered by Bot. Commit: |
|
PR_Github #56285 [ run ] completed with state |
xxi-nv
left a comment
xxi-nv
left a comment
There was a problem hiding this comment.
Thanks for the great optimization and please help to address the comments from codderabbitai.
Signed-off-by: qgai <qgai@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #57303 [ run ] triggered by Bot. Commit: |
|
PR_Github #57303 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57680 [ run ] triggered by Bot. Commit: |
|
PR_Github #57680 [ run ] completed with state |
Signed-off-by: qgai <qgai@nvidia.com>
4 participants
Summary
Changes
tests/unittest/_torch/modules/moe/test_moe_comm.py: restructure MoE communication test case generation and grouped execution.tests/unittest/_torch/modules/moe/test_moe_comm.py: add valid receive-row filtering for combine reference accumulation.tests/unittest/_torch/modules/moe/test_moe_comm.py: make low-precision combine reference precomputed on worker GPUs and guard NVFP4 scale underflow.Test plan
pre-commit run --files tests/unittest/_torch/modules/moe/test_moe_comm.pyNCCL_NVLS_ENABLE=0 python3 -u -m pytest -q -s tests/unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm[NVLinkTwoSided_ep4_n18] --tb=short --disable-warnings --durations=0passes: 1 passed in 48.72s pytest time.NCCL_NVLS_ENABLE=0 python3 -u -m pytest -q -s tests/unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm --tb=short --disable-warnings --durations=0passes: 8 passed in 351.73s pytest time.