Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / sharding_category.py: 0%
282 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# ============================================================================
16"""Category parameter with dtensor."""
18import itertools
19from dataclasses import dataclass, field
20from typing import Dict, List, Sequence, Tuple, Optional, Any
22import torch
23import torch.distributed as dist
25from hyper_parallel.core.optimizer.dtensor_compat import (
26 DTensor,
27 DeviceMesh,
28 Shard,
29 StridedShard,
30)
32@dataclass(frozen=True)
33class ParamLayoutSpec:
34 """Shape-free but ndim-aware parameter layout.
36 Same layout means:
37 1. same tensor ndim;
38 2. same shard mesh dims;
39 3. same shard tensor dims;
40 4. same replicate mesh dims.
42 Example:
43 tensor_ndim = 3
44 placements = (Shard(0), Replicate(), Replicate(), Replicate())
46 Then:
47 shard_axes = ((0, 0),)
48 replicate_mesh_dims = (1, 2, 3)
50 Meaning:
51 mesh dim 0 shards tensor dim 0;
52 mesh dim 1/2/3 are replicated.
53 """
55 tensor_ndim: int
57 # Each item is:
58 # (mesh_dim, tensor_dim)
59 shard_axes: Tuple[Tuple[int, int], ...]
61 # Mesh dims that are replicated.
62 replicate_mesh_dims: Tuple[int, ...]
64 @property
65 def shard_mesh_dims(self) -> Tuple[int, ...]:
66 """Return the mesh dimensions used for sharding."""
67 return tuple(mesh_dim for mesh_dim, _ in self.shard_axes)
69 @property
70 def shard_tensor_dims(self) -> Tuple[int, ...]:
71 """Return the tensor dimensions that are sharded."""
72 return tuple(tensor_dim for _, tensor_dim in self.shard_axes)
74 @property
75 def is_last2d_sharded(self) -> bool:
76 """Whether any shard axis falls on the last 2 tensor dimensions.
78 Newton-Schulz iteration operates on the last 2 dims.
79 If either dim is sharded, allgather is needed before NS.
80 """
81 for _, tensor_dim in self.shard_axes:
82 if tensor_dim >= self.tensor_ndim - 2:
83 return True
84 return False
87@dataclass(frozen=True)
88class CommDomainKey:
89 """Communication-domain key based on mesh dimensions.
91 This key describes the communication domain logically by relying on the
92 underlying DeviceMesh identifiers and the specific dims involved.
93 """
95 mesh_shape: Tuple[int, ...]
96 mesh_rank_list: Tuple[int, ...]
98 replicate_mesh_dims: Tuple[int, ...]
99 shard_mesh_dims: Tuple[int, ...]
101 @property
102 def has_replicate_redundancy(self) -> bool:
103 """Check if there is any replicate redundancy in the mesh."""
104 return len(self.replicate_mesh_dims) > 0
106 @property
107 def has_shard_group(self) -> bool:
108 """Check if there are any sharded mesh dimensions."""
109 return len(self.shard_mesh_dims) > 0
112@dataclass(frozen=True)
113class HSDPGroupKey:
114 """Final grouping key for HSDP.
116 Same key means:
117 1. same communication domain;
118 2. same tensor ndim;
119 3. same shard / replicate axis layout.
121 Shape is not part of this key.
122 """
124 comm_key: CommDomainKey
125 axis_spec: ParamLayoutSpec
128@dataclass
129class ParamShardSpec:
130 """Per-parameter shard metadata used during HSDP grouping."""
132 device_mesh: DeviceMesh
133 shard_mesh_dims: Tuple[int, ...]
134 replicate_mesh_dims: Tuple[int, ...]
135 replicate_pgs: Tuple[dist.ProcessGroup, ...]
136 shard_pgs: Tuple[dist.ProcessGroup, ...]
137 axis_spec: ParamLayoutSpec
140@dataclass(frozen=True)
141class ParamRecord:
142 """A lightweight parameter record."""
144 index: int
145 param: DTensor
148@dataclass
149class HSDPCommGroup:
150 """HSDP communication group.
152 Same group means same domain, ndim, and layout.
153 Stores sequences of ProcessGroups aligned with mesh dims.
154 """
156 comm_key: CommDomainKey
157 layout_spec: ParamLayoutSpec
159 # Tuple of runtime process groups, one for each relevant mesh dimension
160 replicate_pgs: Tuple[dist.ProcessGroup, ...] = ()
161 shard_pgs: Tuple[dist.ProcessGroup, ...] = ()
163 records: List[ParamRecord] = field(default_factory=list)
165 @property
166 def params(self) -> List[DTensor]:
167 """Return all DTensors within this communication group."""
168 return [record.param for record in self.records]
170 def add_param(self, index: int, param: DTensor) -> None:
171 """Add a parameter record to the group."""
172 self.records.append(
173 ParamRecord(
174 index=index,
175 param=param,
176 )
177 )
179 def __len__(self) -> int:
180 return len(self.records)
183def _normalize_tensor_dim(dim: int, tensor_ndim: int) -> int:
184 """Normalize tensor dim to non-negative dim."""
185 original_dim = dim
187 if dim < 0:
188 dim += tensor_ndim
190 if dim < 0 or dim >= tensor_ndim:
191 raise ValueError(
192 f"Invalid shard dim {original_dim} for tensor ndim {tensor_ndim}."
193 )
195 return dim
198def extract_param_shard_spec(dtensor: DTensor) -> ParamShardSpec:
199 """Extract all shard metadata and native process groups from one DTensor."""
200 device_mesh = dtensor.device_mesh
201 placements = dtensor.placements
202 tensor_ndim = len(dtensor.shape)
204 shard_mesh_dims: List[int] = []
205 replicate_mesh_dims: List[int] = []
206 shard_axes: List[Tuple[int, int]] = []
207 replicate_pgs: List[dist.ProcessGroup] = []
208 shard_pgs: List[dist.ProcessGroup] = []
210 for mesh_dim_idx, placement in enumerate(placements):
211 pg = device_mesh.get_group(mesh_dim_idx) if hasattr(device_mesh, "get_group") else None
213 if placement.is_replicate():
214 replicate_mesh_dims.append(mesh_dim_idx)
215 replicate_pgs.append(pg)
217 elif placement.is_shard():
218 if isinstance(placement, StridedShard):
219 placement_for_grouping = Shard(placement.dim)
220 else:
221 placement_for_grouping = placement
223 tensor_dim = _normalize_tensor_dim(
224 placement_for_grouping.dim,
225 tensor_ndim,
226 )
228 shard_mesh_dims.append(mesh_dim_idx)
229 shard_axes.append((mesh_dim_idx, tensor_dim))
230 shard_pgs.append(pg)
232 else:
233 raise ValueError(
234 f"Unsupported placement type in HSDP parameter grouping: "
235 f"{type(placement).__name__}."
236 )
238 axis_spec = ParamLayoutSpec(
239 tensor_ndim=tensor_ndim,
240 shard_axes=tuple(shard_axes),
241 replicate_mesh_dims=tuple(replicate_mesh_dims),
242 )
244 return ParamShardSpec(
245 device_mesh=device_mesh,
246 shard_mesh_dims=tuple(shard_mesh_dims),
247 replicate_mesh_dims=tuple(replicate_mesh_dims),
248 replicate_pgs=tuple(replicate_pgs),
249 shard_pgs=tuple(shard_pgs),
250 axis_spec=axis_spec,
251 )
254def build_comm_domain_key(shard_spec: ParamShardSpec) -> CommDomainKey:
255 """Build grouping key utilizing DeviceMesh properties."""
256 mesh_rank_list = ()
257 if hasattr(shard_spec.device_mesh, "rank_list"):
258 mesh_rank_list = tuple(shard_spec.device_mesh.rank_list)
260 return CommDomainKey(
261 mesh_shape = tuple(getattr(shard_spec.device_mesh, "mesh_shape", None) or shard_spec.device_mesh.mesh.shape),
262 mesh_rank_list=mesh_rank_list,
263 replicate_mesh_dims=shard_spec.replicate_mesh_dims,
264 shard_mesh_dims=shard_spec.shard_mesh_dims,
265 )
268def group_parameters_for_hsdp(
269 params: List[DTensor],
270) -> Tuple[List[DTensor], List[HSDPCommGroup]]:
271 """Group parameters relying on native DeviceMesh topology."""
272 no_comm_params: List[DTensor] = []
273 groups: Dict[HSDPGroupKey, HSDPCommGroup] = {}
275 for param_index, param in enumerate(params):
276 if not isinstance(param, DTensor):
277 no_comm_params.append(param)
278 continue
280 shard_spec = extract_param_shard_spec(param)
281 comm_key = build_comm_domain_key(shard_spec)
283 if not comm_key.has_replicate_redundancy and not comm_key.has_shard_group:
284 no_comm_params.append(param)
285 continue
287 group_key = HSDPGroupKey(
288 comm_key=comm_key,
289 axis_spec=shard_spec.axis_spec, # ParamLayoutSpec
290 )
292 if group_key not in groups:
293 groups[group_key] = HSDPCommGroup(
294 comm_key=comm_key,
295 layout_spec=shard_spec.axis_spec,
296 replicate_pgs=shard_spec.replicate_pgs,
297 shard_pgs=shard_spec.shard_pgs,
298 )
300 groups[group_key].add_param(param_index, param)
302 return no_comm_params, list(groups.values())
305@dataclass
306class HSDPGroupAssignment:
307 """Optimizer assignment for one HSDP communication group."""
309 owned_records: List[ParamRecord]
310 all_records: List[ParamRecord]
312 # param_index -> (dim_0_rank, dim_1_rank, ...)
313 owner_by_index: Dict[int, Tuple[int, ...]]
315 # Record the rank and size for each dimension, and the list of cur_rank within replicate_groups.
316 replicate_group_ranks: Tuple[int, ...]
317 replicate_sizes: Tuple[int, ...]
319 replicate_pgs: Tuple[dist.ProcessGroup, ...] = ()
320 shard_pgs: Tuple[dist.ProcessGroup, ...] = ()
322 is_shard: bool = False
323 layout_spec: Optional[ParamLayoutSpec] = None
325 @property
326 def owned_params(self) -> List[DTensor]:
327 """Return the parameters owned by the current rank."""
328 return [record.param for record in self.owned_records]
330 @property
331 def all_params(self) -> List[DTensor]:
332 """Return all parameters in this assignment group."""
333 return [record.param for record in self.all_records]
335 @property
336 def is_replicated(self) -> bool:
337 """Check if the group spans across multiple ranks."""
338 return any(s > 1 for s in self.replicate_sizes)
340 def owner_rank_coord(self, record: ParamRecord) -> Tuple[int, ...]:
341 """Get the rank coordinates of the owner of a given record."""
342 return self.owner_by_index.get(record.index, ())
344 def is_owned(self, record: ParamRecord) -> bool:
345 """Check if a given record is owned by the current replicate group rank."""
346 return self.owner_rank_coord(record) == self.replicate_group_ranks
348 def __str__(self) -> str:
349 shard_ranks = [list(dist.get_process_group_ranks(pg)) for pg in self.shard_pgs if pg is not None]
350 replicate_ranks = [list(dist.get_process_group_ranks(pg)) for pg in self.replicate_pgs if pg is not None]
352 owned_names = [getattr(p, "model_name", f"p_{i}") for i, p in enumerate(self.owned_params)]
353 all_names = [getattr(p, "model_name", f"p_{i}") for i, p in enumerate(self.all_params)]
355 return (
356 f"HSDPGroupAssignment( \n"
357 f"owned_params={owned_names}, \n"
358 f"all_params={all_names}, \n"
359 f"replicate_group_ranks={self.replicate_group_ranks}, \n"
360 f"replicate_sizes={self.replicate_sizes}, \n"
361 f"is_shard={self.is_shard}, \n"
362 f"shard_pg_ranks={shard_ranks}, \n"
363 f"replicate_pg_ranks={replicate_ranks}, \n"
364 f"layout_spec={self.layout_spec}) \n"
365 )
367 __repr__ = __str__
370def get_multi_dim_logical_info(
371 device_mesh: DeviceMesh,
372 mesh_dims: Sequence[int]
373) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
374 """Obtain the independent relative ranks and sizes of the parameters across multiple replicate dimensions."""
375 if not mesh_dims:
376 return (), ()
378 coords = device_mesh.get_coordinate()
379 if coords is None:
380 return (-1,) * len(mesh_dims), (1,) * len(mesh_dims)
382 ranks = tuple(coords[dim] for dim in mesh_dims)
383 sizes = tuple(device_mesh.size(dim) for dim in mesh_dims)
385 return ranks, sizes
388def build_owner_by_size(
389 records: List[ParamRecord],
390 replicate_sizes: Tuple[int, ...],
391) -> Dict[int, Tuple[int, ...]]:
392 """Build deterministic owner map across a multi-dimensional replicate grid."""
393 valid_records = [
394 record for record in records
395 if getattr(record.param, "requires_grad", True)
396 ]
398 if not valid_records:
399 return {}
401 if not replicate_sizes:
402 return {record.index: () for record in valid_records}
404 dim_ranges = [range(s) for s in replicate_sizes]
405 all_coords = list(itertools.product(*dim_ranges))
407 sorted_records = sorted(
408 valid_records,
409 key=lambda record: (-record.param.numel(), record.index),
410 )
412 # The greedy strategy assigns the task to the node with the lowest load.
413 coord_loads = {coord: 0 for coord in all_coords}
414 owner_by_index: Dict[int, Tuple[int, ...]] = {}
416 for record in sorted_records:
417 best_coord = min(
418 all_coords,
419 key=lambda c: (coord_loads[c], c),
420 )
422 owner_by_index[record.index] = best_coord
423 coord_loads[best_coord] += record.param.numel()
425 return owner_by_index
428def select_owned_records(
429 records: List[ParamRecord],
430 owner_by_index: Dict[int, Tuple[int, ...]],
431 replicate_group_ranks: Tuple[int, ...],
432) -> List[ParamRecord]:
433 """Select records owned by current multi-dimensional replicate rank."""
434 if any(r < 0 for r in replicate_group_ranks):
435 return []
437 return [
438 record for record in records
439 if owner_by_index.get(record.index) == replicate_group_ranks
440 ]
443def chunk_update_by_layout(
444 global_update: torch.Tensor,
445 param: "DTensor",
446 layout_spec: "ParamLayoutSpec",
447) -> torch.Tensor:
448 """Slice a full update back to the local shard using narrow."""
449 if not hasattr(param, "device_mesh") or layout_spec is None or not layout_spec.shard_axes:
450 return global_update
452 device_mesh = param.device_mesh
453 mesh_coordinates = device_mesh.get_coordinate()
455 shard_axes = layout_spec.shard_axes
456 local_update = global_update
458 # Apply each shard axis in order. `narrow` returns a view and avoids
459 # creating all chunks when only the local rank's chunk is needed.
460 for mesh_dim, tensor_dim in shard_axes:
461 num_chunks = device_mesh.size(mesh_dim)
463 if num_chunks <= 1:
464 continue
466 local_rank = mesh_coordinates[mesh_dim]
467 chunk_size = local_update.size(tensor_dim) // num_chunks
469 local_update = local_update.narrow(tensor_dim, local_rank * chunk_size, chunk_size)
471 # Non-dim0 slicing may produce a non-contiguous view.
472 if not local_update.is_contiguous():
473 local_update = local_update.contiguous()
475 return local_update
478def _get_or_alloc_buffer(
479 cache: Optional[Dict],
480 key: Any,
481 numel: int,
482 dtype: torch.dtype,
483 device: torch.device,
484) -> torch.Tensor:
485 """Return a cached buffer with at least `numel` elements."""
486 if cache is not None and key in cache:
487 buf = cache[key]
488 if buf.numel() >= numel:
489 return buf
490 buf = torch.empty(numel, dtype=dtype, device=device)
491 cache[key] = buf
492 return buf
494 buf = torch.empty(numel, dtype=dtype, device=device)
495 if cache is not None:
496 cache[key] = buf
497 return buf
500def _early_return_tensors(
501 local_tensors: List[torch.Tensor],
502 keep_indices: Optional[set],
503) -> List[Optional[torch.Tensor]]:
504 """Return tensors directly when no communication is needed."""
505 if keep_indices is None:
506 return list(local_tensors)
507 return [t if i in keep_indices else None for i, t in enumerate(local_tensors)]
510def _prepare_gather_inputs(
511 current_tensors: List[torch.Tensor],
512 tensor_dim: int,
513 alignment_elements: int,
514) -> Tuple[List[torch.Tensor], List[Tuple[int, int, int, Tuple[int, ...]]], int]:
515 """Move shard dim to dim0 and compute padding metadata.
517 Returns:
518 (gather_inputs, param_meta, total_padded_numel)
519 param_meta item: (offset, actual_numel, padded_numel, rest_shape)
520 """
521 gather_inputs: List[torch.Tensor] = []
522 param_meta: List[Tuple[int, int, int, Tuple[int, ...]]] = []
523 total_padded_numel = 0
525 for t in current_tensors:
526 tensor_dim_norm = tensor_dim % t.dim()
528 # Put shard dim at dim0 so all-gather can concatenate along dim0.
529 if tensor_dim_norm == 0 and t.is_contiguous():
530 gi = t
531 else:
532 gi = t.movedim(tensor_dim_norm, 0).contiguous()
534 actual_numel = gi.numel()
535 padded_numel = ((actual_numel + alignment_elements - 1) // alignment_elements) * alignment_elements
537 gather_inputs.append(gi)
538 param_meta.append((total_padded_numel, actual_numel, padded_numel, tuple(gi.shape[1:])))
539 total_padded_numel += padded_numel
541 return gather_inputs, param_meta, total_padded_numel
544def _pack_and_allgather(
545 gather_inputs: List[torch.Tensor],
546 param_meta: List[Tuple[int, int, int, Tuple[int, ...]]],
547 total_padded_numel: int,
548 axis_idx: int,
549 dtype: torch.dtype,
550 device: torch.device,
551 shard_pg: dist.ProcessGroup,
552 shard_size: int,
553 buffer_cache: Optional[Dict],
554) -> torch.Tensor:
555 """Pack local shards into one buffer, all-gather, return gathered view.
557 Returns:
558 gathered_view with shape [shard_size, total_padded_numel]
559 """
560 cache_key = ("fused_allgather", axis_idx, dtype, device)
561 pack_buffer = _get_or_alloc_buffer(
562 buffer_cache, cache_key, total_padded_numel,
563 dtype, device,
564 )[:total_padded_numel]
566 # Copy only real data. Padding is not zeroed because it is never read.
567 for gi, (offset, actual_numel, _, _) in zip(gather_inputs, param_meta):
568 pack_buffer[offset:offset + actual_numel].copy_(gi.view(-1))
570 gathered_numel = total_padded_numel * shard_size
571 cache_key_out = ("fused_allgather_out", axis_idx, dtype, device)
572 gathered_buffer = _get_or_alloc_buffer(
573 buffer_cache, cache_key_out, gathered_numel,
574 dtype, device,
575 )[:gathered_numel]
577 dist.all_gather_into_tensor(gathered_buffer, pack_buffer, group=shard_pg)
579 # Rank-major layout: [rank0_pack | rank1_pack | ...]
580 return gathered_buffer.view(shard_size, total_padded_numel)
583def _unpack_gathered_results(
584 gathered_view: torch.Tensor,
585 gather_inputs: List[torch.Tensor],
586 param_meta: List[Tuple[int, int, int, Tuple[int, ...]]],
587 current_tensors: List[torch.Tensor],
588 tensor_dim: int,
589 shard_size: int,
590 n_params: int,
591 is_last_axis: bool,
592 keep_indices: Optional[set],
593) -> List[Optional[torch.Tensor]]:
594 """Slice gathered buffer back to per-parameter full tensors."""
595 new_tensors: List[Optional[torch.Tensor]] = []
596 for i in range(n_params):
597 # Only the final output can be skipped.
598 # Earlier axis results may be needed by the next shard-axis gather.
599 if is_last_axis and keep_indices is not None and i not in keep_indices:
600 new_tensors.append(None)
601 continue
603 offset, actual_numel, _, rest_shape = param_meta[i]
604 dim0_size = gather_inputs[i].shape[0]
606 # Pick this parameter from every rank.
607 # This slice is usually non-contiguous because data is rank-major.
608 param_slice = gathered_view[:, offset:offset + actual_numel]
610 # Materialize as contiguous: [rank0_param | rank1_param | ...]
611 param_data = param_slice.contiguous()
613 # Restore full tensor with gathered dim0.
614 result = param_data.view(dim0_size * shard_size, *rest_shape)
616 # Move dim0 back to the original shard dimension.
617 tensor_dim_norm = tensor_dim % current_tensors[i].dim()
618 if tensor_dim_norm == 0:
619 new_tensors.append(result)
620 else:
621 new_tensors.append(result.movedim(0, tensor_dim_norm))
623 return new_tensors
626def fused_allgather_dtensor_params(
627 local_tensors: List[torch.Tensor],
628 shard_pgs: Sequence[dist.ProcessGroup],
629 layout_spec: ParamLayoutSpec,
630 buffer_cache: Optional[Dict] = None,
631 keep_indices: Optional[set] = None,
632) -> List[Optional[torch.Tensor]]:
633 """Fuse many parameter shards into one all-gather per shard axis.
635 Flow:
636 1. Move shard dim to dim0 if needed.
637 2. Pack all local shards into one flat buffer.
638 3. Run one all-gather on the packed buffer.
639 4. Slice gathered buffer back to per-parameter full tensors.
640 5. On the last shard axis, only unpack `keep_indices`.
641 """
642 if not shard_pgs or not local_tensors:
643 return _early_return_tensors(local_tensors, keep_indices)
645 n_params = len(local_tensors)
646 device = local_tensors[0].device
647 dtype = local_tensors[0].dtype
648 alignment_bytes = 512
649 element_size = local_tensors[0].element_size()
650 alignment_elements = max(1, alignment_bytes // element_size)
652 # Keep only real shard axes that need communication.
653 active_axes = []
654 for axis_idx, ((_, tensor_dim), shard_pg) in enumerate(zip(layout_spec.shard_axes, shard_pgs)):
655 if shard_pg is None:
656 continue
657 shard_size = dist.get_world_size(shard_pg)
658 if shard_size <= 1:
659 continue
660 active_axes.append((axis_idx, tensor_dim, shard_pg, shard_size))
662 if not active_axes:
663 return _early_return_tensors(local_tensors, keep_indices)
665 # After each shard axis, this becomes the partially gathered result.
666 current_tensors: List[torch.Tensor] = list(local_tensors)
668 for active_pos, (axis_idx, tensor_dim, shard_pg, shard_size) in enumerate(active_axes):
669 is_last_axis = active_pos == len(active_axes) - 1
671 gather_inputs, param_meta, total_padded_numel = _prepare_gather_inputs(
672 current_tensors, tensor_dim, alignment_elements,
673 )
675 gathered_view = _pack_and_allgather(
676 gather_inputs, param_meta, total_padded_numel,
677 axis_idx, dtype, device, shard_pg, shard_size, buffer_cache,
678 )
680 new_tensors = _unpack_gathered_results(
681 gathered_view, gather_inputs, param_meta,
682 current_tensors, tensor_dim, shard_size,
683 n_params, is_last_axis, keep_indices,
684 )
686 if is_last_axis:
687 return new_tensors
689 # Safe because keep_indices is only applied on the last axis.
690 current_tensors = new_tensors # type: ignore[assignment]
692 return current_tensors