From ac76dbf6903677599f14a9bcf5300bbc6583c385 Mon Sep 17 00:00:00 2001 From: lzy0920232 Date: Mon, 19 Jan 2026 22:49:26 +0800 Subject: [PATCH] code_hyper_parallel_dtensor --- hyper_parallel/core/device_mesh.py | 121 +++++-- hyper_parallel/core/dtensor.py | 244 ++++++++++++-- hyper_parallel/core/layout.py | 314 +++--------------- hyper_parallel/core/shard/_op_dispatch.py | 16 +- hyper_parallel/core/shard/api.py | 131 +++++++- hyper_parallel/core/shard/local_func.py | 77 +++-- hyper_parallel/core/tensor_redistribution.py | 8 +- hyper_parallel/platform/mindspore/dtensor.py | 66 +++- hyper_parallel/platform/mindspore/platform.py | 4 +- tests/mindspore/st/shard/base_shard.py | 47 ++- tests/mindspore/st/shard/utils.py | 12 +- .../mindspore/st/test_placement_conversion.py | 5 +- 12 files changed, 621 insertions(+), 424 deletions(-) diff --git a/hyper_parallel/core/device_mesh.py b/hyper_parallel/core/device_mesh.py index ec132af..41a287f 100644 --- a/hyper_parallel/core/device_mesh.py +++ b/hyper_parallel/core/device_mesh.py @@ -70,13 +70,13 @@ class DeviceMesh: Topological abstraction describing cluster devices. Args: - mesh (Union[Tensor, list, np.ndarray]): A multi-dimensional array, list, or integer + mesh (Union[Tensor, list, tuple, np.ndarray]): A multi-dimensional array, list, or integer tensor describing the device layout. The IDs in the mesh are global IDs of the default process group, representing the multi-dimensional networking structure of devices in distributed training (e.g., [[0,1],[2,3]] represents a 2x2 device mesh). If a list or non-int32 tensor is provided, it will be automatically converted to an int32 tensor. - alias_name (Tuple[str]): A tuple of alias names for each dimension of mesh. + alias_name (tuple[str]): A tuple[str] of alias names for each dimension of mesh. Attributes: ndim (int): Number of dimensions in the mesh. @@ -94,12 +94,12 @@ class DeviceMesh: >>> # Get sub mesh >>> dp_mesh = device_mesh["dp"] >>> # Access ndim - >>> >>> print(device_mesh.ndim) # Output: 2. + >>> print(device_mesh.ndim) # Output: 2 >>> print(device_mesh.mesh_shape) # Output: (2, 2) >>> print(device_mesh.rank_list) # Output: (0, 1, 2, 3) """ def __init__(self, - mesh: Union[Tensor, list, np.ndarray], + mesh: Union[Tensor, list, tuple, np.ndarray], alias_name: Tuple[str]): # Convert mesh to Tensor with int32 dtype mesh = self._convert_mesh_to_tensor(mesh) @@ -159,9 +159,9 @@ class DeviceMesh: self._sub_mesh: List['DeviceMesh'] = [] @staticmethod - def _convert_mesh_to_tensor(mesh: Union[Tensor, list, np.ndarray]) -> Tensor: + def _convert_mesh_to_tensor(mesh: Union[Tensor, list, tuple, np.ndarray]) -> Tensor: """Convert mesh to Tensor with int32 dtype.""" - if isinstance(mesh, (list, np.ndarray)): + if isinstance(mesh, (list, tuple, np.ndarray)): mesh = Tensor(mesh) elif not isinstance(mesh, Tensor): raise TypeError( @@ -222,16 +222,21 @@ class DeviceMesh: """ Get a sub DeviceMesh based on the specified dimension names. + This method supports both original dimension names and flattened dimension names. + For example, if a mesh has dimensions ("dp", "cp", "tp") and a flattened mesh + "dp_cp" was created via flatten(), both mesh["dp"] and mesh["dp_cp"] are valid. + Args: sub_alias_name: A string or tuple of strings specifying the dimension names - for the sub mesh. + for the sub mesh. Can be original dimension names or flattened + dimension names registered in the root mesh's flatten_mapping. Returns: - DeviceMesh A new DeviceMesh representing the sub mesh. + DeviceMesh: A new DeviceMesh representing the sub mesh. Raises: ValueError: If sub_alias_name is invalid or not a contiguous prefix. - KeyError: If sub_alias_name contains names not in alias_name. + KeyError: If sub_alias_name contains names not in alias_name or flatten_mapping. Examples: >>> mesh = platform.tensor([[0, 1], [2, 3]]) @@ -239,6 +244,10 @@ class DeviceMesh: >>> dp_mesh = device_mesh["dp"] >>> print(dp_mesh.mesh_shape) # Output: (2,) >>> print(dp_mesh.alias_name) # Output: ("dp",) + >>> # After creating a flattened mesh: + >>> flat_mesh = device_mesh.flatten() + >>> # Can also access via flattened name: + >>> same_flat_mesh = device_mesh["dp_tp"] """ # Convert string to tuple for uniform handling if isinstance(sub_alias_name, str): @@ -252,13 +261,48 @@ class DeviceMesh: if len(sub_alias_name) == 0: raise ValueError("sub_alias_name cannot be empty") - # Validate all names exist in alias_name + # Get root mesh and its flatten_mapping + root_mesh = self._get_root_mesh() + flatten_mapping = root_mesh.get_flatten_mapping() + + # Check if requesting a single flattened dimension name + # If so, return the flattened mesh directly from flatten_mapping + if len(sub_alias_name) == 1 and sub_alias_name[0] in flatten_mapping: + return flatten_mapping[sub_alias_name[0]] + + # Build valid dimension names: original alias_name + flattened dimension names + valid_dim_names = list(self._alias_name) + list(flatten_mapping.keys()) + + # Validate all names exist in valid_dim_names for name in sub_alias_name: - if name not in self._alias_name: + if name not in valid_dim_names: raise KeyError( - f"Dimension name '{name}' not found in alias_name {self._alias_name}" + f"Dimension name '{name}' not found in alias_name {self._alias_name} " + f"or flatten_mapping keys {list(flatten_mapping.keys())}" ) + # Separate original dimensions and flattened dimensions + original_dims = [name for name in sub_alias_name if name in self._alias_name] + flattened_dims = [name for name in sub_alias_name if name in flatten_mapping] + + # If all requested dimensions are flattened dimensions + if len(flattened_dims) == len(sub_alias_name): + if len(flattened_dims) == 1: + return flatten_mapping[flattened_dims[0]] + # Multiple flattened dimensions case - not commonly supported + raise ValueError( + f"Slicing multiple flattened dimensions {flattened_dims} simultaneously " + f"is not supported. Please slice them separately." + ) + + # If mixed original and flattened dimensions + if flattened_dims: + raise ValueError( + f"Cannot mix original dimensions {original_dims} with flattened dimensions " + f"{flattened_dims} in a single slice operation." + ) + + # From here, all dimensions are original dimensions # Check if sub_alias_name is a contiguous prefix of alias_name # Get the indices of sub_alias_name in alias_name indices = [self._alias_name.index(name) for name in sub_alias_name] @@ -300,7 +344,7 @@ class DeviceMesh: alias_name=sub_alias_name ) # Set root mesh reference - sub_mesh.root_mesh = self + sub_mesh.root_mesh = self._get_root_mesh() # Cache the sub mesh self._sub_mesh_cache[cache_key] = sub_mesh @@ -314,9 +358,9 @@ class DeviceMesh: Get the communication group for a specific mesh dimension. Args: - mesh_dim (Optional[Union[int, str]], optional): The dimension index or name. If None and mesh is 1D, + mesh_dim: The dimension index or name. If None and mesh is 1D, returns the only group. If None and mesh is multi-dimensional, - raises an error. Default: None. + raises an error. Returns: The process group for the specified dimension. @@ -376,7 +420,7 @@ class DeviceMesh: raises an error. Returns: - int The local rank within the specified dimension. + int: The local rank within the specified dimension. Raises: RuntimeError: If mesh_dim is None and mesh has more than 1 dimension. @@ -428,24 +472,33 @@ class DeviceMesh: return coord[dim_index] - def flatten(self) -> 'DeviceMesh': + def flatten(self, mesh_dim_name: Optional[str] = None) -> 'DeviceMesh': """ Returns a 1D DeviceMesh by flattening the current DeviceMesh. + Args: + mesh_dim_name (str, optional): The name for the flattened dimension. + If not provided, the name will be generated by joining the original + alias names with underscore (e.g., "dp_tp" for ("dp", "tp")). + This name will be used as the key in the root mesh's flatten_mapping. + Returns: DeviceMesh: A 1D DeviceMesh with flattened dimensions. Raises: - ValueError: If generated mesh_dim_name conflicts with existing alias names. + ValueError: If mesh_dim_name conflicts with existing alias names. Examples: >>> mesh = Tensor([[0, 1], [2, 3]]) >>> device_mesh = DeviceMesh(mesh, ("dp", "tp")) + >>> # Using default name >>> flat_mesh = device_mesh.flatten() - >>> print(flat_mesh.mesh_shape) # Output: (4,) >>> print(flat_mesh.alias_name) # Output: ("dp_tp",) + >>> # Using custom name + >>> flat_mesh = device_mesh.flatten(mesh_dim_name="custom_name") + >>> print(flat_mesh.alias_name) # Output: ("custom_name",) """ - return self._create_flatten_mesh() + return self._create_flatten_mesh(mesh_dim_name) def _get_root_mesh(self) -> 'DeviceMesh': """Get the root mesh of this DeviceMesh.""" @@ -454,12 +507,18 @@ class DeviceMesh: # pylint: disable=protected-access return self._root_mesh._get_root_mesh() - def _create_flatten_mesh(self) -> 'DeviceMesh': - """Create a flattened 1D mesh from the current mesh.""" + def _create_flatten_mesh(self, mesh_dim_name: Optional[str] = None) -> 'DeviceMesh': + """Create a flattened 1D mesh from the current mesh. + + Args: + mesh_dim_name (str, optional): The name for the flattened dimension. + If not provided, defaults to joining alias names with underscore. + """ root_mesh = self._get_root_mesh() - # Generate mesh_dim_name by joining alias names - mesh_dim_name = "_".join(self._alias_name) + # Generate mesh_dim_name by joining alias names if not provided + if mesh_dim_name is None: + mesh_dim_name = "_".join(self._alias_name) # Flatten a 1D device mesh into its original alias_name will return itself if self.ndim == 1 and mesh_dim_name in self._alias_name: @@ -694,7 +753,7 @@ class DeviceMesh: _DEVICE_MESH_MAP = {} -def _create_device_mesh(mesh: Tensor, alias_name: Tuple[str]): +def _create_device_mesh(mesh_shape: Tuple[int], alias_name: Tuple[str], rank_list: Tuple[int]): """ Create or retrieve a cached DeviceMesh. @@ -705,8 +764,7 @@ def _create_device_mesh(mesh: Tensor, alias_name: Tuple[str]): Returns: DeviceMesh: A DeviceMesh object. """ - mesh_shape = tuple(mesh.shape) - rank_list = tuple(mesh.flatten().tolist()) + mesh = Tensor(rank_list).reshape(mesh_shape) rank_ids = (rank_list[0], rank_list[-1]) map_key = hash((mesh_shape, alias_name, rank_ids)) if map_key not in _DEVICE_MESH_MAP: @@ -731,10 +789,10 @@ def init_device_mesh( - Validation of parameters Args: - mesh_shape (Tuple[int, ...]): A tuple describing the dimensions of the multi-dimensional + mesh_shape (tuple[int]): A tuple describing the dimensions of the multi-dimensional array that describes the layout of devices. For example, (2, 4) creates a 2D mesh with 2 rows and 4 columns. - alias_name (Tuple[str, ...]): A tuple of names to assign to each dimension of the mesh. + alias_name (tuple[str]): A tuple of names to assign to each dimension of the mesh. Its length must match the length of mesh_shape. Each string must be unique. Returns: @@ -779,8 +837,7 @@ def init_device_mesh( # Generate mesh tensor from mesh_shape with sequential ranks using Tensor() total_devices = int(np.prod(np.array(mesh_shape))) - rank_list = list(range(total_devices)) - mesh = Tensor(rank_list).reshape(mesh_shape) + rank_list = tuple(range(total_devices)) # Use the caching mechanism - return _create_device_mesh(mesh, alias_name) + return _create_device_mesh(mesh_shape, alias_name, rank_list) diff --git a/hyper_parallel/core/dtensor.py b/hyper_parallel/core/dtensor.py index e63f686..ef81a26 100644 --- a/hyper_parallel/core/dtensor.py +++ b/hyper_parallel/core/dtensor.py @@ -14,8 +14,11 @@ # ============================================================================ """dtensor""" import copy as cp -from hyper_parallel.core.layout import Layout +from typing import Sequence, Tuple +from hyper_parallel.core.layout import Layout, DeviceMesh, _get_slice_tensor_by_layout +from hyper_parallel.core.placement_types import Placement, Replicate from hyper_parallel.platform import get_platform + platform = get_platform() DTensorBase = platform.DTensorBase Tensor = platform.Tensor @@ -33,60 +36,261 @@ class SkipDTensorDispatch(): enable_dtensor_dispatch() +def _build_layout( + device_mesh: DeviceMesh, + placements: Sequence[Placement], + tensor_dim: int +) -> Layout: + """ + Build Layout from device_mesh and placements. + + Args: + device_mesh: The device mesh describing the device topology. + placements: Sequence of Placement objects (Shard, Replicate, etc.). + tensor_dim: Number of dimensions in the tensor. + + Returns: + Layout: The built layout object. + """ + layout = Layout.from_device_mesh(device_mesh) + result = layout(placements) + result.placement_to_tensor_map(tensor_dim) + return result + + class DTensor(DTensorBase): """ - DTensor + DTensor - Distributed Tensor + + A DTensor represents a tensor that is distributed across multiple devices + according to a DeviceMesh and placement specifications. + + Args: + local_tensor (Tensor): The local tensor shard on this device. + device_mesh (DeviceMesh): The device mesh describing the device topology. + placements (Sequence[Placement]): The placement strategy for each mesh dimension. + Each element should be a Placement object (Shard, Replicate, Partial, etc.). + + Example: + >>> from hyper_parallel.core.placement_types import Shard, Replicate + >>> mesh = DeviceMesh((2, 2), ("dp", "tp")) + >>> local_tensor = Tensor(np.ones((4, 4))) + >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()]) """ _local_tensor: Tensor - _layout: Layout + _device_mesh: DeviceMesh + _placements: Tuple[Placement, ...] - def __init_data__(self, local_tensor, layout): + def __init_data__( + self, + local_tensor: Tensor, + device_mesh: DeviceMesh, + placements: Sequence[Placement] + ): self._local_tensor = local_tensor - self._layout = layout + self._device_mesh = device_mesh + self._placements = tuple(placements) + # Build internal layout for redistribution operations + self._layout = _build_layout( + device_mesh, placements, len(local_tensor.shape) + ) + + @property + def device_mesh(self) -> DeviceMesh: + """The device mesh of this DTensor.""" + return self._device_mesh + + @property + def placements(self) -> Tuple[Placement, ...]: + """The placements of this DTensor.""" + return self._placements @property - def layout(self): - """Sharding state for dtensor""" + def layout(self) -> Layout: + """Internal layout for redistribution (for backward compatibility).""" if not hasattr(self, '_layout'): return None return self._layout @staticmethod - def from_local(local_tensor: Tensor, layout: Layout): - d_tensor = DTensor(local_tensor, layout) - return d_tensor + def from_local( + local_tensor: Tensor, + device_mesh: DeviceMesh, + placements: Sequence[Placement] + ) -> 'DTensor': + """ + Create a DTensor from a local tensor with device mesh and placements. + + Args: + local_tensor (Tensor): The local tensor shard on this device. + device_mesh (DeviceMesh): The device mesh describing the device topology. + placements (Sequence[Placement]): The placement strategy. Each element + should be a Placement object (Shard, Replicate, Partial, etc.). - def to_local(self): - """covert global_tensor to local_tensor""" + Returns: + DTensor: A new DTensor instance. + + Example: + >>> from hyper_parallel.core.placement_types import Shard, Replicate + >>> mesh = DeviceMesh((2, 2), ("dp", "tp")) + >>> local_tensor = Tensor(np.ones((4, 4))) + >>> dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()]) + """ + return DTensor(local_tensor, device_mesh, placements) + + def to_local(self) -> Tensor: + """ + Convert DTensor to local tensor. + + Returns: + Tensor: The local tensor shard on this device. + """ return self._local_tensor @property - def shape(self): + def shape(self) -> Tuple[int, ...]: + """ + The global shape of this DTensor. + + Returns: + Tuple[int, ...]: The global tensor shape. + """ return self._layout.get_global_shape(self._local_tensor.shape) @property - def local_shape(self): + def local_shape(self) -> Tuple[int, ...]: + """ + The local shape of this DTensor on this device. + + Returns: + Tuple[int, ...]: The local tensor shape. + """ return self._local_tensor.shape - def redistribute(self, dst_layout): + def redistribute( + self, + device_mesh: DeviceMesh, + placements: Sequence[Placement] + ) -> 'DTensor': """ - Redistribute dtensor to destination layout. + Redistribute this DTensor to a new device mesh and placements. + + Args: + device_mesh (DeviceMesh): The target device mesh. + placements (Sequence[Placement]): The target placements. Each element + should be a Placement object (Shard, Replicate, Partial, etc.). + + Returns: + DTensor: A new DTensor with the specified distribution. + + Example: + >>> from hyper_parallel.core.placement_types import Shard, Replicate + >>> new_dtensor = dtensor.redistribute(mesh, [Replicate(), Shard(1)]) """ + # Build dst_layout from device_mesh and placements + dst_layout = _build_layout( + device_mesh, placements, len(self._local_tensor.shape) + ) + # pylint: disable=C0415 from hyper_parallel.core.tensor_redistribution import _tensor_redistribution out = _tensor_redistribution.redistribution(self, dst_layout) return out - def reduce_partial(self): + def reduce_partial(self) -> 'DTensor': """ - Reduce partial sharding state for dtensor. + Reduce partial sharding state for this DTensor. + Returns: + DTensor: A new DTensor with partial state reduced. """ - if not self.layout: + if not self._layout: return self - to_layout = cp.deepcopy(self.layout) + to_layout = cp.deepcopy(self._layout) to_layout.reset_partial() # pylint: disable=C0415 from hyper_parallel.core.tensor_redistribution import _tensor_redistribution out = _tensor_redistribution.reduce_partial(self, to_layout) return out + + def full_tensor(self) -> Tensor: + """ + Return the full tensor of this DTensor. + + Returns: + Tensor: A Tensor object that represents the full tensor of this DTensor. + The returned tensor contains the complete data gathered from + all ranks. + + Note: + This operation involves communication across all ranks in the DeviceMesh, + which may be expensive for large tensors. Use with caution in + performance-critical code paths. + + Example: + >>> # Assume dtensor is sharded across multiple devices + >>> local_tensor = dtensor.to_local() # Returns only the local shard + >>> full_tensor = dtensor.full_tensor() # Returns the complete tensor + """ + if not self._layout: + return self._local_tensor + + # Create a fully replicated layout + replicated_layout = cp.deepcopy(self._layout) + + # Set all placements to Replicate and convert to tensor_map + replicated_placements = [Replicate()] * len(replicated_layout.mesh_shape) + replicated_layout.set_placements(replicated_placements) + replicated_layout.placement_to_tensor_map(len(self._local_tensor.shape)) + + # Clear partial status from original layout since Replicate has no partial + replicated_layout.reset_partial() + + # Update compact_str since tensor_map changed (used as cache key in redistribution) + replicated_layout.update_compact_str() + + # Redistribute to the replicated layout and return local tensor + # pylint: disable=C0415 + from hyper_parallel.core.tensor_redistribution import _tensor_redistribution + out = _tensor_redistribution.redistribution(self, replicated_layout) + return out.to_local() + + @staticmethod + def distribute_tensor( + tensor: Tensor, + device_mesh: DeviceMesh, + placements: Sequence[Placement] + ) -> 'DTensor': + """ + Distribute a global tensor to the device mesh according to the placements. + + Args: + tensor (Tensor): The global tensor to be distributed. All ranks + should have the same tensor data. + device_mesh (DeviceMesh): The device mesh describing the device topology. + placements (Sequence[Placement]): The placement strategy for distribution. + Each element should be a Placement object (Shard, Replicate, etc.). + + Returns: + DTensor: A new DTensor with the local shard on each rank. + + Note: + This method assumes all ranks have the same global tensor. It slices + the tensor locally without communication. If ranks have different + data, use `from_local` instead. + + Example: + >>> from hyper_parallel.core.placement_types import Shard, Replicate + >>> mesh = DeviceMesh((2, 2), ("dp", "tp")) + >>> global_tensor = Tensor(np.arange(16).reshape(4, 4)) + >>> dtensor = DTensor.distribute_tensor(global_tensor, mesh, [Shard(0), Replicate()]) + >>> # rank 0 gets: [[0,1,2,3], [4,5,6,7]] + >>> # rank 1 gets: [[8,9,10,11], [12,13,14,15]] + """ + # Build layout from device_mesh and placements + layout = _build_layout(device_mesh, placements, len(tensor.shape)) + + # Slice the global tensor to get local shard based on layout + local_tensor = _get_slice_tensor_by_layout(tensor, layout) + + return DTensor(local_tensor, device_mesh, placements) diff --git a/hyper_parallel/core/layout.py b/hyper_parallel/core/layout.py index 0e5f479..4fb790d 100644 --- a/hyper_parallel/core/layout.py +++ b/hyper_parallel/core/layout.py @@ -15,16 +15,16 @@ """layout""" import copy -from typing import Tuple +from typing import Optional, Tuple, Any, List +from hyper_parallel.core.placement_types import Placement, Shard, Replicate, Partial +from hyper_parallel.core.device_mesh import DeviceMesh, _create_device_mesh + import functools import numpy as np -from hyper_parallel.core.placement_types import Placement, Shard, Replicate, Partial from hyper_parallel.platform import get_platform platform = get_platform() -_group_map = {} - def _infer_slice_area_by_rank(mesh_shape, tensor_map, rank_id: int, full_shape: tuple): # -> tuple[tuple[int]]: """Return the range of each axis from full tensor for slice in current rank.""" @@ -97,267 +97,6 @@ def _infer_slice_shape_by_layout(global_shape, layout): return slice_shape -class _DeviceMesh: - """ - Topological abstraction describing cluster devices. - Args: - mesh_shape: - alias_name: - rank_list: - """ - - def __init__(self, - mesh_shape: Tuple[int], - alias_name: Tuple[str], - rank_list: Tuple[int]): - if not isinstance(mesh_shape, tuple): - raise TypeError(f'mesh_shape must be tuple type, but got:{type(mesh_shape)}') - if not isinstance(alias_name, tuple): - raise TypeError(f'alias_name must be tuple type, but got:{type(alias_name)}') - if len(mesh_shape) != len(alias_name): - raise ValueError('mesh_shape length should be equal to alias_name length') - for in_ele in mesh_shape: - if not isinstance(in_ele, int): - raise TypeError(f'The element of mesh_shape must be int type, but got:{type(in_ele)}') - for in_ele in alias_name: - if not isinstance(in_ele, str): - raise TypeError(f'The element of alias_name must be str type, but got:{type(in_ele)}') - if not in_ele: - raise ValueError("The element of alias_name can not be empty.") - if in_ele == "None": - raise ValueError("The element of alias_name can not set 'None', because 'None' means no sharding.") - if len(set(alias_name)) != len(alias_name): - raise ValueError(f'Each element of alias_name {alias_name} should be different') - inter_key = "interleaved_parallel" - if inter_key in alias_name and alias_name.index(inter_key) != len(alias_name) - 1: - raise ValueError(f"When alias_name {alias_name} contains keyword 'interleaved_parallel'," - f" it should be at the last dim of alias_name, which means the virtual sharding.") - self._mesh_shape = mesh_shape - self._alias_name = alias_name - self._dev_num = np.prod(np.array(self._mesh_shape)) - self._global_shape_map = {} - self._rank = platform.get_rank() - self._dev_rank = len(self._mesh_shape) - self._dev_name_to_dev_id = {name: self._dev_rank - i - 1 for i, name in enumerate(self.alias_name)} - self._dev_name_to_index = {name: i for i, name in enumerate(self.alias_name)} - self._cache_rank_list_along_axis = {} - self._rank_list = rank_list - self._check_rank_list() - self._global_shape_map = {} - - def _check_rank_list(self): - """_check_rank_list""" - if not isinstance(self._rank_list, tuple): - raise TypeError(f"The rank_list should be a tuple, but got {type(self._rank_list).__name__}.") - for in_ele in self._rank_list: - if not isinstance(in_ele, int): - raise TypeError(f"The element of rank_list should be int, but got {type(in_ele).__name__}.") - if len(np.array(self._rank_list).shape) != 1: - raise ValueError( - f"The rank_list should be a 1-D list, but got {len(np.array(self._rank_list).shape)}-D list.") - if len(self._rank_list) != np.prod(np.array(self._mesh_shape)): - raise ValueError(f"The length of rank_list should be equal to the product of mesh_shape, " - f"but got {len(self._rank_list)} and {np.prod(np.array(self._mesh_shape))}.") - - @property - def rank(self): - return self._rank - - @property - def mesh_shape(self): - return self._mesh_shape - - @property - def alias_name(self): - return self._alias_name - - @property - def rank_list(self): - return self._rank_list - - def axis_id(self, axis): - if axis == "None": - return -1 - if axis not in self.alias_name: - raise ValueError(f"The axis name must be one of mesh shape alias name {self.alias_name}), " - f"but got {axis}") - return self._dev_name_to_dev_id[axis] - - def axis_index(self, axis): - if axis not in self.alias_name: - raise ValueError(f"The axis name must be one of mesh shape alias name {self.alias_name}), " - f"but got {axis}") - return self._dev_name_to_index[axis] - - def get_device_num_along_axis(self, axis): - """Return device num along specify device axis""" - if axis not in self.alias_name: - raise ValueError(f"The axis must be one of device alias name: {self.alias_name}, but got {axis}") - return self.mesh_shape[self.alias_name.index(axis)] - - def get_rank_list_along_axis(self, axis): - """ - Get the repeat rank list when the axis is not shard. - - Args: - layout (Layout): Layout - axis (str): Axis name. - rank (int): Global rank - - Returns: - list: reduce rank list - """ - if axis in self._cache_rank_list_along_axis: - # short cut, get rank list from cache - return self._cache_rank_list_along_axis[axis] - - mesh_shape = self.mesh_shape - alias_name = self.alias_name - rank_list = self.rank_list - rank = self.rank - - if axis not in alias_name: - raise ValueError(f"Axis '{axis}' not found in alias_name {alias_name}") - - if rank not in rank_list: - raise ValueError(f"Rank {rank} not found in rank_list") - - idx = rank_list.index(rank) - coord = [0] * len(mesh_shape) - temp = idx - for i in range(len(mesh_shape) - 1, -1, -1): - coord[i] = temp % mesh_shape[i] - temp //= mesh_shape[i] - - dim_index = alias_name.index(axis) - strides = [1] * len(mesh_shape) - for i in range(len(mesh_shape) - 2, -1, -1): - strides[i] = strides[i + 1] * mesh_shape[i + 1] - - result_ranks = [] - for v in range(mesh_shape[dim_index]): - new_coord = coord.copy() - new_coord[dim_index] = v - new_idx = 0 - for i in range(len(mesh_shape)): - new_idx += new_coord[i] * strides[i] - - result_ranks.append(rank_list[new_idx]) - - self._cache_rank_list_along_axis[axis] = result_ranks - return result_ranks - - def get_global_shape(self, slice_shape, tensor_map): - """get global shape""" - map_key = hash((slice_shape, tensor_map)) - if map_key in self._global_shape_map: - return self._global_shape_map[map_key] - if tensor_map is None: - raise ValueError("tensor_map is not set. Please configure the tensor map by calling the layout.") - if len(slice_shape) != len(tensor_map): - raise ValueError(f"Length of slice_shape ({len(slice_shape)}) must match " - f"the length of tensor_map ({len(tensor_map)}).") - - n_dims = len(self._mesh_shape) - factors = [1] * len(slice_shape) - - for dev_idx, size in enumerate(self._mesh_shape): - reverse_idx = n_dims - 1 - dev_idx - for axis_idx, mapping in enumerate(tensor_map): - if isinstance(mapping, int): - if mapping == -1: - continue - if mapping == reverse_idx: - factors[axis_idx] *= size - break - elif isinstance(mapping, tuple): - if reverse_idx in mapping: - factors[axis_idx] *= size - break - - global_shape = [] - for i, dim in enumerate(slice_shape): - global_shape.append(dim * factors[i]) - self._global_shape_map[map_key] = tuple(global_shape) - return tuple(global_shape) - - def get_comm_group_by_axis(self, axis, rank): - if (axis, rank) in _group_map: - return _group_map[(axis, rank)] - rank_list = self.get_rank_list_along_axis(axis=axis) - group = platform.create_group(rank_list) - _group_map[(axis, rank)] = group - return group - - def get_devices_for_axis(self, axis, rank): - """ - Get the repeat rank list when the axis is not shard. - - Args: - layout (Layout): Layout - axis (str): Axis name. - rank (int): Global rank - - Returns: - list: reduce rank list - """ - mesh_shape = self._mesh_shape - alias_name = self._alias_name - rank_list = self._rank_list - - if axis not in alias_name: - raise ValueError(f"Axis '{axis}' not found in alias_name {alias_name}") - - if rank not in rank_list: - raise ValueError(f"Rank {rank} not found in rank_list") - - if rank not in rank_list: - raise ValueError(f"Rank {rank} not found in rank_list") - - idx = rank_list.index(rank) - coord = [0] * len(mesh_shape) - temp = idx - for i in range(len(mesh_shape) - 1, -1, -1): - coord[i] = temp % mesh_shape[i] - temp //= mesh_shape[i] - - dim_index = alias_name.index(axis) - strides = [1] * len(mesh_shape) - for i in range(len(mesh_shape) - 2, -1, -1): - strides[i] = strides[i + 1] * mesh_shape[i + 1] - - result_ranks = [] - for v in range(mesh_shape[dim_index]): - new_coord = coord.copy() - new_coord[dim_index] = v - new_idx = 0 - for i in range(len(mesh_shape)): - new_idx += new_coord[i] * strides[i] - - result_ranks.append(rank_list[new_idx]) - - return result_ranks - - def to_hash(self): - rank_ids = (self.rank_list[0], self.rank_list[-1]) - map_key = (self.mesh_shape, self.alias_name, rank_ids) - return map_key - - -_DEVICE_MESH_MAP = {} - - -def _create_device_mesh(mesh_shape: Tuple[int], alias_name: Tuple[str], rank_list: Tuple[int]): - """ - create_device_mesh - """ - rank_ids = (rank_list[0], rank_list[-1]) - map_key = hash((mesh_shape, alias_name, rank_ids)) - if map_key not in _DEVICE_MESH_MAP: - _DEVICE_MESH_MAP[map_key] = _DeviceMesh(mesh_shape, alias_name, rank_list) - return _DEVICE_MESH_MAP.get(map_key, None) - - class Layout: """ Topological abstraction describing cluster devices for tensor slice placement on the cluster. @@ -416,21 +155,15 @@ class Layout: self._placements = None @classmethod - def from_device_mesh(cls, device_mesh): + def from_device_mesh(cls, device_mesh: DeviceMesh) -> 'Layout': """ - Create a Layout object from a device mesh. + Create a Layout from an existing DeviceMesh. Args: - device_mesh (_DeviceMesh): The device mesh object containing topological information about the - cluster devices. + device_mesh (DeviceMesh): The device mesh to create layout from. Returns: - Layout: A new Layout instance initialized with the properties of the provided device mesh. - - Examples: - >>> from hyper_parallel.core.layout import Layout, _DeviceMesh - >>> device_mesh = _DeviceMesh((2, 2), ("dp", "mp"), tuple(range(4))) - >>> layout = Layout.from_device_mesh(device_mesh) + Layout: A new Layout instance. """ obj = cls.__new__(cls) obj._mesh = device_mesh @@ -491,6 +224,7 @@ class Layout: tensor_map += (len(obj.alias_name) - 1 - obj.alias_name.index(ele),) writed_map += (ele,) obj.set_tensor_map(tensor_map) + obj.tensor_map_to_placement() obj.update_compact_str() return copy.deepcopy(obj) @@ -531,6 +265,15 @@ class Layout: raise ValueError(f"Tensor dimension must be positive, but got {dim}") if dim == 0: self.set_tensor_map(()) + self._alias_tensor_map =() + for mesh_idx, placement in enumerate(self.placements): + if isinstance(placement, Partial): + op_name = getattr(placement, "reduce_op", "sum") + if isinstance(op_name, str): + op_name = op_name.lower() + self._partial[mesh_idx] = op_name + return [] + dim_map = [-1] * dim self.partial_ops = {} for mesh_idx, placement in enumerate(self.placements): @@ -542,14 +285,30 @@ class Layout: raise ValueError(f"Dimension {shard_dim} has been sharded by Mesh axis {dim_map[shard_dim]}") dim_map[shard_dim] = mesh_idx elif isinstance(placement, Partial): - op_name = getattr(placement, "reduce_op", "sum").lower() - self.partial[mesh_idx] = op_name + op_name = getattr(placement, "reduce_op", "sum") + if isinstance(op_name, str): + op_name = op_name.lower() + self._partial[mesh_idx] = op_name device_dim_count = len(self.mesh_shape) tensor_map = [ device_dim_count - 1 - mesh_idx if mesh_idx != -1 else -1 for mesh_idx in dim_map ] self.set_tensor_map(tuple(tensor_map)) + readable_map = [] + for item in self._tensor_map: + if isinstance(item, tuple): + # 处理嵌套元组 + mapped_tuple = tuple( + self._mesh.alias_name[len(self._mesh.alias_name) - 1 - dim] if dim != -1 else "None" + for dim in item + ) + readable_map.append(mapped_tuple) + else: + readable_map.append( + self._mesh.alias_name[len(self._mesh.alias_name) - 1 - item] if item != -1 else "None" + ) + self._alias_tensor_map = tuple(readable_map) return tensor_map def tensor_map_to_placement(self): @@ -650,6 +409,7 @@ class Layout: if self.is_dev_axis_apply_shard(axis): raise ValueError("Partial dim must be replicate.") self._partial[self._mesh.axis_index(axis)] = op + self.tensor_map_to_placement() def get_partial_by_dev_id(self, axis): """Get the partial status for the specified dev id""" diff --git a/hyper_parallel/core/shard/_op_dispatch.py b/hyper_parallel/core/shard/_op_dispatch.py index c607198..257d581 100644 --- a/hyper_parallel/core/shard/_op_dispatch.py +++ b/hyper_parallel/core/shard/_op_dispatch.py @@ -249,7 +249,7 @@ class OpDispatcher: if isinstance(output_layout, (tuple, list)): if len(py_output) == len(output_layout): for i, output_item in enumerate(py_output): - output += (DTensor.from_local(output_item, output_layout[i]),) + output += (DTensor.from_local(output_item, output_layout[i].mesh, output_layout[i].placements),) else: raise RuntimeError(f"Output tuple size ({len(py_output)}) " f"does not match layout tuple size ({len(output_layout)})") @@ -257,7 +257,7 @@ class OpDispatcher: raise RuntimeError("Output is a tuple but layout is not") return output - return DTensor.from_local(py_output, output_layout) + return DTensor.from_local(py_output, output_layout.mesh, output_layout.placements) def _with_layout_infer_with_tuple_expand(self, func: callable, *args, **kwargs) -> Tensor: """_with_layout_infer_with_tuple_expand""" @@ -317,7 +317,7 @@ class OpDispatcher: if isinstance(output_layout, (tuple, list)): if len(py_output) == len(output_layout): for i, output_item in enumerate(py_output): - output += (DTensor.from_local(output_item, output_layout[i]),) + output += (DTensor.from_local(output_item, output_layout[i].mesh, output_layout[i].placements),) else: raise RuntimeError(f"Output tuple size ({len(py_output)}) " f"does not match layout tuple size ({len(output_layout)})") @@ -325,7 +325,7 @@ class OpDispatcher: raise RuntimeError("Output is a tuple but layout is not") return output - return DTensor.from_local(py_output, output_layout) + return DTensor.from_local(py_output, output_layout.mesh, output_layout.placements) def _with_layout_infer_reshape(self, func: callable, *args) -> Tensor: """_with_layout_infer_reshape""" @@ -367,7 +367,7 @@ class OpDispatcher: local_shape = infer_output_tuple[1] py_output = func(input_tensor.to_local(), local_shape) - return DTensor.from_local(py_output, infer_output_tuple[0]) + return DTensor.from_local(py_output, infer_output_tuple[0].mesh, infer_output_tuple[0].placements) def _with_layout_infer_with_shape(self, func: callable, *args, **kwargs) -> Tensor: """_with_layout_infer_with_shape""" @@ -433,7 +433,7 @@ class OpDispatcher: if isinstance(output_layout, (tuple, list)): if len(py_output) == len(output_layout): for i, output_item in enumerate(py_output): - output += (DTensor.from_local(output_item, output_layout[i]),) + output += (DTensor.from_local(output_item, output_layout[i].mesh, output_layout[i].placements),) else: raise RuntimeError(f"Output tuple size ({len(py_output)}) " f"does not match layout tuple size ({len(output_layout)})") @@ -441,7 +441,7 @@ class OpDispatcher: raise RuntimeError("Output is a tuple but layout is not") return output - return DTensor.from_local(py_output, output_layout) + return DTensor.from_local(py_output, output_layout.mesh, output_layout.placements) def _with_layout_infer_slice(self, func: callable, *args) -> Tensor: """_with_layout_infer_slice""" @@ -488,7 +488,7 @@ class OpDispatcher: new_end = infer_output_tuple[2] py_output = func(input_tensor.to_local(), new_begin, new_end) - return DTensor.from_local(py_output, infer_output_tuple[0]) + return DTensor.from_local(py_output, infer_output_tuple[0].mesh, infer_output_tuple[0].placements) def safe_load_yaml_from_dir(self): """ diff --git a/hyper_parallel/core/shard/api.py b/hyper_parallel/core/shard/api.py index b5e25fe..4319594 100644 --- a/hyper_parallel/core/shard/api.py +++ b/hyper_parallel/core/shard/api.py @@ -16,8 +16,9 @@ import inspect from typing import Union, Callable, Dict from functools import wraps -from hyper_parallel.core.layout import Layout +from hyper_parallel.core.layout import Layout, DeviceMesh from hyper_parallel.core.dtensor import DTensor +from hyper_parallel.core.placement_types import Placement from hyper_parallel.platform import get_platform platform = get_platform() @@ -41,6 +42,69 @@ def _get_param_name(func): return list(sig.parameters.keys()) +def _convert_sharding_plan(sharding_plan: Dict, device_mesh: DeviceMesh) -> Dict: + """ + Convert sharding_plan values to Layout objects. + + Args: + sharding_plan: The original sharding plan with tuple specifications + device_mesh: The DeviceMesh to use for conversion + + Returns: + Dict: Converted sharding plan with Layout objects + """ + + def _to_layout(value): + """Convert a single sharding specification to Layout.""" + layout = Layout.from_device_mesh(device_mesh) + result = layout(value) + return result + + converted_plan = {} + + # Convert parameter sharding plan + param_plan = sharding_plan.get("parameter") + if param_plan is not None: + converted_param_plan = {} + for param_name, value in param_plan.items(): + converted_param_plan[param_name] = _to_layout(value) + converted_plan["parameter"] = converted_param_plan + + # Convert forward sharding plan + forward_plan = sharding_plan.get("forward") + if forward_plan is not None: + converted_forward_plan = {} + for key, value in forward_plan.items(): + if value is None: + converted_forward_plan[key] = None + elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], (tuple, list)): + converted_forward_plan[key] = [_to_layout(v) for v in value] + else: + converted_forward_plan[key] = [_to_layout(value)] + converted_plan["forward"] = converted_forward_plan + + # Convert input/output for callable sharding + if "input" in sharding_plan: + value = sharding_plan["input"] + if value is None: + converted_plan["input"] = None + elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], (tuple, list)): + converted_plan["input"] = [_to_layout(v) for v in value] + else: + converted_plan["input"] = [_to_layout(value)] + + if "output" in sharding_plan: + value = sharding_plan["output"] + if value is None: + converted_plan["output"] = None + elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], (tuple, list)): + converted_plan["output"] = [_to_layout(v) for v in value] + else: + converted_plan["output"] = [_to_layout(value)] + + return converted_plan + + def _parallel_in(func, args, kwargs, layouts): """_parallel_in""" if not isinstance(layouts, (list, dict, tuple)): @@ -63,13 +127,13 @@ def _parallel_in(func, args, kwargs, layouts): continue to_layout = _get_layout(i, is_list) - processed_args[i] = arg.redistribute(to_layout) + processed_args[i] = arg.redistribute(to_layout.mesh, to_layout.placements) for k, v in kwargs.items(): if not isinstance(v, DTensor) or layouts.get(k) is None: processed_kwargs[k] = v continue to_layout = layouts[k] - processed_kwargs[k] = v.redistribute(to_layout) + processed_kwargs[k] = v.redistribute(to_layout.mesh, to_layout.placements) return tuple(processed_args), processed_kwargs @@ -88,12 +152,13 @@ def _parallel_out(outputs, layouts): new_outputs.append(arg) continue to_layout = layouts[i] - new_outputs.append(arg.redistribute(to_layout)) + new_outputs.append(arg.redistribute(to_layout.mesh, to_layout.placements)) return tuple(new_outputs) if len(layouts) != 1: raise ValueError(f"The size of outputs and out_layout must be equal, but got 1 and " f"{len(layouts)}") - return outputs.redistribute(layouts[0]) if isinstance(outputs, DTensor) else outputs + + return outputs.redistribute(layouts[0].mesh, layouts[0].placements) if isinstance(outputs, DTensor) else outputs def _forward_pre_hook(cell, args): @@ -190,28 +255,58 @@ def _shard_callable(func: Callable, sharding_plan: Dict): return _shard_wrapper -def shard(model: Union[Module, Callable], sharding_plan: Dict): +def shard(model: Union[Module, Callable], device_mesh: DeviceMesh, sharding_plan: Dict): """ - Defining the input, output and parameters layouts of this cell or Callable. + Defining the input, output and parameters layouts of this cell or Callable. + + Note: + - It is valid only in pynative mode. - Note: - - It is valid only in pynative mode. + .. warning:: + The method is currently not supported in Graph mode. - .. warning:: - The method is currently not supported in Graph mode. + Args: + model (Module or Callable): The model to be sharded. + device_mesh (DeviceMesh): The device mesh for sharding. + sharding_plan (Dict): Define the layout for the specified parameters, inputs or outputs. + The sharding specification can be: + - tuple of strings for alias format, e.g., ("dp", "None") + - tuple of Placements, e.g., (Shard(0), Replicate()) - Args: - model (Module or Callable): The model to be sharded. - sharding_plan (Dict): Define the layout for the specified parameters, inputs or outputs. + Returns: + Module or Callable: The sharded model. + + Examples: + >>> # Usage with device_mesh and alias format + >>> mesh = DeviceMesh((2, 2), ("dp", "tp")) + >>> sharding_plan = { + ... "parameter": {"mlp.weight": ("None", "tp")}, + ... "forward": {"input": ("dp", "None"), + ... "output": ("dp", "tp")} + ... } + >>> model = shard(model, mesh, sharding_plan) + + >>> # Usage with device_mesh and Placement format + >>> mesh = DeviceMesh((2, 2), ("dp", "tp")) + >>> sharding_plan = { + ... "parameter": {"mlp.weight": (Replicate(), Shard(1))}, + ... "forward": {"input": (Shard(0), Replicate()), + ... "output": (Shard(0), Shard(1))} + ... } + >>> model = shard(model, mesh, sharding_plan) """ if platform.get_world_size() == 1: return None + + # Convert sharding_plan to Layout objects + converted_plan = _convert_sharding_plan(sharding_plan, device_mesh) + if not isinstance(model, Module): - return _shard_callable(model, sharding_plan) + return _shard_callable(model, converted_plan) - param_sharding_plan = sharding_plan.get("parameter") - forward_sharding_plan = sharding_plan.get("forward") + param_sharding_plan = converted_plan.get("parameter") + forward_sharding_plan = converted_plan.get("forward") if param_sharding_plan is not None: for param_name, layout in param_sharding_plan.items(): @@ -222,7 +317,7 @@ def shard(model: Union[Module, Callable], sharding_plan: Dict): if not result: raise ValueError(f"{param_name} is configured with a layout, but no instance was found.") _, _, param = result - + layout.placement_to_tensor_map(param.dim()) param = platform.set_layout_into_parameter(param, layout) platform.update_parameter_by_name(model, result, param) diff --git a/hyper_parallel/core/shard/local_func.py b/hyper_parallel/core/shard/local_func.py index 2b3f685..8b5d8fe 100644 --- a/hyper_parallel/core/shard/local_func.py +++ b/hyper_parallel/core/shard/local_func.py @@ -15,17 +15,20 @@ """shard""" import queue -from typing import Callable, Tuple, Optional -from hyper_parallel.core.layout import Layout +from typing import Callable, Tuple, Optional, List +from hyper_parallel.core.layout import DeviceMesh from hyper_parallel.core.dtensor import DTensor +from hyper_parallel.core.placement_types import Placement from hyper_parallel.platform import get_platform platform = get_platform() Tensor = platform.Tensor + def custom_shard( func: Callable, - out_layouts: Tuple[Layout, ...], - in_layouts: Optional[Tuple[Optional[Layout], ...]] = None, + device_mesh: DeviceMesh, + out_placements: Tuple[Tuple[Placement, ...], ...], + in_placements: Optional[Tuple[Optional[Tuple[Placement, ...]], ...]] = None, redistribute_inputs: bool = True, ) -> Callable: """ @@ -33,18 +36,30 @@ def custom_shard( Args: func (Callable): The function to be wrapped. - out_layouts (Tuple[Layout, ...]): Layouts for each output tensor. - in_layouts (Optional[Tuple[Optional[Layout], ...]], optional): - Layouts for each input argument. None entries indicate non-tensor inputs. - redistribute_inputs (bool): Whether to redistribute inputs to required layouts. + device_mesh (DeviceMesh): The device mesh for sharding. + out_placements (Tuple[Tuple[Placement, ...], ...]): Placements for each output tensor. + in_placements (Optional[Tuple[Optional[Tuple[Placement, ...]], ...]], optional): + Placements for each input argument. None entries indicate non-tensor inputs. + redistribute_inputs (bool): Whether to redistribute inputs to required placements. Returns: Callable: Wrapped function that handles distributed tensors. + + Examples: + >>> mesh = DeviceMesh((2, 2), ("dp", "tp")) + >>> @custom_shard( + ... device_mesh=mesh, + ... out_placements=((Shard(0), Replicate()),), + ... in_placements=((Shard(0), Replicate()), (Replicate(), Shard(1))) + ... ) + ... def my_func(x, y): + ... return x + y """ + def wrapped(*args, **kwargs): - if in_layouts is not None: - assert len(in_layouts) == len(args), ( - f"in_placements length {len(in_layouts)} does not match " + if in_placements is not None: + assert len(in_placements) == len(args), ( + f"in_placements length {len(in_placements)} does not match " f"the number of input args {len(args)}!" ) @@ -54,18 +69,18 @@ def custom_shard( args_layout = queue.Queue(len(args)) for i, arg in enumerate(args): if isinstance(arg, DTensor): - if in_layouts is None: - raise RuntimeError("Found Tensor input but in_layouts is None") + if in_placements is None: + raise RuntimeError("Found Tensor input but in_placements is None") - required_in_layout = in_layouts[i] - if required_in_layout is None: + required_in_placement = in_placements[i] + if required_in_placement is None: raise TypeError( - f"Tensor input at position {i} requires Layout, " - "but corresponding in_layouts entry is None!" + f"Tensor input at position {i} requires Placement, " + "but corresponding in_placements entry is None!" ) if redistribute_inputs: - arg = arg.redistribute(required_in_layout) + arg = arg.redistribute(device_mesh, list(required_in_placement)) args_layout.put(arg.layout) local_tensor = arg.to_local() @@ -73,10 +88,10 @@ def custom_shard( contain_distributed_arg = True else: - if in_layouts is not None and in_layouts[i] is not None: + if in_placements is not None and in_placements[i] is not None: raise TypeError( - f"Non-DTensor input at position {i} requires None in_layouts, " - f"but received {in_layouts[i]}!" + f"Non-DTensor input at position {i} requires None in_placements, " + f"but received {in_placements[i]}!" ) local_args.append(arg) @@ -88,26 +103,28 @@ def custom_shard( out_is_tuple = isinstance(out, tuple) out_tuple = (out,) if not out_is_tuple else out - assert len(out_tuple) == len(out_layouts), ( + assert len(out_tuple) == len(out_placements), ( f"Output count {len(out_tuple)} does not match " - f"out_layouts count {len(out_layouts)}!" + f"out_placements count {len(out_placements)}!" ) dist_output = [] - for item, out_layout in zip(out_tuple, out_layouts): + for item, out_placement in zip(out_tuple, out_placements): if isinstance(item, Tensor): - if out_layout is None: + if out_placement is None: raise TypeError( - "Tensor output requires non-None out_layout!" + "Tensor output requires non-None out_placements!" ) - dist_output.append(DTensor.from_local(item, out_layout)) + dist_output.append( + DTensor.from_local(item, device_mesh=device_mesh, placements=out_placement) + ) else: - if out_layout is not None: + if out_placement is not None: raise TypeError( - f"Non-tensor output requires None out_layout, got {out_layout}!" + f"Non-tensor output requires None out_placements, got {out_placement}!" ) dist_output.append(item) return dist_output[0] if not out_is_tuple else tuple(dist_output) - return wrapped + return wrapped \ No newline at end of file diff --git a/hyper_parallel/core/tensor_redistribution.py b/hyper_parallel/core/tensor_redistribution.py index aa6c5c1..6f03907 100644 --- a/hyper_parallel/core/tensor_redistribution.py +++ b/hyper_parallel/core/tensor_redistribution.py @@ -199,7 +199,7 @@ class TensorRedistribution: transform_operator_list = self._transform_cache[key] for transform_operator in transform_operator_list: x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1]) - return DTensor.from_local(x, to_layout) + return DTensor.from_local(x, to_layout.mesh, to_layout.placements) full_shape = x.shape key_and_shape = key + str(full_shape) @@ -208,7 +208,7 @@ class TensorRedistribution: transform_operator_list = self._transform_cache[key_and_shape] for transform_operator in transform_operator_list: x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1]) - return DTensor.from_local(x, to_layout) + return DTensor.from_local(x, to_layout.mesh, to_layout.placements) if self._apply_eazy_redistribute(from_layout, to_layout): if from_layout.is_partial: @@ -219,7 +219,7 @@ class TensorRedistribution: full_shape, key_and_shape) for transform_operator in transform_operator_list: x = self._construct_op_operator[transform_operator[0]](x, *transform_operator[1]) - return DTensor.from_local(x, to_layout) + return DTensor.from_local(x, to_layout.mesh, to_layout.placements) def _infer_transform_operator_list(self, from_layout, to_layout, from_full_shape, key): """infer transform operator list""" @@ -314,7 +314,7 @@ class TensorRedistribution: output_layout = from_layout(*output_alias_tensor_map) output_layout.reset_partial() - return DTensor.from_local(x, output_layout) + return DTensor.from_local(x, output_layout.mesh, output_layout.placements) _tensor_redistribution = TensorRedistribution() diff --git a/hyper_parallel/platform/mindspore/dtensor.py b/hyper_parallel/platform/mindspore/dtensor.py index 6ec8773..6351acd 100644 --- a/hyper_parallel/platform/mindspore/dtensor.py +++ b/hyper_parallel/platform/mindspore/dtensor.py @@ -20,24 +20,38 @@ from mindspore._c_expression import NoFallbackGuard class DTensorBase(Tensor): """ - DTensorBase + DTensorBase - Base class for distributed tensors in MindSpore. + + This class extends Tensor to support distributed tensor operations with + device mesh and placement specifications. """ - def __new__(cls, local_tensor, layout=None, device="Ascend"): + def __new__(cls, local_tensor, device_mesh=None, placements=None, device="Ascend"): + """ + Create a new DTensorBase instance. + + Args: + local_tensor: The local tensor shard or another DTensorBase instance. + device_mesh: The device mesh describing the device topology. + placements: The placement strategy for each mesh dimension. + device: The device type (default: "Ascend"). + """ if isinstance(local_tensor, DTensorBase): - device_local_tensor = local_tensor.to_local() if local_tensor.to_local().has_init else\ + device_local_tensor = local_tensor.to_local() if local_tensor.to_local().has_init else \ local_tensor.to_local().to(device) t = Tensor._make_subclass(cls, device_local_tensor) - t.__init_data__(device_local_tensor, local_tensor.layout) + t.__init_data__(device_local_tensor, local_tensor.device_mesh, local_tensor.placements) t._device = device return t - if not layout: - raise ValueError("Layout is None") + if device_mesh is None: + raise ValueError("device_mesh is None") + if placements is None: + raise ValueError("placements is None") device_local_tensor = local_tensor if local_tensor.has_init else local_tensor.to(device) if local_tensor.has_init: local_tensor.init_device = device t = Tensor._make_subclass(cls, device_local_tensor) - t.__init_data__(device_local_tensor, layout) + t.__init_data__(device_local_tensor, device_mesh, placements) t._device = device return t @@ -51,11 +65,39 @@ class DTensorBase(Tensor): return str(self._local_tensor) def __copy__(self): + """ + Create a shallow copy of the DTensorBase instance. + + This method ensures that device_mesh and placements are correctly + propagated when creating a copy (e.g., for optimizer states). + """ + # Get device_mesh and placements from either direct attributes or from layout + device_mesh = getattr(self, '_device_mesh', None) + placements = getattr(self, '_placements', None) + + # If not found directly, try to get from layout + if device_mesh is None and hasattr(self, '_layout') and self._layout is not None: + device_mesh = self._layout.mesh + if placements is None and hasattr(self, '_layout') and self._layout is not None: + placements = self._layout.placements + + if device_mesh is None or placements is None: + raise ValueError("Cannot copy DTensorBase: device_mesh or placements is None") + if self._local_tensor.has_init: - obj = DTensorBase.__new__(type(self), initializer(self._local_tensor.init, self._local_tensor.shape, - self._local_tensor.dtype), self._layout) + obj = DTensorBase.__new__( + type(self), + initializer(self._local_tensor.init, self._local_tensor.shape, self._local_tensor.dtype), + device_mesh, + placements + ) else: - obj = DTensorBase.__new__(type(self), self._local_tensor.clone(), self._layout) + obj = DTensorBase.__new__( + type(self), + self._local_tensor.clone(), + device_mesh, + placements + ) filtered_dict = {k: v for k, v in self.__dict__.items() if k != '_local_tensor'} obj.__dict__.update(filtered_dict) return obj @@ -93,6 +135,8 @@ class DTensorBase(Tensor): data = data.to(self._device) if isinstance(data, DTensorBase): self._local_tensor._update_data(data.to_local()) + self._device_mesh = data.device_mesh + self._placements = data.placements self._layout = data.layout self._update_data(self._local_tensor) return @@ -125,7 +169,6 @@ class DTensorBase(Tensor): return None return self._local_tensor.init - @init.setter def init(self, init_value): """ @@ -149,7 +192,6 @@ class DTensorBase(Tensor): return None return self._local_tensor.param_info - @local_param_info.setter def local_param_info(self, local_param_info_value): """ diff --git a/hyper_parallel/platform/mindspore/platform.py b/hyper_parallel/platform/mindspore/platform.py index e19d020..0ad964d 100644 --- a/hyper_parallel/platform/mindspore/platform.py +++ b/hyper_parallel/platform/mindspore/platform.py @@ -188,13 +188,13 @@ class MindSporePlatform(Platform): if not param.has_init: # has been init, get slice data - param_dtensor = DTensor.from_local(_get_slice_tensor_by_layout(param, layout).value(), layout) + param_dtensor = DTensor.from_local(_get_slice_tensor_by_layout(param, layout).value(), layout.mesh, layout.placements) param = Parameter(param_dtensor, name=name, requires_grad=requires_grad) param.param_info = param_info else: # has not been init, need to modify init shape param.init_mode.shape = slice_shape - param_dtensor = DTensor.from_local(param.init_mode, layout) + param_dtensor = DTensor.from_local(param.init_mode, layout.mesh, layout.placements) param = Parameter(param_dtensor, name=name, requires_grad=requires_grad) param.param_info = param_info return param diff --git a/tests/mindspore/st/shard/base_shard.py b/tests/mindspore/st/shard/base_shard.py index 108d060..c337cd7 100644 --- a/tests/mindspore/st/shard/base_shard.py +++ b/tests/mindspore/st/shard/base_shard.py @@ -22,7 +22,8 @@ import mindspore.communication.management as D from mindspore import nn, Tensor from mindspore.nn.utils import no_init_parameters from mindspore.common.initializer import initializer -from hyper_parallel import Layout, hsdp, init_parameters, shard, parallelize_value_and_grad +from hyper_parallel import init_device_mesh, hsdp, init_parameters, shard, parallelize_value_and_grad +from hyper_parallel.core.placement_types import Shard, Replicate from tests.mindspore.st.shard.utils import create_dtensor learning_rate = 0.01 @@ -86,23 +87,43 @@ def base_case(dp, mp, hsdp_shard_size): local_input_size = input_size // mp local_output_size = output_size local_x = np.ones([local_batch_size, local_input_size]).astype(np.float32) - layout = Layout((dp, mp), ("dp", "mp")) - x_layout = layout("dp", "mp") - w_layout = layout("mp", "None") - out_layout = layout() - relu_strategy = ((layout("dp", "None"),), (layout("dp", "None"),)) + + # Create DeviceMesh + mesh = init_device_mesh( + mesh_shape=(dp, mp), + alias_name=("dp", "tp") + ) + + # Define placements using Placement format + x_placements = (Shard(0), Shard(1)) + w_placements = (Replicate(), Shard(0)) + out_placements = (Replicate(), Replicate()) + relu_input_placements = (Shard(0), Replicate()) + relu_output_placements = (Shard(0), Replicate()) # step 1: define network with no init parameters with no_init_parameters(): model = SimpleModel(input_size, output_size) # step 2: shard - model_stra = { "forward": { "input": (x_layout,), "output": (out_layout,)}, - "parameter": {"weight": w_layout}} - shard(model, model_stra) - - model_relu_stra = { "forward": { "input": relu_strategy[0], "output": relu_strategy[1]}} - shard(model.relu, model_relu_stra) + model_sharding_plan = { + "forward": { + "input": x_placements, + "output": out_placements + }, + "parameter": { + "weight": w_placements + } + } + shard(model, device_mesh=mesh, sharding_plan=model_sharding_plan) + + model_relu_sharding_plan = { + "forward": { + "input": relu_input_placements, + "output": relu_output_placements, + } + } + shard(model.relu, device_mesh=mesh, sharding_plan=model_relu_sharding_plan) # step 3: hsdp model = hsdp(model, shard_size=hsdp_shard_size, threshold=0) @@ -110,7 +131,7 @@ def base_case(dp, mp, hsdp_shard_size): # step 4: init parameters model = init_parameters(model) - x = create_dtensor(local_x, x_layout) + x = create_dtensor(local_x, mesh, x_placements) parallel_loss, parallel_grads = run_model(x, model, parallel=True) # compare loss diff --git a/tests/mindspore/st/shard/utils.py b/tests/mindspore/st/shard/utils.py index 8526050..bb3ff8d 100644 --- a/tests/mindspore/st/shard/utils.py +++ b/tests/mindspore/st/shard/utils.py @@ -18,14 +18,14 @@ from typing import List, Optional from mindspore import Tensor, mint from mindspore.communication import get_rank, create_group from hyper_parallel import Layout, DTensor -from hyper_parallel.core.layout import _DeviceMesh +from hyper_parallel.core.device_mesh import DeviceMesh -def create_dtensor(data, layout, dtype=None): +def create_dtensor(data, device_mesh, placements, dtype=None): """create_dtensor""" if dtype: - return DTensor.from_local(Tensor(data, dtype=dtype), layout) - return DTensor.from_local(Tensor(data), layout) + return DTensor.from_local(Tensor(data, dtype=dtype), device_mesh, placements) + return DTensor.from_local(Tensor(data), device_mesh, placements) def _infer_slice_area_by_rank(mesh_shape, tensor_map, rank_id: int, full_shape: tuple): # -> tuple[tuple[int]]: @@ -99,7 +99,7 @@ def local_to_global(local_tensor): return local_tensor.redistribute(to_layout).to_local() -def mesh_scatter(output: Tensor, scatter_list: List[Tensor], dev_mesh: _DeviceMesh, +def mesh_scatter(output: Tensor, scatter_list: List[Tensor], dev_mesh: DeviceMesh, mesh_alias: str, group_src_rank: Optional[int] = 0): """mesh_scatter""" rank_id = get_rank() @@ -114,7 +114,7 @@ def mesh_scatter(output: Tensor, scatter_list: List[Tensor], dev_mesh: _DeviceMe return output -def mesh_broadcast(tensor: Tensor, dev_mesh: _DeviceMesh, +def mesh_broadcast(tensor: Tensor, dev_mesh: DeviceMesh, mesh_alias: str, group_src_rank: Optional[int] = 0): """mesh_broadcast""" rank_id = get_rank() diff --git a/tests/mindspore/st/test_placement_conversion.py b/tests/mindspore/st/test_placement_conversion.py index c5f7fe4..85e0d58 100644 --- a/tests/mindspore/st/test_placement_conversion.py +++ b/tests/mindspore/st/test_placement_conversion.py @@ -15,7 +15,8 @@ """test placement and tensor_map conversion""" import pytest from tests.common.mark_utils import arg_mark -from hyper_parallel.core.layout import Layout, _DeviceMesh +from hyper_parallel.core.layout import Layout +from hyper_parallel.core.device_mesh import DeviceMesh from hyper_parallel.core.placement_types import Shard, Replicate, Partial @@ -184,7 +185,7 @@ def test_from_device_mesh(): mesh_shape = (4, 2) alias_name = ("dp", "mp") rank_list = tuple(range(8)) - device_mesh = _DeviceMesh(mesh_shape, alias_name, rank_list) + device_mesh = DeviceMesh(mesh_shape, alias_name, rank_list) layout = Layout.from_device_mesh(device_mesh) -- Gitee