Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
trtllm-agent wants to merge 1 commit into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions tensorrt_llm/_torch/modules/linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3306,6 +3306,16 @@ def __init__(
f"Invalid tp_mode {self.tp_mode!r}; expected ROW, COLUMN, or None."
)

# 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}."
)
Comment on lines +3309 to +3317

@coderabbitai coderabbitai Bot Jul 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


# Init TP sharding either from override or auto generated
_uneven_tp_unsupported = {QuantAlgo.NVFP4_ARC}
_quant_algo = quant_config.quant_algo if quant_config else None
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,6 @@ full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[
full:B300/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760)
full:DGX_B200/accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] SKIP (https://nvbugs/6501837)
full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm SKIP (https://nvbugs/6506920)
full:DGX_H100/unittest/_torch/multi_gpu/test_linear.py::test_column_linear[2-unbalanced] SKIP (https://nvbugs/6506918)
full:DGX_H200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[disable_skip_indexer] SKIP (https://nvbugs/6476233)
full:DGX_H200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[latency_default] SKIP (https://nvbugs/6476233)
full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy SKIP (https://nvbugs/6276923)
Expand Down
Loading