[TRTLLM-12762][fix] Enable multi-node TP for MiniMax-M2

[pcicotti]

Background

MiniMaxM2's QK-norm path uses an IPC-based fused all-reduce kernel
(MiniMaxAllReduceRMS) for the variance sum-of-squares reduction. IPC
is only available between GPUs with peer-to-peer access, so MiniMax-M2
cannot run with tensor parallelism that spans more than one node.
Additionally, MiniMaxRMSNorm.load_weights does not replicate heads
when num_kv_heads < tp_size, so the k_norm weights fail to load in
typical cross-node configurations (e.g. 8 KV heads with TP=16).

JIRA: TRTLLM-12762

Summary

Detect cross-node TP via can_access_peer(mapping) at construction.
When peer access is unavailable, MiniMaxRMSNorm.forward falls back to
an NCCL all-reduce of the partial sum-of-squares followed by a local
RMS normalization, and MiniMaxM2Attention.apply_qk_norm falls back to
separate per-tensor q_norm(q) / k_norm(k) calls instead of the
fused IPC kernel. In MiniMaxRMSNorm.load_weights, when the checkpoint
tensor is smaller than tp_size * hidden_size, the weight is replicated
at the head level using repeat_interleave before passing to
load_weight_shard, mirroring duplicate_kv_weight for k_proj/v_proj.

Impact

• Intra-node TP behavior is unchanged: the fast IPC-based fused kernel
  is still used when can_access_peer(mapping) is true.
• Enables MiniMax-M2 deployment on multi-node TP (e.g. TP=16 across 2
  nodes) with num_kv_heads < tp_size.
• No new dependencies or API changes; no perf change to single-node path.

Summary by CodeRabbit

Release Notes

• Improvements
  • Enhanced distributed tensor support with automatic peer-to-peer capability detection
  • Added fallback optimization ensuring compatibility across diverse hardware configurations

[pcicotti]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR extends MiniMaxRMSNorm to detect peer-to-peer (IPC) availability and fall back to NCCL all-reduce when direct peer communication is unavailable. The constructor accepts an optional head_dim parameter for weight shaping. The attention layer now passes this dimension and conditionally applies separate or combined QK normalization based on p2p support.

Changes

RMS Normalization P2P/IPC Support

Layer / File(s)	Summary
Import and RMSNorm constructor update tensorrt_llm/_torch/models/modeling_minimaxm2.py (lines 22, 123–129)	Added can_access_peer import and extended MiniMaxRMSNorm.__init__ with optional head_dim keyword parameter to track head-level weight dimensions during tensor parallel sharding.
P2P support detection and weight replication tensorrt_llm/_torch/models/modeling_minimaxm2.py (lines 138–161)	Initialized p2p support detection using can_access_peer, prepared NCCL all-reduce buffers, and implemented conditional weight replication in load_weights to expand checkpoint tensors when head_dim is smaller than the TP-sharded dimension.
RMSNorm forward with NCCL all-reduce fallback tensorrt_llm/_torch/models/modeling_minimaxm2.py (lines 170–178)	Modified forward to branch on p2p availability: when unsupported, compute RMS via partial sum-of-squares reduction, all-reduce across ranks, inverse computation, and scaled output; otherwise use the existing p2p path.
Attention integration with p2p-aware QK normalization tensorrt_llm/_torch/models/modeling_minimaxm2.py (lines 223–230, 246–248)	Updated MiniMaxM2Attention to pass head_dim to q_norm and k_norm instances during TP sharding. Modified apply_qk_norm to conditionally apply separate q/k norms when p2p is unsupported, triggering the NCCL fallback path instead of the combined forward_qk path.

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: enabling multi-node tensor parallelism support for MiniMax-M2 by fixing IPC limitations.
Description check	✅ Passed	The PR description comprehensively covers background, summary, and impact sections that explain the issue, solution, and consequences. It addresses all critical aspects of the change.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/models/modeling_minimaxm2.py (1)

144-166: ⚠️ Potential issue | 🟠 Major | ⚖️ Poor tradeoff

Add validation for weight replication divisibility constraints.

The weight replication logic uses integer division without validation:

• Line 153: num_total_heads = src.shape[0] // self.head_dim assumes src.shape[0] is divisible by head_dim
• Line 154: reps = self.mapping.tp_size // num_total_heads assumes tp_size is divisible by num_total_heads

