From 5a64b05a2fd7cfe5c4d6269764202e116ad03624 Mon Sep 17 00:00:00 2001 From: handongl Date: Sat, 25 Jul 2026 08:12:35 -0700 Subject: [PATCH] [nvbugs/6506918][fix] Assert gather_output requires even TP shard 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 --- tensorrt_llm/_torch/modules/linear.py | 10 ++++++++++ tests/integration/test_lists/waives.txt | 1 - 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 10504cf6ecd7..e6ef13820f24 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -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}." + ) + # 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 diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index b498d3553f2e..68a73f125ac5 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -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)