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

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)
 

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)

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,