Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / muon_shard.py: 0%
264 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"""Shard function of muon."""
18import math
19import logging
21from dataclasses import dataclass
22from typing import Any, Dict, List, Optional, Sequence, Tuple
24import torch
25import torch.distributed as dist
27from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor
28from hyper_parallel.core.optimizer.sharding_category import (
29 HSDPCommGroup,
30 ParamLayoutSpec,
31 ParamShardMeta,
32 ParamShardStageMeta,
33)
35__all__ = [
36 "_debug_param_shard_metadata",
37 "build_param_shard_metadata_for_group",
38 "build_pad_ns_inputs",
39 "chunk_update_by_layout",
40 "fused_allgather_dtensor_params",
41]
43logger = logging.getLogger(__name__)
46def make_param_shard_meta(
47 param: torch.Tensor,
48 stage_metas: List[ParamShardStageMeta],
49) -> ParamShardMeta:
50 """Construct shard metadata from real local-shape plans."""
51 global_shape = tuple(param.shape)
52 local_shape = tuple(to_local_if_dtensor(param.data).shape)
53 return ParamShardMeta(
54 global_shape=global_shape,
55 global_numel=math.prod(global_shape),
56 local_shape=local_shape,
57 local_numel=math.prod(local_shape),
58 stage_metas=tuple(stage_metas),
59 )
62def _get_mesh_shape(device_mesh) -> Tuple[int, ...]:
63 """Get mesh shape."""
64 mesh = getattr(device_mesh, "mesh", None)
65 if mesh is not None:
66 return tuple(mesh.shape)
67 mesh_shape = getattr(device_mesh, "mesh_shape", None)
68 if mesh_shape is not None:
69 return tuple(mesh_shape)
70 return ()
73def _get_local_shape_kinds(global_shape: Tuple[int, ...], shard_meta: ParamShardMeta) -> List[Tuple[int, ...]]:
74 """Collect all distinct local shapes implied by the shard metadata."""
75 dim_size_options: Dict[int, List[int]] = {}
76 for stage_meta in shard_meta.stage_metas:
77 tensor_dim = stage_meta.tensor_dim
78 seen_sizes = []
79 for split_size in stage_meta.split_sizes:
80 if split_size not in seen_sizes:
81 seen_sizes.append(split_size)
82 dim_size_options[tensor_dim] = seen_sizes
84 shape_kinds = [list(global_shape)]
85 for tensor_dim, size_options in sorted(dim_size_options.items()):
86 next_shape_kinds = []
87 for shape_kind in shape_kinds:
88 for size in size_options:
89 next_shape = list(shape_kind)
90 next_shape[tensor_dim] = size
91 next_shape_kinds.append(next_shape)
92 shape_kinds = next_shape_kinds
94 unique_shape_kinds = []
95 seen = set()
96 for shape_kind in shape_kinds:
97 shape_tuple = tuple(shape_kind)
98 if shape_tuple in seen:
99 continue
100 seen.add(shape_tuple)
101 unique_shape_kinds.append(shape_tuple)
102 return unique_shape_kinds
105def _debug_param_shard_metadata(hsdp_group, param_to_meta: Dict[torch.nn.Parameter, ParamShardMeta]) -> None:
106 """Log shard metadata for the group on rank 0 when distributed is ready."""
107 if not dist.is_available() or not dist.is_initialized():
108 return
110 rank = dist.get_rank()
111 if rank != 0:
112 return
114 for record in hsdp_group.records:
115 param = record.param
116 shard_meta = param_to_meta.get(param)
117 if shard_meta is None:
118 continue
119 local_shape_kinds = _get_local_shape_kinds(shard_meta.global_shape, shard_meta)
120 logger.debug_rank0(
121 "name=%s, placements=%s, mesh_shape=%s, global_shape=%s, local_shapes=%s",
122 getattr(param, "model_name", f"param_{record.index}"),
123 [str(placement) for placement in param.placements],
124 _get_mesh_shape(param.device_mesh),
125 shard_meta.global_shape,
126 local_shape_kinds,
127 )
130def _allgather_shapes_for_group(
131 local_shapes: List[Tuple[int, ...]],
132 shard_pg: Optional[dist.ProcessGroup],
133 shard_size: int,
134 device: torch.device,
135) -> List[List[Tuple[int, ...]]]:
136 """All-gather per-parameter shapes within one shard process group."""
137 if shard_pg is None or shard_size <= 1:
138 return [[tuple(shape) for shape in local_shapes]]
140 shape_tensor = torch.tensor(local_shapes, dtype=torch.int64, device=device)
141 gathered = torch.empty((shard_size, *shape_tensor.shape), dtype=torch.int64, device=device)
142 dist.all_gather_into_tensor(gathered.view(-1), shape_tensor.contiguous().view(-1), group=shard_pg)
143 gathered_cpu = gathered.cpu().tolist()
144 return [
145 [tuple(shape) for shape in rank_shapes]
146 for rank_shapes in gathered_cpu
147 ]
150def build_param_shard_metadata_for_group(
151 hsdp_group: HSDPCommGroup,
152) -> Dict[torch.Tensor, ParamShardMeta]:
153 """Build per-parameter shard metadata before HSDP batching."""
154 layout_spec = hsdp_group.layout_spec
155 shard_pgs = hsdp_group.shard_pgs
156 if layout_spec is None or not layout_spec.shard_axes or not hsdp_group.records:
157 return {}
159 current_shapes = [tuple(to_local_if_dtensor(record.param.data).shape) for record in hsdp_group.records]
160 stage_metas_per_param: List[List[ParamShardStageMeta]] = [[] for _ in hsdp_group.records]
161 local_device = to_local_if_dtensor(hsdp_group.records[0].param.data).device
163 for (_, tensor_dim), shard_pg in zip(layout_spec.shard_axes, shard_pgs):
164 shard_size = dist.get_world_size(shard_pg) if shard_pg is not None else 1
165 gathered_shapes = _allgather_shapes_for_group(current_shapes, shard_pg, shard_size, local_device)
166 cur_rank = dist.get_rank(shard_pg) if shard_pg is not None and shard_size > 1 else 0
167 next_shapes: List[Tuple[int, ...]] = []
169 for record_idx, record in enumerate(hsdp_group.records):
170 rank_shapes = tuple(
171 tuple(gathered_shapes[rank_idx][record_idx])
172 for rank_idx in range(shard_size)
173 )
174 tensor_dim_norm = tensor_dim % len(rank_shapes[0])
175 split_sizes = tuple(shape[tensor_dim_norm] for shape in rank_shapes)
176 pad_shape = tuple(max(shape[dim] for shape in rank_shapes) for dim in range(len(rank_shapes[0])))
177 stage_metas_per_param[record_idx].append(
178 ParamShardStageMeta(
179 tensor_dim=tensor_dim_norm,
180 shard_size=shard_size,
181 cur_rank_in_shard_group=cur_rank,
182 rank_shapes=rank_shapes,
183 split_sizes=split_sizes,
184 pad_shape=pad_shape,
185 )
186 )
188 output_shape = list(rank_shapes[cur_rank])
189 if shard_size > 1:
190 output_shape[tensor_dim_norm] = sum(split_sizes)
191 next_shapes.append(tuple(output_shape))
193 current_shapes = next_shapes
195 param_to_meta: Dict[torch.Tensor, ParamShardMeta] = {}
196 updated_records = []
197 for record, stage_metas in zip(hsdp_group.records, stage_metas_per_param):
198 shard_meta = make_param_shard_meta(record.param, stage_metas)
199 param_to_meta[record.param] = shard_meta
200 updated_records.append(record.__class__(index=record.index, param=record.param, shard_meta=shard_meta))
202 hsdp_group.records = updated_records
203 return param_to_meta
206def build_pad_ns_inputs(
207 comm_params: List[torch.Tensor],
208 param_to_ns_input: Dict[torch.Tensor, torch.Tensor],
209 param_shard_metadata: Optional[Dict[torch.Tensor, ParamShardMeta]] = None,
210) -> List[torch.Tensor]:
211 """Build aligned local NS inputs for all communication params."""
212 if not comm_params:
213 return []
215 ref_tensor = next(iter(param_to_ns_input.values()), None)
216 default_dtype = ref_tensor.dtype if ref_tensor is not None else to_local_if_dtensor(comm_params[0].data).dtype
217 local_inputs: List[torch.Tensor] = []
219 for param in comm_params:
220 ns_input = param_to_ns_input.get(param)
221 if ns_input is not None:
222 local_inputs.append(ns_input)
223 continue
225 shard_meta = None
226 if param_shard_metadata is not None:
227 shard_meta = param_shard_metadata.get(param)
228 if shard_meta is None and hasattr(param, "shard_meta"):
229 shard_meta = getattr(param, "shard_meta")
230 if shard_meta is None:
231 local_shape = tuple(to_local_if_dtensor(param.data).shape)
232 else:
233 local_shape = shard_meta.local_shape
235 local_inputs.append(
236 torch.zeros(
237 local_shape,
238 dtype=default_dtype,
239 device=to_local_if_dtensor(param.data).device,
240 )
241 )
243 return local_inputs
246def chunk_update_by_layout(
247 global_update: torch.Tensor,
248 param: torch.Tensor,
249 layout_spec: ParamLayoutSpec,
250 param_shard_meta: Optional[ParamShardMeta] = None,
251) -> torch.Tensor:
252 """Slice a full update back to the local shard using narrow."""
253 if not hasattr(param, "device_mesh") or layout_spec is None or not layout_spec.shard_axes:
254 return global_update
256 device_mesh = param.device_mesh
257 mesh_coordinates = device_mesh.get_coordinate()
259 shard_axes = layout_spec.shard_axes
260 local_update = global_update
262 # ``stage_metas`` are built while all-gathering local shards back to the
263 # global tensor shape, so slicing a global update back to the local shard
264 # must reverse that order.
265 for axis_idx, (mesh_dim, tensor_dim) in reversed(list(enumerate(shard_axes))):
266 num_chunks = device_mesh.size(mesh_dim)
268 if num_chunks <= 1:
269 continue
271 local_rank = mesh_coordinates[mesh_dim]
272 if param_shard_meta is not None and axis_idx < len(param_shard_meta.stage_metas):
273 split_sizes = param_shard_meta.stage_metas[axis_idx].split_sizes
274 else:
275 full_chunk_size = (local_update.size(tensor_dim) + num_chunks - 1) // num_chunks
276 split_sizes = tuple(
277 min(full_chunk_size, max(local_update.size(tensor_dim) - rank * full_chunk_size, 0))
278 for rank in range(num_chunks)
279 )
281 start = sum(split_sizes[:local_rank])
282 chunk_size = split_sizes[local_rank]
283 local_update = local_update.narrow(tensor_dim, start, chunk_size)
285 if not local_update.is_contiguous():
286 local_update = local_update.contiguous()
288 return local_update
291def _get_or_alloc_buffer(
292 cache: Optional[Dict],
293 key: Any,
294 numel: int,
295 dtype: torch.dtype,
296 device: torch.device,
297) -> torch.Tensor:
298 """Return a cached buffer with at least `numel` elements."""
299 if cache is not None and key in cache:
300 buf = cache[key]
301 if buf.numel() >= numel:
302 return buf
303 buf = torch.empty(numel, dtype=dtype, device=device)
304 cache[key] = buf
305 return buf
307 buf = torch.empty(numel, dtype=dtype, device=device)
308 if cache is not None:
309 cache[key] = buf
310 return buf
313def _early_return_tensors(
314 local_tensors: List[torch.Tensor],
315 keep_indices: Optional[set],
316) -> List[Optional[torch.Tensor]]:
317 """Return tensors directly when no communication is needed."""
318 if keep_indices is None:
319 return list(local_tensors)
320 return [t if i in keep_indices else None for i, t in enumerate(local_tensors)]
323@dataclass(frozen=True)
324class _GatherParamMeta:
325 offset: int
326 actual_numel: int
327 padded_numel: int
328 rest_shape: Tuple[int, ...]
329 split_sizes: Tuple[int, ...]
332def _prepare_gather_inputs(
333 current_tensors: List[torch.Tensor],
334 tensor_dim: int,
335 alignment_elements: int,
336 shard_size: int,
337 stage_metas: Optional[List[Optional[ParamShardStageMeta]]] = None,
338) -> Tuple[List[torch.Tensor], List[_GatherParamMeta], int]:
339 """Move shard dim to dim0 and compute padding metadata."""
340 gather_inputs: List[torch.Tensor] = []
341 param_meta: List[_GatherParamMeta] = []
342 total_padded_numel = 0
344 for idx, t in enumerate(current_tensors):
345 tensor_dim_norm = tensor_dim % t.dim()
347 if tensor_dim_norm == 0 and t.is_contiguous():
348 gi = t
349 else:
350 gi = t.movedim(tensor_dim_norm, 0).contiguous()
352 stage_meta = stage_metas[idx] if stage_metas is not None else None
353 actual_numel = gi.numel()
354 if stage_meta is not None:
355 split_sizes = stage_meta.split_sizes
356 pad_shape_moved = list(stage_meta.pad_shape)
357 if tensor_dim_norm != 0:
358 pad_shape_moved[0], pad_shape_moved[tensor_dim_norm] = pad_shape_moved[tensor_dim_norm], \
359 pad_shape_moved[0]
360 padded_numel_raw = math.prod(pad_shape_moved)
361 else:
362 split_sizes = tuple(gi.shape[0] for _ in range(shard_size))
363 padded_numel_raw = gi.numel()
364 padded_numel = ((padded_numel_raw + alignment_elements - 1) // alignment_elements) * alignment_elements
366 gather_inputs.append(gi)
367 param_meta.append(
368 _GatherParamMeta(
369 offset=total_padded_numel,
370 actual_numel=actual_numel,
371 padded_numel=padded_numel,
372 rest_shape=tuple(gi.shape[1:]),
373 split_sizes=split_sizes,
374 )
375 )
376 total_padded_numel += padded_numel
378 return gather_inputs, param_meta, total_padded_numel
381def _pack_and_allgather(
382 gather_inputs: List[torch.Tensor],
383 param_meta: List[_GatherParamMeta],
384 total_padded_numel: int,
385 axis_idx: int,
386 dtype: torch.dtype,
387 device: torch.device,
388 shard_pg: dist.ProcessGroup,
389 shard_size: int,
390 buffer_cache: Optional[Dict],
391) -> torch.Tensor:
392 """Pack local shards into one buffer, all-gather, return gathered view."""
393 cache_key = ("fused_allgather", axis_idx, dtype, device)
394 pack_buffer = _get_or_alloc_buffer(
395 buffer_cache, cache_key, total_padded_numel,
396 dtype, device,
397 )[:total_padded_numel]
399 pack_buffer.zero_()
400 for gi, meta in zip(gather_inputs, param_meta):
401 pack_buffer[meta.offset:meta.offset + meta.actual_numel].copy_(gi.view(-1))
403 gathered_numel = total_padded_numel * shard_size
404 cache_key_out = ("fused_allgather_out", axis_idx, dtype, device)
405 gathered_buffer = _get_or_alloc_buffer(
406 buffer_cache, cache_key_out, gathered_numel,
407 dtype, device,
408 )[:gathered_numel]
410 dist.all_gather_into_tensor(gathered_buffer, pack_buffer, group=shard_pg)
411 return gathered_buffer.view(shard_size, total_padded_numel)
414def _unpack_gathered_results(
415 gathered_view: torch.Tensor,
416 gather_inputs: List[torch.Tensor],
417 param_meta: List[_GatherParamMeta],
418 current_tensors: List[torch.Tensor],
419 tensor_dim: int,
420 shard_size: int,
421 n_params: int,
422 is_last_axis: bool,
423 keep_indices: Optional[set],
424) -> List[Optional[torch.Tensor]]:
425 """Slice gathered buffer back to per-parameter full tensors."""
426 new_tensors: List[Optional[torch.Tensor]] = []
427 for i in range(n_params):
428 if is_last_axis and keep_indices is not None and i not in keep_indices:
429 new_tensors.append(None)
430 continue
432 meta = param_meta[i]
433 rest_shape = meta.rest_shape
434 rest_numel = math.prod(rest_shape) if rest_shape else 1
435 param_slice = gathered_view[:, meta.offset:meta.offset + meta.padded_numel]
437 param_chunks = []
438 for rank in range(shard_size):
439 valid_numel = meta.split_sizes[rank] * rest_numel
440 if valid_numel == 0:
441 continue
442 param_chunks.append(param_slice[rank, :valid_numel].contiguous().view(-1))
444 if param_chunks:
445 param_data = torch.cat(param_chunks, dim=0)
446 else:
447 param_data = gather_inputs[i].new_empty((0,))
449 result = param_data.view(sum(meta.split_sizes), *rest_shape)
450 tensor_dim_norm = tensor_dim % current_tensors[i].dim()
451 if tensor_dim_norm == 0:
452 new_tensors.append(result)
453 else:
454 new_tensors.append(result.movedim(0, tensor_dim_norm))
456 return new_tensors
459def fused_allgather_dtensor_params(
460 local_tensors: List[torch.Tensor],
461 shard_pgs: Sequence[dist.ProcessGroup],
462 layout_spec: ParamLayoutSpec,
463 param_shard_metadata: Optional[List[Optional[ParamShardMeta]]] = None,
464 buffer_cache: Optional[Dict] = None,
465 keep_indices: Optional[set] = None,
466) -> List[Optional[torch.Tensor]]:
467 """Fuse many parameter shards into one all-gather per shard axis."""
468 if not shard_pgs or not local_tensors:
469 return _early_return_tensors(local_tensors, keep_indices)
471 n_params = len(local_tensors)
472 device = local_tensors[0].device
473 dtype = local_tensors[0].dtype
474 alignment_bytes = 512
475 element_size = local_tensors[0].element_size()
476 alignment_elements = max(1, alignment_bytes // element_size)
478 active_axes = []
479 for axis_idx, ((_, tensor_dim), shard_pg) in enumerate(zip(layout_spec.shard_axes, shard_pgs)):
480 if shard_pg is None:
481 continue
482 shard_size = dist.get_world_size(shard_pg)
483 if shard_size <= 1:
484 continue
485 active_axes.append((axis_idx, tensor_dim, shard_pg, shard_size))
487 if not active_axes:
488 return _early_return_tensors(local_tensors, keep_indices)
490 current_tensors: List[torch.Tensor] = list(local_tensors)
492 for active_pos, (axis_idx, tensor_dim, shard_pg, shard_size) in enumerate(active_axes):
493 is_last_axis = active_pos == len(active_axes) - 1
494 stage_metas: Optional[List[Optional[ParamShardStageMeta]]] = None
495 if param_shard_metadata is not None:
496 stage_metas = []
497 for meta in param_shard_metadata:
498 if meta is None or axis_idx >= len(meta.stage_metas):
499 stage_metas.append(None)
500 else:
501 stage_metas.append(meta.stage_metas[axis_idx])
503 gather_inputs, param_meta, total_padded_numel = _prepare_gather_inputs(
504 current_tensors, tensor_dim, alignment_elements, shard_size, stage_metas,
505 )
507 gathered_view = _pack_and_allgather(
508 gather_inputs, param_meta, total_padded_numel,
509 axis_idx, dtype, device, shard_pg, shard_size, buffer_cache,
510 )
512 new_tensors = _unpack_gathered_results(
513 gathered_view, gather_inputs, param_meta,
514 current_tensors, tensor_dim, shard_size,
515 n_params, is_last_axis, keep_indices,
516 )
518 if is_last_axis:
519 return new_tensors
521 current_tensors = new_tensors # type: ignore[assignment]
523 return current_tensors