From f5dd26d5b68045358d6c72c32ff31563d5586caa Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Thu, 12 Nov 2026 17:07:37 -0400 Subject: [PATCH 1/2] fix(pytorch): initialize tp_group_initialized in TransformerEngineBaseModule Linear with tp_size > 1 and no tp_group never called set_tensor_parallel_group, so the forward-time guard read an undefined attribute and raised AttributeError instead of the intended RuntimeError. Signed-off-by: Alex Castillo --- tests/pytorch/test_sanity.py | 18 ++++++++++++++++++ transformer_engine/pytorch/module/base.py | 1 + 2 files changed, 19 insertions(+) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 2627de023ac..3f366b512db 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -1651,3 +1651,21 @@ def check_weights(): with autocast(enabled=with_quantization, recipe=quantization_recipe): y = module(x, **kwargs) check_weights() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_linear_tp_group_not_set_raises_runtime_error(): + """Forward with tp_size > 1 and no TP group must raise the intended RuntimeError.""" + model = te.Linear( + 16, + 16, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + tp_size=2, + tp_group=None, + parallel_mode="column", + ) + x = torch.ones(8, 16, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="Tensor parallel group not initialized"): + model(x) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index bf4f9bba1d5..ebc0d13895f 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1023,6 +1023,7 @@ def __init__(self, name: Optional[str] = None) -> None: self.fp8_meta_tensors_initialized = False self.quantizers = {"scaling_fwd": [], "scaling_bwd": []} self.tp_group = None + self.tp_group_initialized = False self.tp_size = 1 self.sequence_parallel = False self.param_init_meta = {} From 3ab6fbf32114a0ec852f63842743c7486b3d94e7 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Fri, 13 Nov 2026 17:09:45 -0800 Subject: [PATCH 2/2] Remove the tensor parallel group initialization test Signed-off-by: Przemek Tredak --- tests/pytorch/test_sanity.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 3f366b512db..2627de023ac 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -1651,21 +1651,3 @@ def check_weights(): with autocast(enabled=with_quantization, recipe=quantization_recipe): y = module(x, **kwargs) check_weights() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_linear_tp_group_not_set_raises_runtime_error(): - """Forward with tp_size > 1 and no TP group must raise the intended RuntimeError.""" - model = te.Linear( - 16, - 16, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - tp_size=2, - tp_group=None, - parallel_mode="column", - ) - x = torch.ones(8, 16, device="cuda", dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="Tensor parallel group not initialized"): - model(x)