[https://nvbugs/6506918][fix] Reject the incompatible configuration up front in `Linear.__init__` with an… by trtllm-agent · Pull Request #16870 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-07-25T15:20:26Z

Summary

Root cause: Uneven-TP column-parallel Linear + gather_output=True hangs the NCCL AllGather because per-rank local buffer sizes differ but no sizes= argument is passed, and the underlying C++ AllGather with sizes is also broken for this case.
Fix: Reject the incompatible configuration up front in Linear.__init__ with an assert out_features % self.tp_size == 0 guard that fires when gather_output=True and tp_mode==COLUMN and tp_size>1. This restores the pre-TRTLLM-13117 behavior for exactly the code path that gather_output relies on (uniform NCCL AllGather), matches the test's pytest.raises(AssertionError) expectation, and leaves uneven TP working for gather_output=False VisualGen use cases. Also remove the corresponding waives.txt entry.
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6506918

Dev Engineer Review

Added upfront validation in Linear.__init__ to reject gathered outputs for uneven column-parallel sharding, preventing NCCL AllGather buffer-size mismatches and hangs.
Preserves uneven tensor parallelism when gather_output=False.
Removed the corresponding waiver entry; the test-list format and scope are consistent.
No public API declarations were changed.

QA Engineer Review

No test code was modified.
Removed the test_column_linear[2-unbalanced] waiver from tests/integration/test_lists/waives.txt.
CBTS coverage data is unavailable, so the verdict is needs follow-up.

Uneven-TP column-parallel Linear (TRTLLM-13117) produces different local output sizes per rank, which the NCCL AllGather collective cannot handle without per-rank sizes. Combined with gather_output=True, this deadlocks the collective on odd out_features (e.g. hidden_size=127 with tp_size=2), manifesting as "Test terminated unexpectedly" in tests/unittest/_torch/multi_gpu/test_linear.py::test_column_linear[2-unbalanced]. Reject the incompatible config up front in Linear.__init__ with a clear assertion instead of hanging mid-collective. Uneven TP still works for COLUMN Linears that keep the shard local (gather_output=False), which is what the VisualGen feature was designed for. Also remove the now-passing test from waives.txt. Signed-off-by: handongl <handongl@nvidia.com>

coderabbitai · 2026-07-25T15:23:47Z

Walkthrough

Linear now fails fast for unsupported gathered output with uneven column tensor parallelism, and the previously waived unbalanced column-linear integration test is enabled.

Changes

Linear column-sharding validation

Layer / File(s)	Summary
Configuration guard and test enablement `tensorrt_llm/_torch/modules/linear.py`, `tests/integration/test_lists/waives.txt`	`Linear.__init__` requires `out_features` to be divisible by `tp_size` for gathered column-parallel output, and removes the waiver for the unbalanced column-linear test.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: mingyanghao, belgarten-nv, stanleysun639

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is clearly related to the main change: rejecting the invalid Linear.init configuration up front.
Description check	✅ Passed	The description covers the root cause, fix, test plan, and bug link, which is enough for the template despite missing exact headings.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 2

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

Inline comments:
In `@tensorrt_llm/_torch/modules/linear.py`:
- Around line 3309-3317: Replace the assert in the gather_output column-parallel
validation with an explicit ValueError or RuntimeError, preserving the existing
condition and descriptive message. Ensure the unsupported uneven-TP
configuration is rejected unconditionally before reaching the allgather path,
including when Python runs with optimizations.
- Around line 3309-3317: Update the COLUMN gather_output validation near the
linear module initialization to account for override_tp_sharding before
accepting the configuration. Resolve the per-rank shard sizes after applying the
override and reject any uneven sizes, even when out_features is divisible by
tp_size, since the allgather path cannot provide sizes.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ad46ddb9-3697-4f7d-bc13-be768550bb4e

📥 Commits

Reviewing files that changed from the base of the PR and between cf44a1c and 5a64b05.

📒 Files selected for processing (2)

tensorrt_llm/_torch/modules/linear.py
tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

tests/integration/test_lists/waives.txt

coderabbitai · 2026-07-25T15:23:50Z

+        # gather_output requires a uniform NCCL AllGather; uneven column shards
+        # would produce mismatched local buffers and hang the collective. Reject
+        # this configuration up front rather than fail mid-collective.
+        if (gather_output and self.tp_size > 1
+                and self.tp_mode == TensorParallelMode.COLUMN):
+            assert out_features % self.tp_size == 0, (
+                f"gather_output=True is not supported with uneven TP. "
+                f"out_features={out_features} must be divisible by tp_size={self.tp_size}."
+            )


🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash python -O - <<'PY' def guard(): assert False, "must fail" guard() print("assertions are disabled under -O") PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 190

🏁 Script executed:

#!/bin/bash set -euo pipefail # Show the relevant section with line numbers. sed -n '3280,3345p' tensorrt_llm/_torch/modules/linear.py printf '\n--- search for related shard calculations ---\n' rg -n "override_tp_sharding|calculate_local_out_features|gather_output|tp_mode == TensorParallelMode.COLUMN|out_features % self.tp_size" tensorrt_llm/_torch/modules/linear.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 5656

🏁 Script executed:

#!/bin/bash set -euo pipefail printf '%s\n' '--- calculate_local_out_features ---' sed -n '3452,3495p' tensorrt_llm/_torch/modules/linear.py printf '\n%s\n' '--- gather_output runtime path ---' sed -n '3690,3735p' tensorrt_llm/_torch/modules/linear.py printf '\n%s\n' '--- override_tp_sharding references ---' rg -n "override_tp_sharding|calculate_local_out_features|calculate_local_in_features" tensorrt_llm/_torch/modules/linear.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 5487

Use an explicit exception here instead of assert. assert is stripped under python -O, so this unsupported gather_output/uneven-TP path can still reach allgather and hang. Raise ValueError or RuntimeError so the guard always runs.

🤖 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/linear.py` around lines 3309 - 3317, Replace the assert in the gather_output column-parallel validation with an explicit ValueError or RuntimeError, preserving the existing condition and descriptive message. Ensure the unsupported uneven-TP configuration is rejected unconditionally before reaching the allgather path, including when Python runs with optimizations.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash rg -n -C 4 'override_tp_sharding|gather_output|calculate_local_out_features' \ tensorrt_llm tests

Repository: NVIDIA/TensorRT-LLM

Length of output: 50375

🏁 Script executed:

#!/bin/bash set -euo pipefail sed -n '3320,3475p' tensorrt_llm/_torch/modules/linear.py printf '\n---\n' sed -n '3688,3722p' tensorrt_llm/_torch/modules/linear.py printf '\n---\n' sed -n '1,220p' tensorrt_llm/_torch/modules/gated_mlp.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 19337

🏁 Script executed:

#!/bin/bash set -euo pipefail rg -n -C 4 'def allgather|allgather\(' tensorrt_llm/_torch tensorrt_llm/distributed.py tensorrt_llm/distributed -g '!**/__pycache__/**'

Repository: NVIDIA/TensorRT-LLM

Length of output: 50375

🏁 Script executed:

#!/bin/bash set -euo pipefail rg -n -C 3 'override_tp_sharding=' tensorrt_llm/_torch printf '\n---\n' sed -n '120,220p' tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py printf '\n---\n' sed -n '220,320p' tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 19188

🏁 Script executed:

#!/bin/bash set -euo pipefail rg -n -C 2 'gather_output=True[^\n]*\n.*override_tp_sharding=|override_tp_sharding=.*\n.*gather_output=True' tensorrt_llm/_torch -U printf '\n---\n' rg -n -C 3 'tensor_parallel_mode=TensorParallelMode\.COLUMN.*gather_output=True|gather_output=True.*tensor_parallel_mode=TensorParallelMode\.COLUMN' tensorrt_llm/_torch -U

Repository: NVIDIA/TensorRT-LLM

Length of output: 157

Check override_tp_sharding before COLUMN gather_output (tensorrt_llm/_torch/modules/linear.py:3309-3317). An explicit per-rank override can still make local output widths differ even when out_features % tp_size == 0, and this allgather(output, self.mapping) path has no sizes argument. Validate the resolved shard sizes after applying the override, or reject uneven overrides here.

🤖 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/linear.py` around lines 3309 - 3317, Update the COLUMN gather_output validation near the linear module initialization to account for override_tp_sharding before accepting the configuration. Resolve the per-rank shard sizes after applying the override and reject any uneven sizes, even when out_features is divisible by tp_size, since the allgather path cannot provide sizes.

asfiyab-nvidia · 2026-07-28T17:40:49Z

Closing as bug https://nvbugs/6506918 is not reproducible

trtllm-agent requested review from a team as code owners July 25, 2026 15:20

trtllm-agent assigned QiJune Jul 25, 2026

trtllm-agent requested review from JunyiXu-nv, YihuiLu512, brnguyen2, crazydemo, schetlur-nv, yiqingy0 and yuanjingx87 July 25, 2026 15:20

github-actions Bot assigned trtllm-agent Jul 25, 2026

coderabbitai Bot reviewed Jul 25, 2026

View reviewed changes

coderabbitai Bot mentioned this pull request Jul 26, 2026

[https://nvbugs/6501376][fix] Test-only fix — drop the if hidden_size % 2 != 0: with pytest.raises(...)… #16844

Merged

2 tasks

asfiyab-nvidia closed this Jul 28, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[https://nvbugs/6506918][fix] Reject the incompatible configuration up front in `Linear.init` with an… - #16870

[https://nvbugs/6506918][fix] Reject the incompatible configuration up front in `Linear.init` with an…#16870
trtllm-agent wants to merge 1 commit into
NVIDIA:mainfrom
tensorrt-cicd:repair-bot-bug6506918

trtllm-agent commented Jul 25, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 25, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

coderabbitai Bot Jul 25, 2026

Uh oh!

asfiyab-nvidia commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

trtllm-agent commented Jul 25, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Dev Engineer Review

QA Engineer Review

Uh oh!

coderabbitai Bot commented Jul 25, 2026

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Jul 25, 2026

Choose a reason for hiding this comment

Uh oh!

asfiyab-nvidia commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

trtllm-agent commented Jul 25, 2026 •

edited by coderabbitai Bot

Loading