If these constraints don't hold, the replication will silently produce incorrect weights, leading to model corruption. For example, with tp_size=16 and num_total_heads=9, reps would be 1 instead of the required ~1.78, producing only 9 heads instead of 16.

🛡️ Proposed validation

 full_size = self.mapping.tp_size * self.hidden_size
 if src.shape[0] < full_size and self.head_dim is not None:
     num_total_heads = src.shape[0] // self.head_dim
+    if src.shape[0] % self.head_dim != 0:
+        raise ValueError(
+            f"Checkpoint weight size {src.shape[0]} is not divisible by head_dim {self.head_dim}"
+        )
     reps = self.mapping.tp_size // num_total_heads
+    if self.mapping.tp_size % num_total_heads != 0:
+        raise ValueError(
+            f"TP size {self.mapping.tp_size} must be divisible by num_total_heads {num_total_heads} "
+            f"for weight replication. Consider using a different TP configuration."
+        )
     src = (

🤖 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 `@tensorrt_llm/_torch/models/modeling_minimaxm2.py` around lines 144 - 166, In
load_weights, validate divisibility before computing num_total_heads and reps:
ensure src.shape[0] is divisible by self.head_dim (so num_total_heads =
src.shape[0] // self.head_dim is exact) and ensure self.mapping.tp_size is
divisible by num_total_heads (so reps = self.mapping.tp_size // num_total_heads
is exact); if either check fails, raise a clear ValueError mentioning
load_weights, src.shape[0], self.head_dim, num_total_heads and
self.mapping.tp_size so the caller knows why replication cannot proceed safely.

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

Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_minimaxm2.py`:
- Around line 144-166: In load_weights, validate divisibility before computing
num_total_heads and reps: ensure src.shape[0] is divisible by self.head_dim (so
num_total_heads = src.shape[0] // self.head_dim is exact) and ensure
self.mapping.tp_size is divisible by num_total_heads (so reps =
self.mapping.tp_size // num_total_heads is exact); if either check fails, raise
a clear ValueError mentioning load_weights, src.shape[0], self.head_dim,
num_total_heads and self.mapping.tp_size so the caller knows why replication
cannot proceed safely.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: bda3c719-f8ea-4010-ae80-3d7af8de3a9c

📥 Commits

Reviewing files that changed from the base of the PR and between 0435722 and 33d3d64.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/models/modeling_minimaxm2.py

[tensorrt-cicd]

PR_Github #49217 [ run ] triggered by Bot. Commit: 33d3d64 Link to invocation

[tensorrt-cicd]

PR_Github #49217 [ run ] completed with state SUCCESS. Commit: 33d3d64
/LLM/main/L0_MergeRequest_PR pipeline #38889 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

[pcicotti]

/bot run --disable-fail-fast

CI #38889 failures look unrelated to this PR (only touches tensorrt_llm/_torch/models/modeling_minimaxm2.py); addressed CodeRabbit's review with 1064ace (divisibility asserts in MiniMaxRMSNorm.load_weights).

H100_PCIe-AutoDeploy-1 (5 failures, 4 already waived):

• test_e2e[aggr_upload-llama3_1_8b_fp8_ad_hopper-llama3_1_8b_ad_ws1_1k1k] — waived, https://nvbugs/6192201
• test_vision_attention_matches_reference, test_vision_block_matches_reference, test_vlm_wrapper_delta_is_request_scoped_no_cross_call_leakage (qwen3_5_moe) — waived, https://nvbugs/6189450
• test_unittests_v2[unittest/auto_deploy/singlegpu/models] — parent aggregator; its only failing children are the 3 waived qwen3_5_moe tests above (aggregator doesn't inherit child waives).

DGX_B200-PyTorch-3 (2 failures):

• test_attention_no_cache[...num_heads_16, head_dim_128, num_layers_2...]
• test_attention_no_cache[...num_heads_16, head_dim_72, num_layers_16...]

Both "Test terminated unexpectedly" (worker crash). Historically flaky — cf. #3921 / https://nvbugs/5247232.

B300-PyTorch-1, GB300-PyTorch-1: ABORTED, infra.

Re-running with fail-fast disabled.

[tensorrt-cicd]

PR_Github #49466 Bot args parsing error: Failed to parse bot args

Link to invocation

[pcicotti]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49468 [ run ] triggered by Bot. Commit: 1064ace Link to invocation

[tensorrt-cicd]

PR_Github #49468 [ run ] completed with state SUCCESS. Commit: 1064ace
/LLM/main/L0_MergeRequest_PR pipeline #39109 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation