Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesThe PR migrates distributed AllReduce checkpoint testing to MPI/CUDA launcher setup and updates bf16/nvfp4 MoE tests to pass top-k routing fields with explicit dtypes. Distributed AllReduce checkpoint test
MoE routing contract tests
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the TRT-LLM all-reduce checkpoint test to use MPI-based distributed initialization instead of spawning processes manually. The review feedback suggests two improvements to make the test more robust in diverse environments: adding "PMI_SIZE" to the list of environment variables checked for the distributed world size, and using "torch.cuda.current_device()" to determine the local GPU device instead of assuming a simple round-robin mapping of ranks.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/comm/test_trtllm_allreduce_checkpoint.py (1)
11-12: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winUse
PMI_SIZEinstead ofMPI_LOCALNRANKSfor global world size.
MPI_LOCALNRANKSrepresents the local number of ranks on a single node (used by MPICH/MVAPICH/Intel MPI), not the global world size. If this test is executed in a multi-node environment, relying onMPI_LOCALNRANKScould underestimate the world size and cause the test to skip incorrectly.Consider replacing it with
PMI_SIZE, which accurately reflects the global world size across these MPI implementations.♻️ Proposed fix
"""Return world size from launcher env vars without importing mpi4py.""" - for key in ("SLURM_NTASKS", "WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "MPI_LOCALNRANKS"): + for key in ("SLURM_NTASKS", "WORLD_SIZE", "OMPI_COMM_WORLD_SIZE", "PMI_SIZE"):🤖 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/comm/test_trtllm_allreduce_checkpoint.py` around lines 11 - 12, Update the launcher environment variable list in the world-size helper to use PMI_SIZE instead of MPI_LOCALNRANKS, preserving the existing lookup order and behavior while ensuring the value represents the global world size.
🤖 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/comm/test_trtllm_allreduce_checkpoint.py`:
- Around line 11-12: Update the launcher environment variable list in the
world-size helper to use PMI_SIZE instead of MPI_LOCALNRANKS, preserving the
existing lookup order and behavior while ensuring the value represents the
global world size.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 45e275de-a93b-4a25-b983-9c0a8f885225
📥 Commits
Reviewing files that changed from the base of the PR and between 7a4c24b and c09470fd19c08b0613d5476ec47337174b0aea1a.
📒 Files selected for processing (1)
tests/comm/test_trtllm_allreduce_checkpoint.py
Replace mp.spawn + gloo with mpirun-compatible pattern using setup_mpi_and_cuda() and init_torch_distributed_from_mpi() (nccl), matching the style of other multi-node comm tests. Key changes: - Remove mp.spawn/_run_worker/_free_port; test body runs directly in each MPI rank's pytest process - Add _distributed_world_size() to read world size from launcher env vars (SLURM/torchrun/OpenMPI) without importing mpi4py, enabling correct skipif evaluation in multi-node single-GPU-per-node setups - Fix skipif: drop device_count<2 guard (wrong for multi-node 1-GPU/node); skip only when no CUDA AND no distributed launcher detected - Generalize expected values from hardcoded 3/5 (world_size=2 only) to world_size*(world_size+1)//2 and +world_size Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
…ed API selected_experts → topk_ids, final_scales → topk_weights following the rename in MoEActivationPack (fused_moe/api.py, gh flashinfer-ai#2425 naming convention). Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
MPICH, Intel MPI, and SLURM with PMI expose world size via PMI_SIZE. Without it, tests would be incorrectly skipped on common HPC clusters. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
setup_mpi_and_cuda() already calls torch.cuda.set_device(), so querying torch.cuda.current_device() is more robust than computing local_rank manually, which breaks under custom SLURM GPU bindings or non-round-robin rank-to-GPU mappings. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
9ab0a00 to
6a1caf6
Compare
…llreduce_checkpoint Revert our mpirun refactor in favor of the fix in PR flashinfer-ai#4034 by the test owner, who chose to keep the mp.spawn approach. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
…/main state Align with PR flashinfer-ai#4034 which fixes the mp.spawn module import issue via sys.path instead of the mpirun refactor we had attempted. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
nvamyt
changed the title
|
Closing as "already addressed by #4026# |
3 participants
Update MoEActivationPack field accordingly
Key changes:
Rename selected_experts → topk_ids and final_scales → topk_weights in MoEActivationPack constructor calls inside test_split_fused_moe_kernel_vs_reference.py, to match the field rename introduced in flashinfer/fused_moe/api.py.
#3892
📌 Description
🔍 Related Issues
#4065
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Summary of Changes
Tests
Bug Fixes / Test Fixes
topk_idsandtopk_weights.