From 33d3d6469cf5db60d94d64bce35b70dc6261d36b Mon Sep 17 00:00:00 2001 From: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com> Date: Mon, 18 May 2026 11:53:12 -0700 Subject: [PATCH 1/2] [TRTLLM-12762][fix] Enable multi-node TP for MiniMax-M2 The QK-norm path in MiniMaxRMSNorm uses an IPC-based fused all-reduce kernel that is only available between GPUs with peer-to-peer access, preventing MiniMax-M2 from running with tensor parallelism that spans multiple nodes. Additionally, the k_norm weight loader does not replicate heads when num_kv_heads < tp_size, so checkpoints fail to load in typical cross-node configurations (e.g. 8 KV heads with TP=16). This change: - Detects cross-node TP via can_access_peer(mapping) at construction and caches the result on MiniMaxRMSNorm. - 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. 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, replicate at the head level using repeat_interleave before passing to load_weight_shard. This mirrors duplicate_kv_weight behavior for k_proj/v_proj. Intra-node TP behavior is unchanged: the fast IPC-based fused kernel is still used when can_access_peer(mapping) is true. Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm2.py | 41 ++++++++++++++++++- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm2.py b/tensorrt_llm/_torch/models/modeling_minimaxm2.py index 944f20ec77f3..b4cac333d281 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm2.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm2.py @@ -19,6 +19,7 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._ipc_utils import can_access_peer from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType from tensorrt_llm.mapping import Mapping @@ -119,7 +120,13 @@ def forward( # We use all_reduce across all tp gpus to get the rms norm variance sum class MiniMaxRMSNorm(nn.Module): def __init__( - self, *, hidden_size: int, eps: float, mapping: Mapping, dtype: torch.dtype = torch.bfloat16 + self, + *, + hidden_size: int, + eps: float, + mapping: Mapping, + dtype: torch.dtype = torch.bfloat16, + head_dim: Optional[int] = None, ): super().__init__() self.mapping = mapping @@ -128,14 +135,30 @@ def __init__( self.hidden_size = hidden_size self.eps = eps self.dtype = dtype + self.head_dim = head_dim + self.is_p2p_supported = can_access_peer(mapping) self.all_reduce = AllReduce(mapping=self.mapping, strategy=AllReduceStrategy.NCCL) self.minimax_all_reduce_rms = MiniMaxAllReduceRMS(mapping=self.mapping) def load_weights(self, weights: List[Dict]): assert len(weights) == 1 + src = weights[0]["weight"] + # When num_total_heads < tp_size (e.g. 8 KV heads, tp=16), the checkpoint weight + # [num_total_heads * head_dim] is smaller than what TP sharding expects + # [tp_size * local_hidden_size]. Replicate at the head level before sharding, + # consistent with how duplicate_kv_weight handles k_proj/v_proj. + 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 + reps = self.mapping.tp_size // num_total_heads + src = ( + src.reshape(num_total_heads, self.head_dim) + .repeat_interleave(reps, dim=0) + .reshape(-1) + ) weight = load_weight_shard( - weights[0]["weight"], + src, tensor_parallel_size=self.mapping.tp_size, tensor_parallel_rank=self.mapping.tp_rank, tensor_parallel_mode=TensorParallelMode.COLUMN, @@ -144,6 +167,15 @@ def load_weights(self, weights: List[Dict]): def forward(self, hidden_states: torch.Tensor): hidden_states = hidden_states.contiguous() + if not self.is_p2p_supported: + # Inter-node TP: IPC is unavailable, fall back to NCCL all-reduce of + # partial sum-of-squares followed by local RMS normalization. + hidden_f32 = hidden_states.float() + local_sum_sq = hidden_f32.pow(2).sum(-1, keepdim=True) + total_sum_sq = self.all_reduce(local_sum_sq) + total_hidden = self.hidden_size * self.mapping.tp_size + rms_inv = torch.rsqrt(total_sum_sq / total_hidden + self.eps) + return (hidden_f32 * rms_inv).to(hidden_states.dtype) * self.weight rms_norm_out = self.minimax_all_reduce_rms(hidden_states, self.weight, self.eps) return rms_norm_out @@ -188,12 +220,14 @@ def __init__( eps=config.rms_norm_eps, mapping=self.qkv_proj.mapping, dtype=config.torch_dtype, + head_dim=self.head_dim, ) self.k_norm = MiniMaxRMSNorm( hidden_size=self.kv_size, eps=config.rms_norm_eps, mapping=self.qkv_proj.mapping, dtype=config.torch_dtype, + head_dim=self.head_dim, ) else: self.q_norm = RMSNorm( @@ -209,6 +243,9 @@ def __init__( def apply_qk_norm(self, q, k): if self.qkv_proj.mapping.tp_size > 1: + if not self.q_norm.is_p2p_supported: + # Inter-node TP: fall back to separate per-tensor NCCL-based norm. + return self.q_norm(q), self.k_norm(k) q = q.contiguous() k = k.contiguous() q, k = self.q_norm.minimax_all_reduce_rms.forward_qk( From 1064acec336c0007ff3a7cd79519c08289dd7ea1 Mon Sep 17 00:00:00 2001 From: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com> Date: Wed, 20 May 2026 13:10:26 -0400 Subject: [PATCH 2/2] [TRTLLM-12762][fix] Validate head/TP divisibility in MiniMaxRMSNorm weight replication Assert that the checkpoint weight size is divisible by head_dim and that tp_size is divisible by num_total_heads before head-level replication in `MiniMaxRMSNorm.load_weights`. Without these checks, integer-truncated `reps = tp_size // num_total_heads` could silently produce fewer heads than `tp_size`, leading to wrong shard sizes downstream. Signed-off-by: Pietro Cicotti <5833013+pcicotti@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_minimaxm2.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm2.py b/tensorrt_llm/_torch/models/modeling_minimaxm2.py index b4cac333d281..754962ab5ebc 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm2.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm2.py @@ -150,7 +150,14 @@ def load_weights(self, weights: List[Dict]): # consistent with how duplicate_kv_weight handles k_proj/v_proj. full_size = self.mapping.tp_size * self.hidden_size if src.shape[0] < full_size and self.head_dim is not None: + assert src.shape[0] % self.head_dim == 0, ( + f"checkpoint weight size {src.shape[0]} is not divisible by head_dim {self.head_dim}" + ) num_total_heads = src.shape[0] // self.head_dim + assert self.mapping.tp_size % num_total_heads == 0, ( + f"tp_size {self.mapping.tp_size} must be divisible by num_total_heads {num_total_heads} " + f"for head-level weight replication" + ) reps = self.mapping.tp_size // num_total_heads src = ( src.reshape(num_total_heads, self.head_dim)