fix: test cases

Signed-off-by: Mehant Kammakomati <[email protected]>

Author: kmehant

src/accelerate/accelerator.py

@@ -108,7 +108,12 @@
     save_fsdp_optimizer,
     wait_for_everyone,
 )
-from .utils.constants import FSDP_PYTORCH_VERSION, PROFILE_PATTERN_NAME, BETA_TP_AVAILABLE_PYTORCH_VERSION
+from .utils.constants import (
+    BETA_TP_AVAILABLE_PYTORCH_VERSION,
+    BETA_TP_AVAILABLE_TRANSFORMERS_VERSION,
+    FSDP_PYTORCH_VERSION,
+    PROFILE_PATTERN_NAME,
+)
 from .utils.modeling import get_state_dict_offloaded_model
 from .utils.other import is_compiled_module
 
@@ -359,10 +364,15 @@ def __init__(
             if not is_torch_version(">=", FSDP_PYTORCH_VERSION):
                 raise ValueError(f"FSDP requires PyTorch >= {FSDP_PYTORCH_VERSION}")
 
-        if os.environ.get("ACCELERATE_USE_TP", "false") == "true" or isinstance(torch_tp_plugin, TorchTensorParallelPlugin):
+        if os.environ.get("ACCELERATE_USE_TP", "false") == "true" or isinstance(
+            torch_tp_plugin, TorchTensorParallelPlugin
+        ):
             if not is_torch_version(">=", BETA_TP_AVAILABLE_PYTORCH_VERSION):
                 raise ValueError(f"TP requires PyTorch >= {BETA_TP_AVAILABLE_PYTORCH_VERSION}")
 
+            if not compare_versions("transformers", ">=", BETA_TP_AVAILABLE_TRANSFORMERS_VERSION):
+                raise ValueError(f"TP requires transformers >= {BETA_TP_AVAILABLE_TRANSFORMERS_VERSION}")
+
         if fsdp_plugin is None:  # init from env variables
             fsdp_plugin = (
                 FullyShardedDataParallelPlugin() if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" else None
@@ -373,12 +383,14 @@ def __init__(
             os.environ["ACCELERATE_USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
         if torch_tp_plugin is None:
-            torch_tp_plugin = (TorchTensorParallelPlugin() if os.environ.get("ACCELERATE_USE_TP", "false") == "true" else None)
+            torch_tp_plugin = (
+                TorchTensorParallelPlugin() if os.environ.get("ACCELERATE_USE_TP", "false") == "true" else None
+            )
         else:
             if not isinstance(torch_tp_plugin, TorchTensorParallelPlugin):
                 raise TypeError("`torch_tp_plugin` must be a TorchTensorParallelPlugin object.")
             os.environ["ACCELERATE_USE_TP"] = "true"
-        
+
         if megatron_lm_plugin is None:  # init from env variables
             megatron_lm_plugin = (
                 MegatronLMPlugin() if os.environ.get("ACCELERATE_USE_MEGATRON_LM", "false") == "true" else None
@@ -1489,8 +1501,14 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e
                     if self.ddp_handler is not None:
                         self.ddp_handler.register_comm_hook(model)
             elif self.distributed_type == DistributedType.TP:
-                if not model.supports_tp_plan:
-                    raise NotImplementedError("Provided model does not support tensor parallelism")
+                if hasattr(model, "supports_tp_plan") and not model.supports_tp_plan:
+                    if not compare_versions("transformers", ">=", BETA_TP_AVAILABLE_TRANSFORMERS_VERSION):
+                        raise ValueError(f"TP requires transformers >= {BETA_TP_AVAILABLE_TRANSFORMERS_VERSION}")
+                    raise NotImplementedError(
+                        "Provided model does not support tensor parallelism. \
+                        Tensor parallelism plan can be added as base_model_tp_plan to model config class \
+                        and _tp_plan attribute to model class."
+                    )
                 model.tensor_parallel(self.state.torch_tp_plugin.torch_device_mesh["tp"])
             elif self.distributed_type == DistributedType.FSDP:
                 # We need to fix the optimizer *before* sharding the model

src/accelerate/commands/config/cluster.py

@@ -487,7 +487,7 @@ def get_cluster_input():
                 distributed_type = DistributedType.TP
             if distributed_type == DistributedType.TP:
                 tp_config["tp_size"] = _ask_field(
-                    "What should be your Tensor Parallel degree? [1e8]: ",
+                    "What should be your Tensor Parallel degree? [1]: ",
                     int,
                     default=1,
                 )

src/accelerate/commands/launch.py

@@ -594,7 +594,7 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="Decides Whether (true|false) intermediate activations are freed during the forward pass, and a checkpoint is left as a placeholder. (useful only when `use_fsdp` flag is passed).",
     )
-    
+
     # tp args
     tp_args = parser.add_argument_group("TP Arguments", "Arguments related to Tensor Parallelism using PyToch.")
     tp_args.add_argument(

src/accelerate/data_loader.py

@@ -752,11 +752,11 @@ def __init__(
         self.iteration = 0
 
         # if a device mesh is provided extract each dimension (dp, fsdp, tp)
-        # device mesh may hold any number of dimensions, however, 
+        # device mesh may hold any number of dimensions, however,
         # below code is for targetted support for dp, fsdp and tp
-        
-        # device mesh will be used only if there is tp involved 
-        # or any multi-dimensional parallelism involving tp 
+
+        # device mesh will be used only if there is tp involved
+        # or any multi-dimensional parallelism involving tp
         # (dp, tp) (fsdp, tp) (dp, fsdp, tp)
         # otherwise the default behavour not using device mesh should be sufficient
         # since multi dimensional parallelism devoid of tp would anyway need
@@ -789,8 +789,10 @@ def _fetch_batches(self, iterator):
                 if self.split_batches:
                     # One batch of the main iterator is dispatched and split.
                     if self.submesh_tp:
-                        logger.warning("Use of split_batches for TP would need the dataloader to produce duplicate batches,"
-                                       "otherwise, use dispatch_batches=True instead.")
+                        logger.warning(
+                            "Use of split_batches for TP would need the dataloader to produce duplicate batches,"
+                            "otherwise, use dispatch_batches=True instead."
+                        )
                     self._update_state_dict()
                     batch = next(iterator)
                 else:
@@ -996,7 +998,7 @@ def prepare_data_loader(
     data_seed: Optional[int] = None,
     non_blocking: bool = False,
     use_stateful_dataloader: bool = False,
-    torch_device_mesh: torch.distributed.DeviceMesh = None,
+    torch_device_mesh=None,
 ) -> DataLoader:
     """
     Wraps a PyTorch `DataLoader` to generate batches for one of the processes only.
@@ -1090,7 +1092,7 @@ def prepare_data_loader(
     state = PartialState()
     if num_processes is None:
         num_processes = state.num_processes
-    
+
     # when device mesh is used, specifically with TP
     # then there is need to update process_index and num_processes
     # to bring in the effect of generating same batch across TP ranks
@@ -1110,7 +1112,7 @@ def prepare_data_loader(
             submesh_dp_size = torch_device_mesh["dp"].size()
         if "fsdp" in torch_device_mesh.mesh_dim_names:
             submesh_fsdp_size = torch_device_mesh["fsdp"].size()
-    num_processes = (submesh_fsdp_size * submesh_dp_size)
+    num_processes = submesh_fsdp_size * submesh_dp_size
     if process_index is None:
         process_index = state.process_index
     if torch_device_mesh:

src/accelerate/test_utils/scripts/external_deps/test_performance.py

@@ -43,6 +43,7 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name:
         model_name (`str`, *optional*):
     """
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+
     datasets = load_dataset("glue", "mrpc")
 
     def tokenize_function(examples):
@@ -93,6 +94,10 @@ def training_function(config, args):
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True)
 
+    if args.add_pad_token:
+        if model.config.pad_token_id is None:
+            model.config.pad_token_id = 0
+
     # Instantiate optimizer
     optimizer_cls = (
         AdamW
@@ -243,6 +248,12 @@ def main():
         default=3,
         help="Number of train epochs.",
     )
+    parser.add_argument(
+        "--add_pad_token",
+        type=bool,
+        default=False,
+        help="To add pad token if not exists.",
+    )
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
     training_function(config, args)

src/accelerate/test_utils/testing.py

@@ -342,6 +342,13 @@ def require_deepspeed(test_case):
     return unittest.skipUnless(is_deepspeed_available(), "test requires DeepSpeed")(test_case)
 
 
+def require_tp(test_case):
+    """
+    Decorator marking a test that requires TP installed. These tests are skipped when TP isn't installed
+    """
+    return unittest.skipUnless(is_torch_version(">=", "2.3.0"), "test requires torch version >= 2.3.0")(test_case)
+
+
 def require_torch_min_version(test_case=None, version=None):
     """
     Decorator marking that a test requires a particular torch version to be tested. These tests are skipped when an

src/accelerate/utils/constants.py

@@ -47,6 +47,7 @@
 XPU_PROFILING_AVAILABLE_PYTORCH_VERSION = "2.4.0"
 MITA_PROFILING_AVAILABLE_PYTORCH_VERSION = "2.1.0"
 BETA_TP_AVAILABLE_PYTORCH_VERSION = "2.3.0"
+BETA_TP_AVAILABLE_TRANSFORMERS_VERSION = "4.47.0"
 
 STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
 

src/accelerate/utils/dataclasses.py

@@ -1826,10 +1826,14 @@ class TorchTensorParallelPlugin:
         metadata={"help": "tensor parallel size will be used in the device mesh preparation"},
     )
 
-    # type has to be "torch.distributed.DeviceMesh"
-    torch_device_mesh: torch.distributed.DeviceMesh = field(default=None)
+    # torch_device_mesh is fo type "torch.distributed.DeviceMesh"
+    torch_device_mesh: Optional["torch.distributed.DeviceMesh"] = field(default=None)
 
     def __post_init__(self):
+        self.tp_size = self.tp_size if os.environ.get("TP_SIZE", "1") == "1" else int(os.environ.get("TP_SIZE", "1"))
+        if self.tp_size == 1:
+            raise ValueError("Provide TP degree > 1.")
+
         if is_torch_version("<", BETA_TP_AVAILABLE_PYTORCH_VERSION):
             raise ValueError(
                 f"Minimum PyTorch version {BETA_TP_AVAILABLE_PYTORCH_VERSION} needed to use tensor parallel."

src/accelerate/utils/launch.py

@@ -286,6 +286,7 @@ def prepare_multi_gpu_env(args: argparse.Namespace) -> Dict[str, str]:
 
     if args.use_tp:
         current_env["ACCELERATE_USE_TP"] = "true"
+        current_env["TP_SIZE"] = str(args.tp_size)
 
     if args.use_megatron_lm:
         prefix = "MEGATRON_LM_"

tests/tp/test_tp.py

@@ -0,0 +1,61 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from accelerate.test_utils.testing import (
+    TempDirTestCase,
+    execute_subprocess_async,
+    get_launch_command,
+    path_in_accelerate_package,
+    require_multi_device,
+    require_non_torch_xla,
+    require_tp,
+    require_transformers,
+    slow,
+)
+from accelerate.utils import patch_environment
+
+
+@require_non_torch_xla
+@require_tp
+@require_multi_device
+@require_transformers
+@slow
+class TPIntegrationTest(TempDirTestCase):
+    test_scripts_folder = path_in_accelerate_package("test_utils", "scripts", "external_deps")
+
+    def setUp(self):
+        super().setUp()
+        self.test_tp_size = 2
+        self.model_name_or_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+        self.batch_size = 1
+        from transformers.trainer_utils import set_seed
+
+        set_seed(42)
+
+    def test_working_of_tp(self):
+        self.test_file_path = self.test_scripts_folder / "test_performance.py"
+        cmd = get_launch_command(
+            num_processes=self.test_tp_size, num_machines=1, machine_rank=0, use_tp=True, tp_size=self.test_tp_size
+        )
+        cmd.extend(
+            [
+                self.test_file_path,
+                f"--output_dir={self.tmpdir}",
+                f"--model_name_or_path={self.model_name_or_path}",
+                "--add_pad_token=true",
+            ]
+        )
+        with patch_environment(omp_num_threads=1):
+            execute_subprocess_async(cmd)