diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 10504cf6ecd7..40b95df52ba9 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,31 @@ 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 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 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 @@ -3387,34 +3415,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 +3733,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) 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) 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,