Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / sharding_category.py: 0%
208 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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# ============================================================================
16"""Category parameter with dtensor."""
18import itertools
19import logging
20from dataclasses import dataclass, field
21from typing import Dict, List, Sequence, Tuple, Optional
23import torch.distributed as dist
25from hyper_parallel.core.optimizer.dtensor_compat import (
26 DTensor,
27 DeviceMesh,
28 Shard,
29 StridedShard,
30)
32logger = logging.getLogger(__name__)
35@dataclass(frozen=True)
36class ParamLayoutSpec:
37 """Shape-free but ndim-aware parameter layout.
39 Same layout means:
40 1. same tensor ndim;
41 2. same shard mesh dims;
42 3. same shard tensor dims;
43 4. same replicate mesh dims.
45 Example:
46 tensor_ndim = 3
47 placements = (Shard(0), Replicate(), Replicate(), Replicate())
49 Then:
50 shard_axes = ((0, 0),)
51 replicate_mesh_dims = (1, 2, 3)
53 Meaning:
54 mesh dim 0 shards tensor dim 0;
55 mesh dim 1/2/3 are replicated.
56 """
58 tensor_ndim: int
60 # Each item is:
61 # (mesh_dim, tensor_dim)
62 shard_axes: Tuple[Tuple[int, int], ...]
64 # Mesh dims that are replicated.
65 replicate_mesh_dims: Tuple[int, ...]
67 @property
68 def shard_mesh_dims(self) -> Tuple[int, ...]:
69 """Return the mesh dimensions used for sharding."""
70 return tuple(mesh_dim for mesh_dim, _ in self.shard_axes)
72 @property
73 def shard_tensor_dims(self) -> Tuple[int, ...]:
74 """Return the tensor dimensions that are sharded."""
75 return tuple(tensor_dim for _, tensor_dim in self.shard_axes)
77 @property
78 def is_last2d_sharded(self) -> bool:
79 """Whether any shard axis falls on the last 2 tensor dimensions.
81 Newton-Schulz iteration operates on the last 2 dims.
82 If either dim is sharded, allgather is needed before NS.
83 """
84 for _, tensor_dim in self.shard_axes:
85 if tensor_dim >= self.tensor_ndim - 2:
86 return True
87 return False
90@dataclass(frozen=True)
91class CommDomainKey:
92 """Communication-domain key based on mesh dimensions.
94 This key describes the communication domain logically by relying on the
95 underlying DeviceMesh identifiers and the specific dims involved.
96 """
98 mesh_shape: Tuple[int, ...]
99 mesh_rank_list: Tuple[int, ...]
101 replicate_mesh_dims: Tuple[int, ...]
102 shard_mesh_dims: Tuple[int, ...]
104 @property
105 def has_replicate_redundancy(self) -> bool:
106 """Check if there is any replicate redundancy in the mesh."""
107 return len(self.replicate_mesh_dims) > 0
109 @property
110 def has_shard_group(self) -> bool:
111 """Check if there are any sharded mesh dimensions."""
112 return len(self.shard_mesh_dims) > 0
115@dataclass(frozen=True)
116class HSDPGroupKey:
117 """Final grouping key for HSDP.
119 Same key means:
120 1. same communication domain;
121 2. same tensor ndim;
122 3. same shard / replicate axis layout.
124 Shape is not part of this key.
125 """
127 comm_key: CommDomainKey
128 axis_spec: ParamLayoutSpec
131@dataclass
132class ParamShardSpec:
133 """Per-parameter shard metadata used during HSDP grouping."""
135 device_mesh: DeviceMesh
136 shard_mesh_dims: Tuple[int, ...]
137 replicate_mesh_dims: Tuple[int, ...]
138 replicate_pgs: Tuple[dist.ProcessGroup, ...]
139 shard_pgs: Tuple[dist.ProcessGroup, ...]
140 axis_spec: ParamLayoutSpec
143@dataclass(frozen=True)
144class ParamRecord:
145 """A lightweight parameter record."""
147 index: int
148 param: DTensor
149 shard_meta: Optional["ParamShardMeta"] = None
152@dataclass(frozen=True)
153class ParamShardStageMeta:
154 """Per-stage shard metadata for uneven all-gather and slicing."""
156 tensor_dim: int
157 shard_size: int
158 cur_rank_in_shard_group: int
159 rank_shapes: Tuple[Tuple[int, ...], ...]
160 split_sizes: Tuple[int, ...]
161 pad_shape: Tuple[int, ...]
164@dataclass(frozen=True)
165class ParamShardMeta:
166 """Runtime shard metadata for one DTensor parameter."""
168 global_shape: Tuple[int, ...]
169 global_numel: int
170 local_shape: Tuple[int, ...]
171 local_numel: int
172 stage_metas: Tuple[ParamShardStageMeta, ...]
175@dataclass
176class HSDPCommGroup:
177 """HSDP communication group.
179 Same group means same domain, ndim, and layout.
180 Stores sequences of ProcessGroups aligned with mesh dims.
181 """
183 comm_key: CommDomainKey
184 layout_spec: ParamLayoutSpec
186 # Tuple of runtime process groups, one for each relevant mesh dimension
187 replicate_pgs: Tuple[dist.ProcessGroup, ...] = ()
188 shard_pgs: Tuple[dist.ProcessGroup, ...] = ()
190 records: List[ParamRecord] = field(default_factory=list)
192 @property
193 def params(self) -> List[DTensor]:
194 """Return all DTensors within this communication group."""
195 return [record.param for record in self.records]
197 def add_param(self, index: int, param: DTensor) -> None:
198 """Add a parameter record to the group."""
199 self.records.append(
200 ParamRecord(
201 index=index,
202 param=param,
203 )
204 )
206 def __len__(self) -> int:
207 return len(self.records)
210def _normalize_tensor_dim(dim: int, tensor_ndim: int) -> int:
211 """Normalize tensor dim to non-negative dim."""
212 original_dim = dim
214 if dim < 0:
215 dim += tensor_ndim
217 if dim < 0 or dim >= tensor_ndim:
218 raise ValueError(
219 f"Invalid shard dim {original_dim} for tensor ndim {tensor_ndim}."
220 )
222 return dim
225def extract_param_shard_spec(dtensor: DTensor) -> ParamShardSpec:
226 """Extract all shard metadata and native process groups from one DTensor."""
227 device_mesh = dtensor.device_mesh
228 placements = dtensor.placements
229 tensor_ndim = len(dtensor.shape)
231 shard_mesh_dims: List[int] = []
232 replicate_mesh_dims: List[int] = []
233 shard_axes: List[Tuple[int, int]] = []
234 replicate_pgs: List[dist.ProcessGroup] = []
235 shard_pgs: List[dist.ProcessGroup] = []
237 for mesh_dim_idx, placement in enumerate(placements):
238 pg = device_mesh.get_group(mesh_dim_idx) if hasattr(device_mesh, "get_group") else None
240 if placement.is_replicate():
241 replicate_mesh_dims.append(mesh_dim_idx)
242 replicate_pgs.append(pg)
244 elif placement.is_shard():
245 if isinstance(placement, StridedShard):
246 placement_for_grouping = Shard(placement.dim)
247 else:
248 placement_for_grouping = placement
250 tensor_dim = _normalize_tensor_dim(
251 placement_for_grouping.dim,
252 tensor_ndim,
253 )
255 shard_mesh_dims.append(mesh_dim_idx)
256 shard_axes.append((mesh_dim_idx, tensor_dim))
257 shard_pgs.append(pg)
259 else:
260 raise ValueError(
261 f"Unsupported placement type in HSDP parameter grouping: "
262 f"{type(placement).__name__}."
263 )
265 axis_spec = ParamLayoutSpec(
266 tensor_ndim=tensor_ndim,
267 shard_axes=tuple(shard_axes),
268 replicate_mesh_dims=tuple(replicate_mesh_dims),
269 )
271 return ParamShardSpec(
272 device_mesh=device_mesh,
273 shard_mesh_dims=tuple(shard_mesh_dims),
274 replicate_mesh_dims=tuple(replicate_mesh_dims),
275 replicate_pgs=tuple(replicate_pgs),
276 shard_pgs=tuple(shard_pgs),
277 axis_spec=axis_spec,
278 )
281def build_comm_domain_key(shard_spec: ParamShardSpec) -> CommDomainKey:
282 """Build grouping key utilizing DeviceMesh properties."""
283 mesh_rank_list = ()
284 if hasattr(shard_spec.device_mesh, "rank_list"):
285 mesh_rank_list = tuple(shard_spec.device_mesh.rank_list)
287 return CommDomainKey(
288 mesh_shape=tuple(getattr(shard_spec.device_mesh, "mesh_shape", None) or shard_spec.device_mesh.mesh.shape),
289 mesh_rank_list=mesh_rank_list,
290 replicate_mesh_dims=shard_spec.replicate_mesh_dims,
291 shard_mesh_dims=shard_spec.shard_mesh_dims,
292 )
295def group_parameters_for_hsdp(
296 params: List[DTensor],
297) -> Tuple[List[DTensor], List[HSDPCommGroup]]:
298 """Group parameters relying on native DeviceMesh topology."""
299 no_comm_params: List[DTensor] = []
300 groups: Dict[HSDPGroupKey, HSDPCommGroup] = {}
302 for param_index, param in enumerate(params):
303 if not isinstance(param, DTensor):
304 no_comm_params.append(param)
305 continue
307 shard_spec = extract_param_shard_spec(param)
308 comm_key = build_comm_domain_key(shard_spec)
310 if not comm_key.has_replicate_redundancy and not comm_key.has_shard_group:
311 no_comm_params.append(param)
312 continue
314 group_key = HSDPGroupKey(
315 comm_key=comm_key,
316 axis_spec=shard_spec.axis_spec, # ParamLayoutSpec
317 )
319 if group_key not in groups:
320 groups[group_key] = HSDPCommGroup(
321 comm_key=comm_key,
322 layout_spec=shard_spec.axis_spec,
323 replicate_pgs=shard_spec.replicate_pgs,
324 shard_pgs=shard_spec.shard_pgs,
325 )
327 groups[group_key].add_param(param_index, param)
329 return no_comm_params, list(groups.values())
332@dataclass
333class HSDPGroupAssignment:
334 """Optimizer assignment for one HSDP communication group."""
336 owned_records: List[ParamRecord]
337 all_records: List[ParamRecord]
339 # param_index -> (dim_0_rank, dim_1_rank, ...)
340 owner_by_index: Dict[int, Tuple[int, ...]]
342 # Record the rank and size for each dimension, and the list of cur_rank within replicate_groups.
343 replicate_group_ranks: Tuple[int, ...]
344 replicate_sizes: Tuple[int, ...]
346 replicate_pgs: Tuple[dist.ProcessGroup, ...] = ()
347 shard_pgs: Tuple[dist.ProcessGroup, ...] = ()
349 is_shard: bool = False
350 layout_spec: Optional[ParamLayoutSpec] = None
352 @property
353 def owned_params(self) -> List[DTensor]:
354 """Return the parameters owned by the current rank."""
355 return [record.param for record in self.owned_records]
357 @property
358 def all_params(self) -> List[DTensor]:
359 """Return all parameters in this assignment group."""
360 return [record.param for record in self.all_records]
362 @property
363 def is_replicated(self) -> bool:
364 """Check if the group spans across multiple ranks."""
365 return any(s > 1 for s in self.replicate_sizes)
367 def owner_rank_coord(self, record: ParamRecord) -> Tuple[int, ...]:
368 """Get the rank coordinates of the owner of a given record."""
369 return self.owner_by_index.get(record.index, ())
371 def is_owned(self, record: ParamRecord) -> bool:
372 """Check if a given record is owned by the current replicate group rank."""
373 return self.owner_rank_coord(record) == self.replicate_group_ranks
375 @staticmethod
376 def _format_device_mesh(device_mesh: Optional[DeviceMesh]) -> Optional[str]:
377 """Format device mesh as DeviceMesh((name=size, ...)) for debug logs."""
378 if device_mesh is None:
379 return None
381 mesh_shape = tuple(getattr(device_mesh, "mesh_shape", None) or device_mesh.mesh.shape)
382 mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None)
384 if mesh_dim_names:
385 dims_desc = ", ".join(
386 f"{name}={size}" for name, size in zip(mesh_dim_names, mesh_shape)
387 )
388 else:
389 dims_desc = ", ".join(
390 f"dim_{idx}={size}" for idx, size in enumerate(mesh_shape)
391 )
393 return f"DeviceMesh(({dims_desc}))"
395 def __str__(self) -> str:
396 shard_ranks = [list(dist.get_process_group_ranks(pg)) for pg in self.shard_pgs if pg is not None]
397 replicate_ranks = [list(dist.get_process_group_ranks(pg)) for pg in self.replicate_pgs if pg is not None]
399 owned_names = [getattr(p, "model_name", f"p_{i}") for i, p in enumerate(self.owned_params)]
400 all_names = [getattr(p, "model_name", f"p_{i}") for i, p in enumerate(self.all_params)]
401 device_mesh = self.all_records[0].param.device_mesh if self.all_records else None
402 device_mesh_str = self._format_device_mesh(device_mesh)
404 return (
405 f"HSDPGroupAssignment( \n"
406 f"owned_params={owned_names}, \n"
407 f"all_params={all_names}, \n"
408 f"replicate_pg_ranks={replicate_ranks}, \n"
409 f"replicate_sizes={self.replicate_sizes}, \n"
410 f"is_shard={self.is_shard}, \n"
411 f"shard_pg_ranks={shard_ranks}, \n"
412 f"device_mesh={device_mesh_str}, \n"
413 f"layout_spec={self.layout_spec}) \n"
414 )
416 __repr__ = __str__
419def get_multi_dim_logical_info(
420 device_mesh: DeviceMesh,
421 mesh_dims: Sequence[int]
422) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
423 """Obtain the independent relative ranks and sizes of the parameters across multiple replicate dimensions."""
424 if not mesh_dims:
425 return (), ()
427 coords = device_mesh.get_coordinate()
428 if coords is None:
429 return (-1,) * len(mesh_dims), (1,) * len(mesh_dims)
431 ranks = tuple(coords[dim] for dim in mesh_dims)
432 sizes = tuple(device_mesh.size(dim) for dim in mesh_dims)
434 return ranks, sizes
437def build_owner_by_size(
438 records: List[ParamRecord],
439 replicate_sizes: Tuple[int, ...],
440) -> Dict[int, Tuple[int, ...]]:
441 """Build deterministic owner map across a multi-dimensional replicate grid."""
442 valid_records = [
443 record for record in records
444 if getattr(record.param, "requires_grad", True)
445 ]
447 if not valid_records:
448 return {}
450 if not replicate_sizes:
451 return {record.index: () for record in valid_records}
453 dim_ranges = [range(s) for s in replicate_sizes]
454 all_coords = list(itertools.product(*dim_ranges))
456 sorted_records = sorted(
457 valid_records,
458 key=lambda record: (-record.param.numel(), record.index),
459 )
461 # The greedy strategy assigns the task to the node with the lowest load.
462 coord_loads = {coord: 0 for coord in all_coords}
463 owner_by_index: Dict[int, Tuple[int, ...]] = {}
465 for record in sorted_records:
466 best_coord = min(
467 all_coords,
468 key=lambda c: (coord_loads[c], c),
469 )
471 owner_by_index[record.index] = best_coord
472 coord_loads[best_coord] += record.param.numel()
474 return owner_by_index
477def select_owned_records(
478 records: List[ParamRecord],
479 owner_by_index: Dict[int, Tuple[int, ...]],
480 replicate_group_ranks: Tuple[int, ...],
481) -> List[ParamRecord]:
482 """Select records owned by current multi-dimensional replicate rank."""
483 if any(r < 0 for r in replicate_group_ranks):
484 return []
486 return [
487 record for record in records
488 if owner_by_index.get(record.index) == replicate_group_ranks
489 ]