[tensor]add 1D device mesh (#1492)

colossalai/device/device_mesh.py

@@ -25,7 +25,13 @@ class DeviceMesh:
             (default: False)
     """
 
-    def __init__(self, physical_mesh_id, mesh_shape, mesh_alpha=None, mesh_beta=None, init_process_group=False):
+    def __init__(self,
+                 physical_mesh_id,
+                 mesh_shape,
+                 mesh_alpha=None,
+                 mesh_beta=None,
+                 init_process_group=False,
+                 need_flatten=True):
         self.physical_mesh_id = physical_mesh_id
         self.mesh_shape = mesh_shape
         self._logical_mesh_id = self.physical_mesh_id.reshape(self.mesh_shape)
@@ -39,8 +45,12 @@ class DeviceMesh:
             mesh_beta = [1] * len(self.mesh_shape)
         self.mesh_alpha = tuple(mesh_alpha)
         self.mesh_beta = tuple(mesh_beta)
-        if init_process_group:
+        self.init_process_group = init_process_group
+        self.need_flatten = need_flatten
+        if self.init_process_group:
             self.process_groups_dict = self.create_process_groups_for_logical_mesh()
+        if self.need_flatten:
+            self.flatten_device_mesh = self.flatten()
 
     @property
     def shape(self):
@@ -54,6 +64,19 @@ class DeviceMesh:
     def logical_mesh_id(self):
         return self._logical_mesh_id
 
+    def flatten(self):
+        """
+        Flatten the logical mesh into an effective 1d logical mesh,
+        """
+        flatten_mesh_shape_size = len(self.mesh_shape)
+        flatten_mesh_shape = [self.num_devices]
+        return DeviceMesh(self.physical_mesh_id,
+                          tuple(flatten_mesh_shape),
+                          mesh_alpha=[max(self.mesh_alpha)] * (flatten_mesh_shape_size - 1),
+                          mesh_beta=[min(self.mesh_beta)] * (flatten_mesh_shape_size - 1),
+                          init_process_group=self.init_process_group,
+                          need_flatten=False)
+
     def _global_rank_to_logical_rank_map(self, tensor, index_list):
         '''
         This method is a helper function to build convert_map recursively.

colossalai/tensor/shape_consistency.py

@@ -3,6 +3,7 @@ from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
 from colossalai.tensor.utils import all_gather_simulator, all_to_all_simulator, shard_simulator
 from enum import Enum
 from copy import deepcopy
+from typing import Dict, List, Optional, Tuple, Union
 import torch.distributed as dist
 import math
 from functools import reduce
@@ -29,9 +30,9 @@ class CommSpec:
     Argument:
         comm_pattern(CollectiveCommPattern): decribe the communication method used in this spec.
         sharding_spec(ShardingSpec): This is sharding spec of the tensor which will join the communication action.
-        gather_dim(int, optional): The gather_dim of the tensor will be gathered.
-        shard_dim(int, optional): The shard_dim of the tensor will be sharded.
-        logical_process_axis(int, optional): The mesh_dim to implement the communication action.
+        gather_dim(int, Optional): The gather_dim of the tensor will be gathered.
+        shard_dim(int, Optional): The shard_dim of the tensor will be sharded.
+        logical_process_axis(Union(int, List[int]), Optional): The mesh_dim to implement the communication action.
     '''
 
     def __init__(self, comm_pattern, sharding_spec, gather_dim=None, shard_dim=None, logical_process_axis=None):
@@ -40,6 +41,11 @@ class CommSpec:
         self.gather_dim = gather_dim
         self.shard_dim = shard_dim
         self.logical_process_axis = logical_process_axis
+        if isinstance(self.logical_process_axis, list):
+            self.device_mesh = self.sharding_spec.device_mesh.flatten_device_mesh
+            self.logical_process_axis = 0
+        else:
+            self.device_mesh = self.sharding_spec.device_mesh
 
     def __repr__(self):
         res_list = ["CommSpec:("]
@@ -70,11 +76,11 @@ class CommSpec:
         '''
         comm_size = reduce(operator.mul, self.sharding_spec.get_sharded_shape_per_device(), 1)
         if self.comm_pattern == CollectiveCommPattern.ALLGATHER:
-            return self.sharding_spec.device_mesh.all_gather_cost(comm_size, self.logical_process_axis)
+            return self.device_mesh.all_gather_cost(comm_size, self.logical_process_axis)
         if self.comm_pattern == CollectiveCommPattern.ALLTOALL:
-            return self.sharding_spec.device_mesh.all_to_all_cost(comm_size, self.logical_process_axis)
+            return self.device_mesh.all_to_all_cost(comm_size, self.logical_process_axis)
         if self.comm_pattern == CollectiveCommPattern.ALLREDUCE:
-            return self.sharding_spec.device_mesh.all_reduce_cost(comm_size, self.logical_process_axis)
+            return self.device_mesh.all_reduce_cost(comm_size, self.logical_process_axis)
         if self.comm_pattern == CollectiveCommPattern.SHARD:
             return 0
         raise RuntimeError(f"Could not find a matching CollectiveCommPattern for {self.comm_pattern}.")
@@ -87,15 +93,14 @@ class CommSpec:
         Argument:
             tensor(torch.Tensor): Tensor stored in each device, which could be different in different ranks.
         '''
-        device_mesh = self.sharding_spec.device_mesh
-        process_groups_list = device_mesh.process_groups_dict[self.logical_process_axis]
+        process_groups_list = self.device_mesh.process_groups_dict[self.logical_process_axis]
 
         if self.comm_pattern == CollectiveCommPattern.ALLGATHER:
             for rank_list, process_group in process_groups_list:
                 if dist.get_rank() in rank_list:
                     tensor_list = [
                         torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device)
-                        for _ in range(self.sharding_spec.device_mesh.mesh_shape[self.logical_process_axis])
+                        for _ in range(self.device_mesh.mesh_shape[self.logical_process_axis])
                     ]
                     tensor = tensor
                     group = process_group