Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / fully_shard / pack_utils.py: 93%
132 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 fully_shard communication buffers."""
17from __future__ import annotations
19import math
20from dataclasses import dataclass
21from typing import Any, Literal, Optional
23import torch
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: torch.Size
40 packed_tensor_shape: torch.Size
41 unpacked_shape: torch.Size
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 _has_strided_shard_layout(hsdp_param: Any) -> bool:
53 placements = getattr(hsdp_param, "_spmd_placements", ()) or ()
54 return any(isinstance(placement, StridedShard) for placement in placements)
57def _resolve_same_dim_strided_context(
58 hsdp_param: Any,
59) -> Optional[_SameDimStridedLayoutContext]:
60 if not _has_strided_shard_layout(hsdp_param):
61 return None
62 if not getattr(hsdp_param, "uses_param_shard", False):
63 return None
64 if not getattr(hsdp_param, "_orig_param_is_dtensor", False):
65 return None
66 target_dim = getattr(getattr(hsdp_param, "hsdp_placement", None), "dim", None)
67 if target_dim is None:
68 return None
69 shard_mesh_dim = getattr(hsdp_param, "_spmd_shard_mesh_dim", None)
70 placements = tuple(getattr(hsdp_param, "_spmd_placements", ()) or ())
71 if shard_mesh_dim is None or shard_mesh_dim >= len(placements):
72 return None
73 if not isinstance(placements[shard_mesh_dim], StridedShard):
74 return None
75 orig_placements = getattr(hsdp_param, "_orig_dtensor_placements", None)
76 if orig_placements is None:
77 return None
78 return _SameDimStridedLayoutContext(
79 target_dim=target_dim,
80 shard_mesh_dim=shard_mesh_dim,
81 placements=placements,
82 orig_placements=tuple(orig_placements),
83 )
86def _placements_match_target_dim_only(
87 placements: tuple[Any, ...],
88 target_dim: int,
89) -> bool:
90 return all(
91 placement.is_replicate() or placement.is_shard(target_dim)
92 for placement in placements
93 )
96def _orig_layout_is_supported(
97 orig_placements: tuple[Any, ...],
98 target_dim: int,
99) -> bool:
100 if not _placements_match_target_dim_only(orig_placements, target_dim):
101 return False
102 return sum(
103 placement.is_shard(target_dim) for placement in orig_placements
104 ) == 1
107def _current_strided_layout_is_supported(
108 placements: tuple[Any, ...],
109 target_dim: int,
110) -> bool:
111 if not _placements_match_target_dim_only(placements, target_dim):
112 return False
113 if sum(placement.is_shard() for placement in placements) != 2:
114 return False
116 strided_placements = [
117 placement for placement in placements if isinstance(placement, StridedShard)
118 ]
119 if len(strided_placements) != 1:
120 return False
121 strided_placement = strided_placements[0]
122 if strided_placement.dim != target_dim or strided_placement.split_factor <= 1:
123 return False
125 plain_shards = [
126 placement
127 for placement in placements
128 if placement.is_shard(target_dim) and not isinstance(placement, StridedShard)
129 ]
130 return len(plain_shards) == 1
133def supports_same_dim_strided_layout(hsdp_param: Any) -> bool:
134 """Check whether the parameter's StridedShard layout is supported for same-dim packing."""
135 ctx = _resolve_same_dim_strided_context(hsdp_param)
136 if ctx is None:
137 return False
138 if not _orig_layout_is_supported(ctx.orig_placements, ctx.target_dim):
139 return False
140 return _current_strided_layout_is_supported(ctx.placements, ctx.target_dim)
143def _resolve_unpacked_shape(
144 hsdp_param: Optional[Any],
145 local_tensor: torch.Tensor,
146) -> torch.Size:
147 if hsdp_param is not None and getattr(hsdp_param, "_orig_size", None) is not None:
148 return torch.Size(getattr(hsdp_param, "_orig_size"))
149 return torch.Size(local_tensor.size())
152def _get_packed_tensor_shape(
153 unpacked_shape: torch.Size,
154 shard_dim: int,
155 world_size: int,
156) -> torch.Size:
157 if world_size == 1 or shard_dim == 0:
158 return unpacked_shape
159 packed_tensor_shape = list(unpacked_shape)
160 packed_tensor_shape[0] *= world_size
161 packed_tensor_shape[shard_dim] //= world_size
162 return torch.Size(packed_tensor_shape)
165def build_rs_plan(
166 hsdp_param: Optional[Any],
167 local_tensor: torch.Tensor,
168 world_size: int,
169 *,
170 shard_dim: Optional[int] = None,
171) -> ReduceScatterPlan:
172 """Build the V1 reduce-scatter packing plan for a local gradient tensor."""
174 if world_size <= 0:
175 raise ValueError(f"world_size must be positive, but got {world_size}")
177 resolved_shard_dim = getattr(getattr(hsdp_param, "hsdp_placement", None), "dim", shard_dim)
178 if resolved_shard_dim is None:
179 raise ValueError("build_rs_plan requires either hsdp_param or shard_dim")
180 unpacked_shape = _resolve_unpacked_shape(hsdp_param, local_tensor)
181 if resolved_shard_dim < 0 or resolved_shard_dim >= len(unpacked_shape):
182 raise ValueError(
183 f"Invalid shard dim {resolved_shard_dim} for tensor shape {tuple(unpacked_shape)}"
184 )
185 if world_size == 1:
186 if not local_tensor.is_contiguous():
187 raise NotImplementedError(
188 "reduce_scatter_grad currently expects contiguous local gradients before packing."
189 )
190 return ReduceScatterPlan(
191 pack_kind="identity_dim0",
192 shard_dim=resolved_shard_dim,
193 world_size=world_size,
194 packed_shape=torch.Size((1, math.prod(unpacked_shape))),
195 packed_tensor_shape=unpacked_shape,
196 unpacked_shape=unpacked_shape,
197 )
198 if local_tensor.dim() == 0:
199 raise NotImplementedError("reduce_scatter_grad does not support scalar gradients.")
200 if unpacked_shape[resolved_shard_dim] % world_size != 0:
201 raise NotImplementedError(
202 f"reduce_scatter_grad currently only supports even sharding on dim={resolved_shard_dim}."
203 )
204 if not local_tensor.is_contiguous():
205 raise NotImplementedError(
206 "reduce_scatter_grad currently expects contiguous local gradients before packing."
207 )
209 pack_kind: Literal[
210 "identity_dim0",
211 "same_dim_strided_identity_dim0",
212 "chunk_cat_non_dim0",
213 ] = "identity_dim0"
214 if hsdp_param is not None and _has_strided_shard_layout(hsdp_param):
215 if not supports_same_dim_strided_layout(hsdp_param):
216 raise NotImplementedError(
217 "reduce_scatter_grad only supports same-dim StridedShard layouts "
218 "that restore a single contiguous TP-local shard on the fully_shard dimension."
219 )
220 if resolved_shard_dim == 0:
221 pack_kind = "same_dim_strided_identity_dim0"
222 else:
223 pack_kind = "chunk_cat_non_dim0"
224 elif resolved_shard_dim != 0:
225 pack_kind = "chunk_cat_non_dim0"
227 packed_tensor_shape = _get_packed_tensor_shape(
228 unpacked_shape,
229 resolved_shard_dim,
230 world_size,
231 )
232 total_numel = math.prod(unpacked_shape)
234 return ReduceScatterPlan(
235 pack_kind=pack_kind,
236 shard_dim=resolved_shard_dim,
237 world_size=world_size,
238 packed_shape=torch.Size((world_size, total_numel // world_size)),
239 packed_tensor_shape=packed_tensor_shape,
240 unpacked_shape=unpacked_shape,
241 )
244def pack_for_reduce_scatter(
245 local_tensor: torch.Tensor,
246 plan: ReduceScatterPlan,
247) -> torch.Tensor:
248 """Pack one local gradient into the row-major reduce-scatter layout."""
250 if plan.pack_kind not in (
251 "identity_dim0",
252 "same_dim_strided_identity_dim0",
253 "chunk_cat_non_dim0",
254 ):
255 raise NotImplementedError(f"Unsupported reduce-scatter pack kind: {plan.pack_kind}")
256 if not local_tensor.is_contiguous():
257 raise NotImplementedError(
258 "reduce_scatter_grad currently expects contiguous local gradients before packing."
259 )
260 if local_tensor.size() != plan.unpacked_shape:
261 raise AssertionError(
262 "pack_for_reduce_scatter expects the unsharded local tensor shape to match "
263 f"plan.unpacked_shape, but got {tuple(local_tensor.size())} and "
264 f"{tuple(plan.unpacked_shape)}"
265 )
266 if plan.pack_kind == "chunk_cat_non_dim0":
267 chunks = torch.chunk(local_tensor, plan.world_size, dim=plan.shard_dim)
268 packed_tensor = torch.cat(chunks, dim=0)
269 return packed_tensor.contiguous().view(plan.packed_shape)
270 return local_tensor.view(plan.packed_shape)
273def unpack_from_all_gather(
274 full_packed: torch.Tensor,
275 plan: ReduceScatterPlan,
276) -> torch.Tensor:
277 """Inverse of the V1 reduce-scatter packing plan for all-gather outputs."""
279 if plan.pack_kind not in (
280 "identity_dim0",
281 "same_dim_strided_identity_dim0",
282 "chunk_cat_non_dim0",
283 ):
284 raise NotImplementedError(f"Unsupported all-gather unpack kind: {plan.pack_kind}")
285 packed_tensor = full_packed.view(plan.packed_tensor_shape)
286 if plan.pack_kind == "chunk_cat_non_dim0":
287 chunks = torch.chunk(packed_tensor, plan.world_size, dim=0)
288 return torch.cat(chunks, dim=plan.shard_dim).contiguous()
289 return packed_tensor.view(plan.unpacked_shape)
292__all__ = [
293 "ReduceScatterPlan",
294 "build_rs_plan",
295 "pack_for_reduce_scatter",
296 "unpack_from_all_gather",
297 "supports_same_dim_strided_layout",
298]