[https://nvbugs/6572835][fix] Use HashStore for single-rank MegaMoE tests

[Barry-Delaney]

Dev Engineer Review

• Single-rank MegaMoE initialization now uses torch.distributed.HashStore.
• This avoids env:// rendezvous failures when MASTER_PORT is occupied.
• Multi-rank initialization remains unchanged.
• The MEGAMOE_CUTEDSL waiver removal is valid and enables existing CI coverage.
• No duplicate waiver entry is present.
• CI failures require investigation before reruns.
• Verdict: needs follow-up.

QA Engineer Review

• Updated MegaMoE initialization helpers in:
  • test_moe_backend.py
  • test_moe_module.py
  • tests/microbenchmarks/bench_moe/utils.py
• No test functions were added, modified, or removed.
• Existing coverage includes:
  • test_configurable_moe_single_gpu -k "MEGAMOE_CUTEDSL" in l0_b200.yml.
  • test_moe_backend -k "MEGAMOE_CUTEDSL" in l0_b200.yml.
• Removed the single-GPU MEGAMOE_CUTEDSL waiver from tests/integration/test_lists/waives.txt.
• Multi-GPU MEGAMOE_CUTEDSL remains covered and waived separately.
• CI results remain inconclusive because multiple runs failed, including failures related to the required ci: full pre-merge approved label.
• Verdict: needs follow-up.

Description

The single-rank MegaMoE unit-test helpers initialize an NCCL process group through the default env:// rendezvous. In CI this inherits MASTER_PORT, so an unrelated listener occupying that port causes c10d TCPStore to fail with EADDRINUSE before the MegaMoE test executes.

Use an in-process torch.distributed.HashStore when world_size == 1. Single-rank tests require a ProcessGroup but do not need cross-process rendezvous, so no network socket is necessary. Multi-rank tests retain the existing shared TCP rendezvous path.

Apply the same behavior to both MegaMoE single-rank test helpers and remove the NVBug 6572835 waiver so the affected test runs in CI again.

Test Coverage

• Controlled A/B with the inherited MASTER_PORT deliberately occupied: the original helper reproduced EADDRINUSE, while both patched helpers initialized NCCL and completed all_reduce
• Python compile check for both modified test modules
• pre-commit run --files tests/unittest/_torch/modules/moe/test_moe_module.py tests/unittest/_torch/modules/moe/test_moe_backend.py tests/integration/test_lists/waives.txt
• git diff --check
• The unwaived MegaMoE-CuteDSL target test will run in PR CI

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit 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.

• No public API, dependency, ownership, or architecture changes are introduced.

• 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.

[Barry-Delaney]

/bot run --disable-fail-fast

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4da9a42b-7295-4511-a195-057a70735499

📥 Commits

Reviewing files that changed from the base of the PR and between 4ae87a5 and 41512e6.

📒 Files selected for processing (1)

• tests/integration/test_lists/waives.txt

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/integration/test_lists/waives.txt

Included review availability: Your plan includes up to 12 reviews per rolling hour; 10 remain after this review.

---

Walkthrough

The changes update single-process MegaMoE tests and benchmarks to use local HashStore-based NCCL initialization. They also remove an unused import and remove the related test waiver.

Changes

MegaMoE test setup

Layer / File(s)	Summary
Single-process distributed initialization tests/unittest/_torch/modules/moe/test_moe_backend.py, tests/unittest/_torch/modules/moe/test_moe_module.py, tests/microbenchmarks/bench_moe/utils.py	Single-rank MegaMoE tests and benchmarks now initialize NCCL with dist.HashStore(), rank 0, and world size 1. The module test also sets the CUDA device. Multi-rank setup remains unchanged. The unused os import is removed.
Test waiver removal tests/integration/test_lists/waives.txt	The waiver for the single-GPU MEGAMOЕ_CUTEDSL MegaMoE test is removed.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 41512

The change limits single-rank MegaMoE tests to in-process rendezvous while preserving multi-rank behavior, with no actionable merge-blocking risk remaining after normal checks and review.

Possibly related PRs

• NVIDIA/TensorRT-LLM#17776: Both changes update tests/integration/test_lists/waives.txt, but they target different tests.

Suggested reviewers: schetlur-nv, qijune, zhanruisunch

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required ticket and type format and clearly summarizes the HashStore fix for single-rank MegaMoE tests.
Description check	✅ Passed	The description explains the problem and solution, lists test coverage, and addresses the relevant checklist items for this test-only change.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[tensorrt-cicd]

PR_Github #65288 [ run ] triggered by Bot. Commit: 06ab6ee Link to invocation

[tensorrt-cicd]

PR_Github #65288 [ run ] completed with state SUCCESS. Commit: 06ab6ee
/LLM/main/L0_MergeRequest_PR pipeline #53064 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[brnguyen2]

Approving — the comments below are optional touch-ups, not blockers.

Right fix for the right reason — a single-rank group needs a ProcessGroup, not a rendezvous, so removing the socket entirely is stronger than probing a free port. Verified that dropping the RANK/WORLD_SIZE/LOCAL_RANK env writes in test_moe_backend.py is safe: both MegaMoE bootstraps (mega_moe_cute_dsl.py:782, mega_moe_deepgemm.py:392) early-return when the MPI world is 1, so nothing downstream reads them in these tests.

Two follow-ups below, neither blocking. Since the waiver removal is the point of the PR, please confirm the unwaived -k "MEGAMOE_CUTEDSL" case actually ran and passed in this PR's B200 stage before merging rather than only that the pipeline was green.

[Barry-Delaney]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65534 [ run ] triggered by Bot. Commit: 118a2b3 Link to invocation

[tensorrt-cicd]

PR_Github #65534 [ run ] completed with state FAILURE. Commit: 118a2b3
/LLM/main/L0_MergeRequest_PR pipeline #53272 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[Barry-Delaney]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65851 [ run ] triggered by Bot. Commit: 168f20c Link to invocation

[tensorrt-cicd]

PR_Github #65851 [ run ] completed with state FAILURE. Commit: 168f20c
/LLM/main/L0_MergeRequest_PR pipeline #53547 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Barry-Delaney]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66618 [ run ] triggered by Bot. Commit: 4ae87a5 Link to invocation

[tensorrt-cicd]

PR_Github #66618 [ run ] completed with state FAILURE. Commit: 4ae87a5
/LLM/main/L0_MergeRequest_PR pipeline #54239 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[Barry-Delaney]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66681 [ run ] triggered by Bot. Commit: 41512e6 Link to invocation

[tensorrt-cicd]

PR_Github #66681 [ run ] completed with state SUCCESS. Commit: 41512e6
/LLM/main/L0_MergeRequest_PR pipeline #54297 completed with status: 'SUCCESS'

CI Report

Link to invocation