Protect import for device_mesh (#3742)

Author: SunMarc

src/accelerate/parallelism_config.py

@@ -17,9 +17,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional, Union
 
-from torch.distributed.device_mesh import init_device_mesh
-
 from accelerate.utils.dataclasses import TorchContextParallelConfig, TorchTensorParallelConfig
+from accelerate.utils.versions import is_torch_version
 
 
 if TYPE_CHECKING:
@@ -191,6 +190,11 @@ def build_device_mesh(self, device_type: str):
         Args:
             device_type (`str`): The type of device for which to build the mesh, e
         """
+        if is_torch_version(">=", "2.2.0"):
+            from torch.distributed.device_mesh import init_device_mesh
+        else:
+            raise RuntimeError("Building a device_mesh requires to have torch>=2.2.0")
+
         mesh = self._get_mesh()
         if len(mesh) == 0:
             return None

tests/test_dataclasses.py

@@ -76,7 +76,7 @@ def mock_flatten(name):
 
             return mesh
 
-        with patch("accelerate.parallelism_config.init_device_mesh", side_effect=mock_init_mesh):
+        with patch("torch.distributed.device_mesh.init_device_mesh", side_effect=mock_init_mesh):
             yield mock_init_mesh
 
     @pytest.mark.parametrize(