From 77df10278f8d3f253b54798342b765121a2465e0 Mon Sep 17 00:00:00 2001 From: handongl Date: Fri, 24 Jul 2026 05:17:05 -0700 Subject: [PATCH 1/4] [nvbugs/6501376][fix] Drop stale pytest.raises for uneven-TP row/column linear PR #14875 added VANILLA-mode uneven-TP support in tensorrt_llm/_torch/modules/linear.py via _auto_tp_sharding, so test_row_linear[2-unbalanced] (hidden=15) and test_column_linear [2-unbalanced] (hidden=127) no longer hit the divisibility assert. The pre-existing tests still expected an AssertionError on the odd branch, producing "Failed: DID NOT RAISE" in pre-merge CI. Unconditionally assert all ranks return True. Signed-off-by: handongl --- tests/unittest/_torch/multi_gpu/test_linear.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/tests/unittest/_torch/multi_gpu/test_linear.py b/tests/unittest/_torch/multi_gpu/test_linear.py index 8452baa19e7c..4db1c6b053ce 100644 --- a/tests/unittest/_torch/multi_gpu/test_linear.py +++ b/tests/unittest/_torch/multi_gpu/test_linear.py @@ -286,13 +286,8 @@ def test_column_linear(hidden_size, mpi_pool_executor): run_single_rank, *zip(*[(tensor_parallel_size, column_linear_forward, x, [l0_weight], hidden_size, dtype)] * 2)) - if hidden_size % 2 != 0: - with pytest.raises(AssertionError): - for r in results: - assert r is True - else: - for r in results: - assert r is True + for r in results: + assert r is True @pytest.mark.skipif(torch.cuda.device_count() < 2, @@ -311,13 +306,8 @@ def test_row_linear(hidden_size, mpi_pool_executor): run_single_rank, *zip(*[(tensor_parallel_size, row_linear_forward, x, [l0_weight], hidden_size, dtype)] * 2)) - if hidden_size % 2 != 0: - with pytest.raises(AssertionError): - for r in results: - assert r is True - else: - for r in results: - assert r is True + for r in results: + assert r is True @pytest.mark.skipif(torch.cuda.device_count() < 2, From 736fc59a5b37372c6bb4f5f38c6a7f3f466bd576 Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:34:29 +0900 Subject: [PATCH 2/4] [nvbugs/6501376][fix] Remove uneven linear test waivers Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6dc4776bb633..c3d432471b29 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -182,7 +182,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) @@ -352,7 +351,6 @@ unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_ unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" SKIP (https://nvbugs/6464169) unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) unittest/_torch/multi_gpu/test_linear.py::test_row_linear[2-balanced] SKIP (https://nvbugs/6507113) -unittest/_torch/multi_gpu/test_linear.py::test_row_linear[2-unbalanced] SKIP (https://nvbugs/6501394) unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2] SKIP (https://nvbugs/6501404) unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py -m "part0" SKIP (https://nvbugs/6490036) unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py -m "part1" SKIP (https://nvbugs/6490036) From 3c09b66a6d3406e5a945e8f46360f5c558de1e2d Mon Sep 17 00:00:00 2001 From: Lori Ren Date: Sat, 25 Jul 2026 19:25:22 -0700 Subject: [PATCH 3/4] [nvbugs/6501376][fix] Pass per-rank sizes to uneven-TP column linear allgather Signed-off-by: Lori Ren --- tensorrt_llm/_torch/modules/linear.py | 45 ++++++++++++++++++++++----- 1 file changed, 37 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 10504cf6ecd7..17ff6bca8f4a 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3314,14 +3314,17 @@ def __init__( raise ValueError( f"Uneven TP is not supported with QuantAlgo {_quant_algo}") self.tp_sharding = override_tp_sharding + self._tp_sharding_is_auto = False elif self.tp_size > 1 and self.tp_mode is not None \ and self.weights_loading_config.weight_mode == WeightMode.VANILLA \ and _quant_algo not in _uneven_tp_unsupported \ and not skip_create_weights_in_init: features = in_features if self.tp_mode == TensorParallelMode.ROW else out_features self.tp_sharding = self._auto_tp_sharding(features, quant_config) + self._tp_sharding_is_auto = True else: self.tp_sharding = None + self._tp_sharding_is_auto = False if self.tp_size > 1 and self.tp_mode is not None: features = in_features if self.tp_mode == TensorParallelMode.ROW else out_features assert features % self.tp_size == 0, ( @@ -3334,6 +3337,28 @@ def __init__( self.in_features = self.calculate_local_in_features(in_features) self.out_features = self.calculate_local_out_features(out_features) + # allgather with sizes=None requires every rank to hold the same shape, so an + # unevenly sharded COLUMN output has to declare the per-rank widths explicitly. + self.gather_output_sizes = None + if self.gather_output and self.tp_size > 1 \ + and self.tp_mode == TensorParallelMode.COLUMN: + if self._tp_sharding_is_auto: + sizes = [ + end - start for start, end in ( + self._auto_tp_sharding(out_features, quant_config, rank) + for rank in range(self.tp_size)) + ] + if len(set(sizes)) > 1: + self.gather_output_sizes = sizes + elif isinstance(self.tp_sharding, tuple) and \ + (self.tp_sharding[1] - self.tp_sharding[0]) * self.tp_size != out_features: + # Only this rank's range is known, so the other ranks' widths that + # allgather needs cannot be derived. + raise ValueError( + f"gather_output=True with an uneven override_tp_sharding " + f"{self.tp_sharding} is not supported (out_features=" + f"{out_features}, tp_size={self.tp_size}).") + if self.tp_mode == TensorParallelMode.COLUMN: reduce_output = False if self.mapping.enable_attention_dp else reduce_output @@ -3387,34 +3412,36 @@ def get_quant_method(self, quant_config: Optional[QuantConfig] = None): def _calc_shard(total, tp_size, rank): return (total // tp_size) * rank + min(total % tp_size, rank) - def _auto_tp_sharding(self, features, quant_config): + def _auto_tp_sharding(self, features, quant_config, rank=None): """Auto-generate tp_sharding tuple based on quant alignment requirements. VANILLA mode only. Fused modes (FUSED_QKV, FUSED_GATE_UP) require explicit override_tp_sharding because individual sub-weight sizes (Q vs K vs V; gate vs up) are not knowable here — they aren't always equal (e.g. GQA), and cross-rank consistency must be decided by the caller. + + `rank` defaults to this module's own TP rank; pass an explicit rank to query + another rank's range (e.g. to build the per-rank sizes an allgather needs). """ assert self.weights_loading_config.weight_mode == WeightMode.VANILLA, ( f"_auto_tp_sharding only supports VANILLA mode, got " f"{self.weights_loading_config.weight_mode}. Fused modes require " f"explicit override_tp_sharding.") + rank = self.tp_rank if rank is None else rank alignment = get_quant_method(quant_config).get_tp_alignment( self.tp_mode, quant_config) if alignment <= 1: # No alignment constraint — use standard element-level distribution - start = self._calc_shard(features, self.tp_size, self.tp_rank) - end = self._calc_shard(features, self.tp_size, self.tp_rank + 1) + start = self._calc_shard(features, self.tp_size, rank) + end = self._calc_shard(features, self.tp_size, rank + 1) else: # Distribute whole alignment-sized blocks across ranks assert features % alignment == 0, ( f"Feature dim ({features}) must be divisible by quant alignment " f"({alignment}) for TP sharding") num_blocks = features // alignment - block_start = self._calc_shard(num_blocks, self.tp_size, - self.tp_rank) - block_end = self._calc_shard(num_blocks, self.tp_size, - self.tp_rank + 1) + block_start = self._calc_shard(num_blocks, self.tp_size, rank) + block_end = self._calc_shard(num_blocks, self.tp_size, rank + 1) start = block_start * alignment end = block_end * alignment return (start, end) @@ -3703,7 +3730,9 @@ def forward( output = self.apply_linear(input, self.bias, lora_params, layer_idx) if self.gather_output: from ..distributed import allgather - output = allgather(output, self.mapping) + output = allgather(output, + self.mapping, + sizes=self.gather_output_sizes) else: output = self.apply_linear(input, self.bias, lora_params, layer_idx) From eb39108058f7fd062941e050a69b876641e9dd07 Mon Sep 17 00:00:00 2001 From: Lori Ren Date: Sat, 25 Jul 2026 19:52:54 -0700 Subject: [PATCH 4/4] [nvbugs/6501376][fix] Reject gather_output combined with override_tp_sharding Signed-off-by: Lori Ren --- tensorrt_llm/_torch/modules/linear.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 17ff6bca8f4a..40b95df52ba9 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3350,14 +3350,17 @@ def __init__( ] if len(set(sizes)) > 1: self.gather_output_sizes = sizes - elif isinstance(self.tp_sharding, tuple) and \ - (self.tp_sharding[1] - self.tp_sharding[0]) * self.tp_size != out_features: - # Only this rank's range is known, so the other ranks' widths that - # allgather needs cannot be derived. + elif self.tp_sharding is not None: + # allgather concatenates in rank order, so gather_output assumes rank i + # holds the i-th ascending, contiguous slice. _auto_tp_sharding + # guarantees that; an arbitrary override guarantees neither the widths + # nor the ordering, and `sizes` carries widths only — it cannot express + # a permuted or non-contiguous layout even when the widths are equal. raise ValueError( - f"gather_output=True with an uneven override_tp_sharding " - f"{self.tp_sharding} is not supported (out_features=" - f"{out_features}, tp_size={self.tp_size}).") + f"gather_output=True is not supported together with " + f"override_tp_sharding ({self.tp_sharding}); gather_output " + f"requires the rank-ordered contiguous layout that only " + f"automatic TP sharding provides.") if self.tp_mode == TensorParallelMode.COLUMN: reduce_output = False if self.mapping.enable_attention_dp else reduce_output