From 77df10278f8d3f253b54798342b765121a2465e0 Mon Sep 17 00:00:00 2001 From: handongl Date: Fri, 24 Jul 2026 05:17:05 -0700 Subject: [PATCH] [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,