Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / fully_shard / pack_utils.py: 91%
135 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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
26from hyper_parallel.core.fully_shard.utils import FSDPMeshInfo
29@dataclass(frozen=True)
30class ReduceScatterPlan:
31 """Describe how local tensors map to packed communication layouts."""
33 pack_kind: Literal[
34 "identity_dim0",
35 "same_dim_strided_identity_dim0",
36 "chunk_cat_non_dim0",
37 ]
38 shard_dim: int
39 world_size: int
40 packed_shape: tuple[int, ...]
41 packed_tensor_shape: tuple[int, ...]
42 unpacked_shape: tuple[int, ...]
45@dataclass(frozen=True)
46class _SameDimStridedLayoutContext:
47 target_dim: int
48 shard_mesh_dim: int
49 placements: tuple[Any, ...]
50 orig_placements: tuple[Any, ...]
53def _shape_tuple(shape) -> tuple[int, ...]:
54 return tuple(int(dim) for dim in shape)
57def _has_strided_shard_layout(hsdp_param: Any) -> bool:
58 placements = getattr(hsdp_param, "_spmd_placements", ()) or ()
59 return any(isinstance(placement, StridedShard) for placement in placements)
62def _resolve_same_dim_strided_context(
63 hsdp_param: Any,
64) -> Optional[_SameDimStridedLayoutContext]:
65 if not _has_strided_shard_layout(hsdp_param):
66 return None
67 if not isinstance(getattr(hsdp_param, "mesh_info", None), FSDPMeshInfo):
68 return None
69 if not getattr(hsdp_param, "_orig_param_is_dtensor", False):
70 return None
71 target_dim = getattr(getattr(hsdp_param, "hsdp_placement", None), "dim", None)
72 if target_dim is None:
73 return None
74 shard_mesh_dim = getattr(hsdp_param, "_spmd_shard_mesh_dim", None)
75 placements = tuple(getattr(hsdp_param, "_spmd_placements", ()) or ())
76 if shard_mesh_dim is None or shard_mesh_dim >= len(placements):
77 return None
78 if not isinstance(placements[shard_mesh_dim], StridedShard):
79 return None
80 orig_placements = getattr(hsdp_param, "_orig_dtensor_placements", None)
81 if orig_placements is None:
82 return None
83 return _SameDimStridedLayoutContext(
84 target_dim=target_dim,
85 shard_mesh_dim=shard_mesh_dim,
86 placements=placements,
87 orig_placements=tuple(orig_placements),
88 )
91def _placements_match_target_dim_only(
92 placements: tuple[Any, ...],
93 target_dim: int,
94) -> bool:
95 return all(
96 placement.is_replicate() or placement.is_shard(target_dim)
97 for placement in placements
98 )
101def _orig_layout_is_supported(
102 orig_placements: tuple[Any, ...],
103 target_dim: int,
104) -> bool:
105 if not _placements_match_target_dim_only(orig_placements, target_dim):
106 return False
107 return sum(
108 placement.is_shard(target_dim) for placement in orig_placements
109 ) == 1
112def _current_strided_layout_is_supported(
113 placements: tuple[Any, ...],
114 target_dim: int,
115) -> bool:
116 if not _placements_match_target_dim_only(placements, target_dim):
117 return False
118 if sum(placement.is_shard() for placement in placements) != 2:
119 return False
121 strided_placements = [
122 placement for placement in placements if isinstance(placement, StridedShard)
123 ]
124 if len(strided_placements) != 1:
125 return False
126 strided_placement = strided_placements[0]
127 if strided_placement.dim != target_dim or strided_placement.split_factor <= 1:
128 return False
130 plain_shards = [
131 placement
132 for placement in placements
133 if placement.is_shard(target_dim) and not isinstance(placement, StridedShard)
134 ]
135 return len(plain_shards) == 1
138def supports_same_dim_strided_layout(hsdp_param: Any) -> bool:
139 """Check whether the parameter's same-dimension StridedShard layout is supported for packing."""
140 ctx = _resolve_same_dim_strided_context(hsdp_param)
141 if ctx is None:
142 return False
143 if not _orig_layout_is_supported(ctx.orig_placements, ctx.target_dim):
144 return False
145 return _current_strided_layout_is_supported(ctx.placements, ctx.target_dim)
148def _resolve_unpacked_shape(
149 hsdp_param: Optional[Any],
150 local_tensor: ms.Tensor,
151) -> tuple[int, ...]:
152 if hsdp_param is not None and getattr(hsdp_param, "_orig_size", None) is not None:
153 return _shape_tuple(getattr(hsdp_param, "_orig_size"))
154 return _shape_tuple(local_tensor.shape)
157def _get_packed_tensor_shape(
158 unpacked_shape: tuple[int, ...],
159 shard_dim: int,
160 world_size: int,
161) -> tuple[int, ...]:
162 if world_size == 1 or shard_dim == 0:
163 return unpacked_shape
164 packed_tensor_shape = list(unpacked_shape)
165 packed_tensor_shape[0] *= world_size
166 packed_tensor_shape[shard_dim] //= world_size
167 return tuple(packed_tensor_shape)
170def build_rs_plan(
171 hsdp_param: Optional[Any],
172 local_tensor: ms.Tensor,
173 world_size: int,
174 *,
175 shard_dim: Optional[int] = None,
176) -> ReduceScatterPlan:
177 """Build the V1 reduce-scatter packing plan for a local gradient tensor."""
179 if world_size <= 0:
180 raise ValueError(f"world_size must be positive, but got {world_size}")
182 resolved_shard_dim = getattr(getattr(hsdp_param, "hsdp_placement", None), "dim", shard_dim)
183 if resolved_shard_dim is None:
184 raise ValueError("build_rs_plan requires either hsdp_param or shard_dim")
185 unpacked_shape = _resolve_unpacked_shape(hsdp_param, local_tensor)
186 if resolved_shard_dim < 0 or resolved_shard_dim >= len(unpacked_shape):
187 raise ValueError(
188 f"Invalid shard dim {resolved_shard_dim} for tensor shape {tuple(unpacked_shape)}"
189 )
190 if world_size == 1:
191 if not local_tensor.is_contiguous():
192 raise NotImplementedError(
193 "reduce_scatter_grad currently expects contiguous local gradients before packing."
194 )
195 return ReduceScatterPlan(
196 pack_kind="identity_dim0",
197 shard_dim=resolved_shard_dim,
198 world_size=world_size,
199 packed_shape=(1, math.prod(unpacked_shape)),
200 packed_tensor_shape=unpacked_shape,
201 unpacked_shape=unpacked_shape,
202 )
203 if len(local_tensor.shape) == 0:
204 raise NotImplementedError("reduce_scatter_grad does not support scalar gradients.")
205 if unpacked_shape[resolved_shard_dim] % world_size != 0:
206 raise NotImplementedError(
207 f"reduce_scatter_grad currently only supports even sharding on dim={resolved_shard_dim}."
208 )
209 if not local_tensor.is_contiguous():
210 raise NotImplementedError(
211 "reduce_scatter_grad currently expects contiguous local gradients before packing."
212 )
214 pack_kind: Literal[
215 "identity_dim0",
216 "same_dim_strided_identity_dim0",
217 "chunk_cat_non_dim0",
218 ] = "identity_dim0"
219 if hsdp_param is not None and _has_strided_shard_layout(hsdp_param):
220 if not supports_same_dim_strided_layout(hsdp_param):
221 raise NotImplementedError(
222 "reduce_scatter_grad only supports same-dim StridedShard layouts "
223 "that restore a single contiguous TP-local shard on the fully_shard dimension."
224 )
225 if resolved_shard_dim == 0:
226 pack_kind = "same_dim_strided_identity_dim0"
227 else:
228 pack_kind = "chunk_cat_non_dim0"
229 elif resolved_shard_dim != 0:
230 pack_kind = "chunk_cat_non_dim0"
232 packed_tensor_shape = _get_packed_tensor_shape(
233 unpacked_shape,
234 resolved_shard_dim,
235 world_size,
236 )
237 total_numel = math.prod(unpacked_shape)
238 return ReduceScatterPlan(
239 pack_kind=pack_kind,
240 shard_dim=resolved_shard_dim,
241 world_size=world_size,
242 packed_shape=(world_size, total_numel // world_size),
243 packed_tensor_shape=packed_tensor_shape,
244 unpacked_shape=unpacked_shape,
245 )
248def pack_for_reduce_scatter(
249 local_tensor: ms.Tensor,
250 plan: ReduceScatterPlan,
251) -> ms.Tensor:
252 """Pack one local gradient into the row-major reduce-scatter layout."""
254 if plan.pack_kind not in (
255 "identity_dim0",
256 "same_dim_strided_identity_dim0",
257 "chunk_cat_non_dim0",
258 ):
259 raise NotImplementedError(f"Unsupported reduce-scatter pack kind: {plan.pack_kind}")
260 if not local_tensor.is_contiguous():
261 raise NotImplementedError(
262 "reduce_scatter_grad currently expects contiguous local gradients before packing."
263 )
264 if _shape_tuple(local_tensor.shape) != plan.unpacked_shape:
265 raise AssertionError(
266 "pack_for_reduce_scatter expects the unsharded local tensor shape to match "
267 f"plan.unpacked_shape, but got {tuple(local_tensor.shape)} and "
268 f"{tuple(plan.unpacked_shape)}"
269 )
270 if plan.pack_kind == "chunk_cat_non_dim0":
271 chunks = ms.mint.chunk(local_tensor, plan.world_size, dim=plan.shard_dim)
272 packed_tensor = ms.mint.cat(chunks, dim=0)
273 return packed_tensor.contiguous().view(plan.packed_shape)
274 return local_tensor.view(plan.packed_shape)
277def unpack_from_all_gather(
278 full_packed: ms.Tensor,
279 plan: ReduceScatterPlan,
280) -> ms.Tensor:
281 """Inverse of the V1 reduce-scatter packing plan for all-gather outputs."""
283 if plan.pack_kind not in (
284 "identity_dim0",
285 "same_dim_strided_identity_dim0",
286 "chunk_cat_non_dim0",
287 ):
288 raise NotImplementedError(f"Unsupported all-gather unpack kind: {plan.pack_kind}")
289 packed_tensor = full_packed.view(plan.packed_tensor_shape)
290 if plan.pack_kind == "chunk_cat_non_dim0":
291 chunks = ms.mint.chunk(packed_tensor, plan.world_size, dim=0)
292 return ms.mint.cat(chunks, dim=plan.shard_dim).contiguous()
293 return packed_tensor.view(plan.unpacked_shape)
296__all__ = [
297 "ReduceScatterPlan",
298 "build_rs_plan",
299 "pack_for_reduce_scatter",
300 "unpack_from_all_gather",
301 "supports_same_dim_strided_layout",
302]