[TRTLLM-14054][fix] LMHead: pass true full dims to Linear for quantized TP>1 - #16352
Conversation
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>
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthrough
ChangesLMHead sizing
Estimated code review effort: 1 (Trivial) | ~5 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/embedding.py (1)
61-73: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd regression coverage for tensor-parallel feature sizing.
The calculation matches
Linear’s split-dimension contract, but this fix should include tests that verify quantized COLUMN TP2 allocations use the true full hidden dimension and padded vocabulary dimension, while TP1 and dense LMHead shapes remain unchanged.🤖 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/modules/embedding.py` around lines 61 - 73, The updated sizing logic in the embedding module lacks regression coverage. Add tests for quantized LMHead/embedding construction verifying COLUMN tensor parallelism at TP2 uses the true full hidden dimension and padded vocabulary size, while TP1 and dense LMHead cases retain their existing shapes.
🤖 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 `@tensorrt_llm/_torch/modules/embedding.py`:
- Around line 61-73: The updated sizing logic in the embedding module lacks
regression coverage. Add tests for quantized LMHead/embedding construction
verifying COLUMN tensor parallelism at TP2 uses the true full hidden dimension
and padded vocabulary size, while TP1 and dense LMHead cases retain their
existing shapes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c27d79cc-cf65-4c19-b6b6-614f3d0f07df
📒 Files selected for processing (1)
tensorrt_llm/_torch/modules/embedding.py
|
PR_Github #59093 [ run ] triggered by Bot. Commit: |
kaiyux
changed the title
|
PR_Github #59093 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59143 [ run ] triggered by Bot. Commit: |
|
PR_Github #59143 [ run ] completed with state |
3 participants
Description
LMHead.__init__passedlocal_{in,out}_features * tp_sizefor both dims toLinear.__init__, butLinearonly re-shards the dim selected bytensor_parallel_mode— the unsharded dim arrived tp_size-× inflated.The dense path never hit this because
LMHeadre-creates a dense weight from corrected local dims aftersuper().__init__. The quantized path, however, sizes its (packed) weight and scales insideLinear.__init__viaquant_method.create_weights, so it saw the inflated dim: under COLUMN (vocab-parallel) TP2, the NVFP4-quantizedlm_headof Qwen3.6-35B-A3B-NVFP4 got its packed weight allocated as[vocab/2, hidden]instead of[vocab/2, hidden/2], and weight load failed on every rank:Fix: pass
local * tp_sizeonly on the sharded dim (it carries the ceil-div vocab/hidden padding); pass the true full size on the other dim.This unblocks TP>1 for any model with a quantized
lm_head(any quant algo whose weights are created inLinear.__init__); it was hit while bringing up Qwen3.6-35B-A3B-NVFP4 under TP2 for the perf work in #16313 / #16314, and the TP2 measurements in #16353 were taken with this fix applied.Test Coverage
lm_head, vocab-parallel): TP2 previously failed weight load on both ranks with theRuntimeErrorabove; with the fix the server loads and serves, and both ranks'lm_headshards were verified to hold the correct packed shape and calibrated scales ([124160, 1024]packed +124160×128scales per rank).[248320, 1024]+248320×128scales, byte-identical to before the fix.tests/unittest/_torch/multi_gpu/test_embedding.pyand are unaffected (the dense weight is re-created from local dims aftersuper().__init__).PR Checklist
PR description clearly explains what and why. If using CodeRabbit's 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 (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
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.🤖 Generated with Claude Code
Summary by CodeRabbit