Don't use TP when tensor_parallel_degree is 1 (#3636)

Co-authored-by: Eitan Turok <[email protected]>

Author: eitanturok

composer/core/state.py

@@ -612,6 +612,12 @@ def _validate_parallelism_configs(self):
                     'Tensor parallelism (TP) currently requires FSDP with use_orig_params=True, '
                     'which is the default and recommended setting.',
                 )
+            if self.tp_config.tensor_parallel_degree == 1:
+                warnings.warn(
+                    'Received tensor_parallel_degree of 1, which is a no-op. Tensor parallelism will not be used.',
+                    UserWarning,
+                )
+                self.tp_config = None
 
         # Load monolith rank0 only
         if self.load_monolith_rank0_only:

tests/trainer/test_tp.py

@@ -1,11 +1,14 @@
 # Copyright 2022 MosaicML Composer authors
 # SPDX-License-Identifier: Apache-2.0
 
+import contextlib
+
 import pytest
 import torch
 from packaging import version
 from torch.utils.data import DataLoader
 
+from composer.optim import DecoupledSGDW
 from composer.trainer.trainer import Trainer
 from composer.utils import dist
 from tests.common import (
@@ -17,12 +20,14 @@
 
 @pytest.mark.gpu
 @world_size(4)
-@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('2.3'), reason='requires PyTorch 2.3+')
 @pytest.mark.filterwarnings(r'ignore:.*\(TP\) is experimental.*:FutureWarning')
-def test_tp_train(world_size: int):
+@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('2.3'), reason='requires PyTorch 2.3+')
+@pytest.mark.parametrize('tensor_parallel_degree', [1, 2])
+def test_tp_train(world_size: int, tensor_parallel_degree: int):
     from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
 
     model = SimpleModel()
+    optimizer = DecoupledSGDW(model.parameters(), lr=0.1)
     dataset = RandomClassificationDataset(size=8)
     dataloader = DataLoader(dataset, batch_size=2, sampler=dist.get_sampler(dataset))
 
@@ -31,18 +36,26 @@ def test_tp_train(world_size: int):
         'fc2': RowwiseParallel(),
     }
 
-    trainer = Trainer(
-        model=model,
-        train_dataloader=dataloader,
-        parallelism_config={
-            'tp': {
-                'layer_plan': layer_plan,
-                'tensor_parallel_degree': 2,
+    if tensor_parallel_degree == 1:
+        expected_warning = 'Received tensor_parallel_degree of 1, which is a no-op. Tensor parallelism will not be used.'
+        ctx = pytest.warns(UserWarning, match=expected_warning)
+    else:
+        ctx = contextlib.nullcontext()
+
+    with ctx:
+        trainer = Trainer(
+            model=model,
+            optimizers=optimizer,
+            train_dataloader=dataloader,
+            parallelism_config={
+                'tp': {
+                    'layer_plan': layer_plan,
+                    'tensor_parallel_degree': tensor_parallel_degree,
+                },
+                'fsdp': {},
             },
-            'fsdp': {},
-        },
-        max_duration='3ba',
-    )
+            max_duration='3ba',
+        )
 
     trainer.fit()