From c82386a57b2c4adb86e58e6bb9914daabda0f52c Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:03:43 -0700 Subject: [PATCH] [None][fix] LMHead: pass true full dims to Linear for quantized TP>1 LMHead.__init__ passed local_{in,out}_features * tp_size for BOTH dims to Linear.__init__, but Linear only re-shards the dim selected by tensor_parallel_mode - the unsharded dim arrived tp_size-x inflated. Under COLUMN (vocab-parallel) TP2 the NVFP4-quantized lm_head of Qwen3.6-35B-A3B-NVFP4 got its packed weight allocated as [vocab/2, hidden] instead of [vocab/2, hidden/2], so weight load failed on every rank: RuntimeError: The size of tensor a (2048) must match the size of tensor b (1024) at non-singleton dimension 1 The dense path never hit this because LMHead re-creates a dense weight from corrected local dims after super().__init__; the quantized path sizes its (packed) weight and scales inside Linear.__init__ via quant_method.create_weights, so it saw the inflated dim. Fix: pass local*tp only on the sharded dim (it carries the ceil-div vocab/hidden padding); pass the true full size on the other dim. TP1 shapes are unchanged; quantized TP2 now allocates the correct per-rank shard ([124160, 1024] packed weight + 124160x128 scales per rank for the model above) and weight load succeeds. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- tensorrt_llm/_torch/modules/embedding.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index c3236464bb18..a707372590a2 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -58,9 +58,19 @@ def __init__( else: self.padding_size = 0 + # Linear re-shards only the dim selected by tensor_parallel_mode, so + # only that dim may be passed as local*tp (to carry the ceil-div + # padding); the other dim must be the true full size — Linear keeps it + # as-is, and quant_method.create_weights sizes the (packed) quantized + # weight and scales from it. + in_features_full = (local_in_features * tp_size if tensor_parallel_mode + == TensorParallelMode.ROW else embedding_dim) + out_features_full = (local_out_features * + tp_size if tensor_parallel_mode + == TensorParallelMode.COLUMN else num_embeddings) super().__init__( - local_in_features * tp_size, - local_out_features * tp_size, + in_features_full, + out_features_full, bias=False, dtype=dtype, mapping=mapping,