Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / fully_shard / pack_utils.py: 91%
134 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Packing helpers for MindSpore fully_shard communication buffers."""
17from __future__ import annotations
19import math
20from dataclasses import dataclass
21from typing import Any, Literal, Optional
23import mindspore as ms
25from hyper_parallel.core.dtensor.placement_types import StridedShard
28@dataclass(frozen=True)
29class ReduceScatterPlan:
30 """Describe how local tensors map to packed communication layouts."""
32 pack_kind: Literal[
33 "identity_dim0",
34 "same_dim_strided_identity_dim0",
35 "chunk_cat_non_dim0",
36 ]
37 shard_dim: int
38 world_size: int
39 packed_shape: tuple[int, ...]
40 packed_tensor_shape: tuple[int, ...]
41 unpacked_shape: tuple[int, ...]
44@dataclass(frozen=True)
45class _SameDimStridedLayoutContext:
46 target_dim: int
47 shard_mesh_dim: int
48 placements: tuple[Any, ...]
49 orig_placements: tuple[Any, ...]
52def _shape_tuple(shape) -> tuple[int, ...]:
53 return tuple(int(dim) for dim in shape)
56def _has_strided_shard_layout(hsdp_param: Any) -> bool:
57 placements = getattr(hsdp_param, "_spmd_placements", ()) or ()
58 return any(isinstance(placement, StridedShard) for placement in placements)
61def _resolve_same_dim_strided_context(
62 hsdp_param: Any,
63) -> Optional[_SameDimStridedLayoutContext]:
64 if not _has_strided_shard_layout(hsdp_param):
65 return None
66 if not getattr(hsdp_param, "uses_param_shard", False):
67 return None
68 if not getattr(hsdp_param, "_orig_param_is_dtensor", False):
69 return None
70 target_dim = getattr(getattr(hsdp_param, "hsdp_placement", None), "dim", None)
71 if target_dim is None:
72 return None
73 shard_mesh_dim = getattr(hsdp_param, "_spmd_shard_mesh_dim", None)
74 placements = tuple(getattr(hsdp_param, "_spmd_placements", ()) or ())
75 if shard_mesh_dim is None or shard_mesh_dim >= len(placements):
76 return None
77 if not isinstance(placements[shard_mesh_dim], StridedShard):
78 return None
79 orig_placements = getattr(hsdp_param, "_orig_dtensor_placements", None)
80 if orig_placements is None:
81 return None
82 return _SameDimStridedLayoutContext(
83 target_dim=target_dim,
84 shard_mesh_dim=shard_mesh_dim,
85 placements=placements,
86 orig_placements=tuple(orig_placements),
87 )
90def _placements_match_target_dim_only(
91 placements: tuple[Any, ...],
92 target_dim: int,
93) -> bool:
94 return all(
95 placement.is_replicate() or placement.is_shard(target_dim)
96 for placement in placements
97 )
100def _orig_layout_is_supported(
101 orig_placements: tuple[Any, ...],
102 target_dim: int,
103) -> bool:
104 if not _placements_match_target_dim_only(orig_placements, target_dim):
105 return False
106 return sum(
107 placement.is_shard(target_dim) for placement in orig_placements
108 ) == 1
111def _current_strided_layout_is_supported(
112 placements: tuple[Any, ...],
113 target_dim: int,
114) -> bool:
115 if not _placements_match_target_dim_only(placements, target_dim):
116 return False
117 if sum(placement.is_shard() for placement in placements) != 2:
118 return False
120 strided_placements = [
121 placement for placement in placements if isinstance(placement, StridedShard)
122 ]
123 if len(strided_placements) != 1:
124 return False
125 strided_placement = strided_placements[0]
126 if strided_placement.dim != target_dim or strided_placement.split_factor <= 1:
127 return False
129 plain_shards = [
130 placement
131 for placement in placements
132 if placement.is_shard(target_dim) and not isinstance(placement, StridedShard)
133 ]
134 return len(plain_shards) == 1
137def supports_same_dim_strided_layout(hsdp_param: Any) -> bool:
138 """Check whether the parameter's same-dimension StridedShard layout is supported for packing."""
139 ctx = _resolve_same_dim_strided_context(hsdp_param)
140 if ctx is None:
141 return False
142 if not _orig_layout_is_supported(ctx.orig_placements, ctx.target_dim):
143 return False
144 return _current_strided_layout_is_supported(ctx.placements, ctx.target_dim)
147def _resolve_unpacked_shape(
148 hsdp_param: Optional[Any],
149 local_tensor: ms.Tensor,
150) -> tuple[int, ...]:
151 if hsdp_param is not None and getattr(hsdp_param, "_orig_size", None) is not None:
152 return _shape_tuple(getattr(hsdp_param, "_orig_size"))
153 return _shape_tuple(local_tensor.shape)
156def _get_packed_tensor_shape(
157 unpacked_shape: tuple[int, ...],
158 shard_dim: int,
159 world_size: int,
160) -> tuple[int, ...]:
161 if world_size == 1 or shard_dim == 0:
162 return unpacked_shape
163 packed_tensor_shape = list(unpacked_shape)
164 packed_tensor_shape[0] *= world_size
165 packed_tensor_shape[shard_dim] //= world_size
166 return tuple(packed_tensor_shape)
169def build_rs_plan(
170 hsdp_param: Optional[Any],
171 local_tensor: ms.Tensor,
172 world_size: int,
173 *,
174 shard_dim: Optional[int] = None,
175) -> ReduceScatterPlan:
176 """Build the V1 reduce-scatter packing plan for a local gradient tensor."""
178 if world_size <= 0:
179 raise ValueError(f"world_size must be positive, but got {world_size}")
181 resolved_shard_dim = getattr(getattr(hsdp_param, "hsdp_placement", None), "dim", shard_dim)
182 if resolved_shard_dim is None:
183 raise ValueError("build_rs_plan requires either hsdp_param or shard_dim")
184 unpacked_shape = _resolve_unpacked_shape(hsdp_param, local_tensor)
185 if resolved_shard_dim < 0 or resolved_shard_dim >= len(unpacked_shape):
186 raise ValueError(
187 f"Invalid shard dim {resolved_shard_dim} for tensor shape {tuple(unpacked_shape)}"
188 )
189 if world_size == 1:
190 if not local_tensor.is_contiguous():
191 raise NotImplementedError(
192 "reduce_scatter_grad currently expects contiguous local gradients before packing."
193 )
194 return ReduceScatterPlan(
195 pack_kind="identity_dim0",
196 shard_dim=resolved_shard_dim,
197 world_size=world_size,
198 packed_shape=(1, math.prod(unpacked_shape)),
199 packed_tensor_shape=unpacked_shape,
200 unpacked_shape=unpacked_shape,
201 )
202 if len(local_tensor.shape) == 0:
203 raise NotImplementedError("reduce_scatter_grad does not support scalar gradients.")
204 if unpacked_shape[resolved_shard_dim] % world_size != 0:
205 raise NotImplementedError(
206 f"reduce_scatter_grad currently only supports even sharding on dim={resolved_shard_dim}."
207 )
208 if not local_tensor.is_contiguous():
209 raise NotImplementedError(
210 "reduce_scatter_grad currently expects contiguous local gradients before packing."
211 )
213 pack_kind: Literal[
214 "identity_dim0",
215 "same_dim_strided_identity_dim0",
216 "chunk_cat_non_dim0",
217 ] = "identity_dim0"
218 if hsdp_param is not None and _has_strided_shard_layout(hsdp_param):
219 if not supports_same_dim_strided_layout(hsdp_param):
220 raise NotImplementedError(
221 "reduce_scatter_grad only supports same-dim StridedShard layouts "
222 "that restore a single contiguous TP-local shard on the fully_shard dimension."
223 )
224 if resolved_shard_dim == 0:
225 pack_kind = "same_dim_strided_identity_dim0"
226 else:
227 pack_kind = "chunk_cat_non_dim0"
228 elif resolved_shard_dim != 0:
229 pack_kind = "chunk_cat_non_dim0"
231 packed_tensor_shape = _get_packed_tensor_shape(
232 unpacked_shape,
233 resolved_shard_dim,
234 world_size,
235 )
236 total_numel = math.prod(unpacked_shape)
237 return ReduceScatterPlan(
238 pack_kind=pack_kind,
239 shard_dim=resolved_shard_dim,
240 world_size=world_size,
241 packed_shape=(world_size, total_numel // world_size),
242 packed_tensor_shape=packed_tensor_shape,
243 unpacked_shape=unpacked_shape,
244 )
247def pack_for_reduce_scatter(
248 local_tensor: ms.Tensor,
249 plan: ReduceScatterPlan,
250) -> ms.Tensor:
251 """Pack one local gradient into the row-major reduce-scatter layout."""
253 if plan.pack_kind not in (
254 "identity_dim0",
255 "same_dim_strided_identity_dim0",
256 "chunk_cat_non_dim0",
257 ):
258 raise NotImplementedError(f"Unsupported reduce-scatter pack kind: {plan.pack_kind}")
259 if not local_tensor.is_contiguous():
260 raise NotImplementedError(
261 "reduce_scatter_grad currently expects contiguous local gradients before packing."
262 )
263 if _shape_tuple(local_tensor.shape) != plan.unpacked_shape:
264 raise AssertionError(
265 "pack_for_reduce_scatter expects the unsharded local tensor shape to match "
266 f"plan.unpacked_shape, but got {tuple(local_tensor.shape)} and "
267 f"{tuple(plan.unpacked_shape)}"
268 )
269 if plan.pack_kind == "chunk_cat_non_dim0":
270 chunks = ms.mint.chunk(local_tensor, plan.world_size, dim=plan.shard_dim)
271 packed_tensor = ms.mint.cat(chunks, dim=0)
272 return packed_tensor.contiguous().view(plan.packed_shape)
273 return local_tensor.view(plan.packed_shape)
276def unpack_from_all_gather(
277 full_packed: ms.Tensor,
278 plan: ReduceScatterPlan,
279) -> ms.Tensor:
280 """Inverse of the V1 reduce-scatter packing plan for all-gather outputs."""
282 if plan.pack_kind not in (
283 "identity_dim0",
284 "same_dim_strided_identity_dim0",
285 "chunk_cat_non_dim0",
286 ):
287 raise NotImplementedError(f"Unsupported all-gather unpack kind: {plan.pack_kind}")
288 packed_tensor = full_packed.view(plan.packed_tensor_shape)
289 if plan.pack_kind == "chunk_cat_non_dim0":
290 chunks = ms.mint.chunk(packed_tensor, plan.world_size, dim=0)
291 return ms.mint.cat(chunks, dim=plan.shard_dim).contiguous()
292 return packed_tensor.view(plan.unpacked_shape)
295__all__ = [
296 "ReduceScatterPlan",
297 "build_rs_plan",
298 "pack_for_reduce_scatter",
299 "unpack_from_all_gather",
300 "supports_same_dim_strided_layout",
301]