Remove ParallelismConfig from PartialState (#3720)

* remove

* style

* fix

* valueerror instead

* add device_mesh

Author: SunMarc

src/accelerate/accelerator.py

@@ -451,23 +451,28 @@ def __init__(
                 if "recipe_handler" in handler_attr and not self.has_fp8_handler:
                     self.has_fp8_handler = True
 
-        parallelism_config = self._setup_parallelism_config(parallelism_config, torch_tp_plugin)
+        if parallelism_config is None:
+            # TODO: Remove after deprecating tp_plugin
+            if torch_tp_plugin is not None:
+                parallelism_config = ParallelismConfig(tp_size=torch_tp_plugin.tp_size)
+            elif os.environ.get("ACCELERATE_USE_PARALLELISM_CONFIG", "false").lower() == "true":
+                parallelism_config = ParallelismConfig()
 
         kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {}
-        kwargs["parallelism_config"] = parallelism_config
         self.state = AcceleratorState(
             mixed_precision=mixed_precision,
             cpu=cpu,
             dynamo_plugin=dynamo_plugin,
             deepspeed_plugin=deepspeed_plugins,
             fsdp_plugin=fsdp_plugin,
             megatron_lm_plugin=megatron_lm_plugin,
+            parallelism_config=parallelism_config,
             _from_accelerator=True,
             **kwargs,
         )
 
         if self.parallelism_config:
-            self._build_torch_device_mesh(self.parallelism_config)
+            self.state.device_mesh = parallelism_config.get_device_mesh(self.device.type)
             self.parallelism_config._validate_accelerator(self)
 
         self.fp8_enabled = self.state.mixed_precision == "fp8" or mixed_precision == "fp8"
@@ -776,23 +781,6 @@ def should_save_model(self):
         # TODO: S1ro - this is a temporary solution until we figure out why `save_safe_file` is slow when not all processes
         return True
 
-    def _setup_parallelism_config(
-        self, parallelism_config: ParallelismConfig | None, torch_tp_plugin: TorchTensorParallelPlugin | None
-    ):
-        if parallelism_config is None:
-            if PartialState._shared_state != {} and PartialState().parallelism_config is not None:
-                if os.environ.get("ACCELERATE_USE_PARALLELISM_CONFIG", "false") == "true":
-                    raise ValueError(
-                        "Partial state contains a `parallelism_config` which is not None, but you configured `parallelism_config` from the `accelerate launch` CLI. We don't know which to use, please remove one of those configuration methods."
-                    )
-                parallelism_config = PartialState().parallelism_config
-            else:
-                # TODO: Remove after deprecating tp_plugin
-                tp_size = None if torch_tp_plugin is None else torch_tp_plugin.tp_size
-                parallelism_config = ParallelismConfig(tp_size=tp_size)
-
-        return parallelism_config
-
     @property
     def tensor_parallel_rank(self) -> int:
         """
@@ -843,14 +831,6 @@ def data_parallel_shard_rank(self) -> int:
             return 0
         raise RuntimeError("Shard-based data parallelism is not configured. Set `parallelism_config` first.")
 
-    def _build_torch_device_mesh(self, parallelism_config):
-        if PartialState._shared_state != {} and getattr(PartialState(), "device_mesh", None) is not None:
-            device_mesh = PartialState().device_mesh
-        else:
-            device_mesh = parallelism_config.build_device_mesh(self.device.type)
-        self.state.device_mesh = device_mesh
-        PartialState().device_mesh = device_mesh
-
     @contextmanager
     def split_between_processes(self, inputs: list | tuple | dict | torch.Tensor, apply_padding: bool = False):
         """

src/accelerate/parallelism_config.py

@@ -15,7 +15,7 @@
 import os
 import warnings
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from torch.distributed.device_mesh import init_device_mesh
 
@@ -66,6 +66,8 @@ class ParallelismConfig:
     tp_handler: Union[None, TorchTensorParallelConfig] = None
     cp_handler: Union[None, TorchContextParallelConfig] = None
 
+    device_mesh = None
+
     def __repr__(self):
         return (
             "ParallelismConfig(\n "
@@ -178,7 +180,7 @@ def build_device_mesh(self, device_type: str):
         """
         mesh = self._get_mesh()
         if len(mesh) == 0:
-            return
+            return None
         mesh_dim_names, mesh_shape = mesh
         device_mesh = init_device_mesh(
             device_type,
@@ -194,6 +196,20 @@ def build_device_mesh(self, device_type: str):
 
         return device_mesh
 
+    def get_device_mesh(self, device_type: Optional[str] = None):
+        if self.device_mesh is None:
+            if device_type is not None:
+                self.device_mesh = self.build_device_mesh(device_type)
+            else:
+                raise ("You need to pass a device_type e.g cuda to build the device mesh")
+        else:
+            if device_type is not None:
+                if self.device_mesh.device_type != device_type:
+                    raise ValueError(
+                        f"The device_mesh is already created with device type {self.device_mesh.device_type}. However, you are trying to get a device mesh with device_type {device_type}. Please check if you correctly initialized your device_mesh"
+                    )
+        return self.device_mesh
+
     def _get_mesh(self) -> tuple[tuple[int, ...], tuple[str, ...]]:
         """Generate mesh shape and dimension names for torch.distributed.init_device_mesh()."""
 

src/accelerate/state.py

@@ -180,8 +180,6 @@ def __init__(self, cpu: bool = False, **kwargs):
         if not self.initialized:
             self._cpu = cpu
             self.backend = None
-            self.parallelism_config = kwargs.pop("parallelism_config", None)
-            self.device_mesh = kwargs.pop("device_mesh", None)
             env_device = os.environ.get("ACCELERATE_TORCH_DEVICE", None)
             self.device = torch.device(env_device) if env_device is not None else None
             self.debug = parse_flag_from_env("ACCELERATE_DEBUG_MODE")
@@ -919,6 +917,7 @@ def __init__(
             self.use_ipex = None
             self.torch_tp_plugin = torch_tp_plugin
             self.parallelism_config = parallelism_config
+            self.device_mesh = None
             mixed_precision = (
                 parse_choice_from_env("ACCELERATE_MIXED_PRECISION", "no")
                 if mixed_precision is None