Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / muon_shard.py: 12%
256 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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."""
18__all__ = [
19 "build_param_shard_metadata_for_group",
20 "build_pad_ns_inputs",
21 "chunk_update_by_layout",
22 "fused_allgather_dtensor_params",
23]
25import math
26import logging
28from dataclasses import dataclass
29from typing import Any, Dict, List, Optional, Sequence, Tuple
31import torch
32import torch.distributed as dist
34from hyper_parallel.core.optimizer.dtensor_compat import to_local_if_dtensor
35from hyper_parallel.core.optimizer.sharding_category import (
36 HSDPCommGroup,
37 ParamLayoutSpec,
38 ParamShardMeta,
39 ParamShardStageMeta,
40)
42logger = logging.getLogger(__name__)
45def make_param_shard_meta(
46 param: torch.Tensor,
47 stage_metas: List[ParamShardStageMeta],
48) -> ParamShardMeta:
49 """Construct shard metadata from real local-shape plans."""
50 global_shape = tuple(param.shape)
51 local_shape = tuple(to_local_if_dtensor(param.data).shape)
52 return ParamShardMeta(
53 global_shape=global_shape,
54 global_numel=math.prod(global_shape),
55 local_shape=local_shape,
56 local_numel=math.prod(local_shape),
57 stage_metas=tuple(stage_metas),
58 )
61def _get_mesh_shape(device_mesh) -> Tuple[int, ...]:
62 """Get mesh shape."""
63 mesh = getattr(device_mesh, "mesh", None)
64 if mesh is not None:
65 return tuple(mesh.shape)
66 mesh_shape = getattr(device_mesh, "mesh_shape", None)
67 if mesh_shape is not None:
68 return tuple(mesh_shape)
69 return ()
72def _get_local_shape_kinds(global_shape: Tuple[int, ...], shard_meta: ParamShardMeta) -> List[Tuple[int, ...]]:
73 """Collect all distinct local shapes implied by the shard metadata."""
74 dim_size_options: Dict[int, List[int]] = {}
75 for stage_meta in shard_meta.stage_metas:
76 tensor_dim = stage_meta.tensor_dim
77 seen_sizes = []
78 for split_size in stage_meta.split_sizes:
79 if split_size not in seen_sizes:
80 seen_sizes.append(split_size)
81 dim_size_options[tensor_dim] = seen_sizes
83 shape_kinds = [list(global_shape)]
84 for tensor_dim, size_options in sorted(dim_size_options.items()):
85 next_shape_kinds = []
86 for shape_kind in shape_kinds:
87 for size in size_options:
88 next_shape = list(shape_kind)
89 next_shape[tensor_dim] = size
90 next_shape_kinds.append(next_shape)
91 shape_kinds = next_shape_kinds
93 unique_shape_kinds = []
94 seen = set()
95 for shape_kind in shape_kinds:
96 shape_kind = tuple(shape_kind)
97 if shape_kind in seen:
98 continue
99 seen.add(shape_kind)
100 unique_shape_kinds.append(shape_kind)
101 return unique_shape_kinds
104def _debug_param_shard_metadata(hsdp_group, param_to_meta: Dict[torch.nn.Parameter, ParamShardMeta]) -> None:
105 """Log shard metadata for the group on rank 0 when distributed is ready."""
106 if not dist.is_available() or not dist.is_initialized():
107 return
109 rank = dist.get_rank()
110 if rank != 0:
111 return
113 for record in hsdp_group.records:
114 param = record.param
115 shard_meta = param_to_meta.get(param)
116 if shard_meta is None:
117 continue
118 local_shape_kinds = _get_local_shape_kinds(shard_meta.global_shape, shard_meta)
119 logger.debug_rank0(
120 "name=%s, placements=%s, mesh_shape=%s, global_shape=%s, local_shapes=%s",
121 getattr(param, "model_name", f"param_{record.index}"),
122 [str(placement) for placement in param.placements],
123 _get_mesh_shape(param.device_mesh),
124 shard_meta.global_shape,
125 local_shape_kinds,
126 )
129def _allgather_shapes_for_group(
130 local_shapes: List[Tuple[int, ...]],
131 shard_pg: Optional[dist.ProcessGroup],
132 shard_size: int,
133 device: torch.device,
134) -> List[List[Tuple[int, ...]]]:
135 """All-gather per-parameter shapes within one shard process group."""
136 if shard_pg is None or shard_size <= 1:
137 return [[tuple(shape) for shape in local_shapes]]
139 shape_tensor = torch.tensor(local_shapes, dtype=torch.int64, device=device)
140 gathered = torch.empty((shard_size, *shape_tensor.shape), dtype=torch.int64, device=device)
141 dist.all_gather_into_tensor(gathered.view(-1), shape_tensor.contiguous().view(-1), group=shard_pg)
142 gathered_cpu = gathered.cpu().tolist()
143 return [
144 [tuple(shape) for shape in rank_shapes]
145 for rank_shapes in gathered_cpu
146 ]
149def build_param_shard_metadata_for_group(
150 hsdp_group: HSDPCommGroup,
151) -> Dict[torch.Tensor, ParamShardMeta]:
152 """Build per-parameter shard metadata before HSDP batching."""
153 # pylint: disable=too-many-locals
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, _ 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 )
280 local_update = local_update.narrow(tensor_dim, sum(split_sizes[:local_rank]), split_sizes[local_rank])
282 if not local_update.is_contiguous():
283 local_update = local_update.contiguous()
285 return local_update
288def _get_or_alloc_buffer(
289 cache: Optional[Dict],
290 key: Any,
291 numel: int,
292 dtype: torch.dtype,
293 device: torch.device,
294) -> torch.Tensor:
295 """Return a cached buffer with at least `numel` elements."""
296 if cache is not None and key in cache:
297 buf = cache[key]
298 if buf.numel() >= numel:
299 return buf
300 buf = torch.empty(numel, dtype=dtype, device=device)
301 cache[key] = buf
302 return buf
304 buf = torch.empty(numel, dtype=dtype, device=device)
305 if cache is not None:
306 cache[key] = buf
307 return buf
310def _early_return_tensors(
311 local_tensors: List[torch.Tensor],
312 keep_indices: Optional[set],
313) -> List[Optional[torch.Tensor]]:
314 """Return tensors directly when no communication is needed."""
315 if keep_indices is None:
316 return list(local_tensors)
317 return [t if i in keep_indices else None for i, t in enumerate(local_tensors)]
320@dataclass(frozen=True)
321class _GatherParamMeta:
322 offset: int
323 actual_numel: int
324 padded_numel: int
325 rest_shape: Tuple[int, ...]
326 split_sizes: Tuple[int, ...]
329def _prepare_gather_inputs(
330 current_tensors: List[torch.Tensor],
331 tensor_dim: int,
332 alignment_elements: int,
333 shard_size: int,
334 stage_metas: Optional[List[Optional[ParamShardStageMeta]]] = None,
335) -> Tuple[List[torch.Tensor], List[_GatherParamMeta], int]:
336 """Move shard dim to dim0 and compute padding metadata."""
337 gather_inputs: List[torch.Tensor] = []
338 param_meta: List[_GatherParamMeta] = []
339 total_padded_numel = 0
341 for idx, gather_input in enumerate(current_tensors):
342 normalized_dim = tensor_dim % gather_input.dim()
343 if normalized_dim != 0 or not gather_input.is_contiguous():
344 gather_input = gather_input.movedim(normalized_dim, 0).contiguous()
345 stage_meta = stage_metas[idx] if stage_metas is not None else None
346 if stage_meta is None:
347 split_sizes = tuple(gather_input.shape[0] for _ in range(shard_size))
348 padded_numel_raw = gather_input.numel()
349 else:
350 split_sizes = stage_meta.split_sizes
351 pad_shape = list(stage_meta.pad_shape)
352 if normalized_dim != 0:
353 pad_shape[0], pad_shape[normalized_dim] = pad_shape[normalized_dim], pad_shape[0]
354 padded_numel_raw = math.prod(pad_shape)
355 param_meta.append(
356 _GatherParamMeta(
357 offset=total_padded_numel,
358 actual_numel=gather_input.numel(),
359 padded_numel=(
360 (padded_numel_raw + alignment_elements - 1) // alignment_elements * alignment_elements
361 ),
362 rest_shape=tuple(gather_input.shape[1:]),
363 split_sizes=split_sizes,
364 )
365 )
366 gather_inputs.append(gather_input)
367 total_padded_numel += param_meta[-1].padded_numel
369 return gather_inputs, param_meta, total_padded_numel
372def _pack_and_allgather(
373 gather_inputs: List[torch.Tensor],
374 param_meta: List[_GatherParamMeta],
375 total_padded_numel: int,
376 axis_idx: int,
377 dtype: torch.dtype,
378 device: torch.device,
379 shard_pg: dist.ProcessGroup,
380 shard_size: int,
381 buffer_cache: Optional[Dict],
382) -> torch.Tensor:
383 """Pack local shards into one buffer, all-gather, return gathered view."""
384 pack_buffer = _get_or_alloc_buffer(
385 buffer_cache, ("fused_allgather", axis_idx, dtype, device), total_padded_numel,
386 dtype, device,
387 )[:total_padded_numel]
389 pack_buffer.zero_()
390 for gi, meta in zip(gather_inputs, param_meta):
391 pack_buffer[meta.offset:meta.offset + meta.actual_numel].copy_(gi.view(-1))
393 gathered_buffer = _get_or_alloc_buffer(
394 buffer_cache,
395 ("fused_allgather_out", axis_idx, dtype, device),
396 total_padded_numel * shard_size,
397 dtype, device,
398 )[:total_padded_numel * shard_size]
400 dist.all_gather_into_tensor(gathered_buffer, pack_buffer, group=shard_pg)
401 return gathered_buffer.view(shard_size, total_padded_numel)
404def _unpack_gathered_results(
405 gathered_view: torch.Tensor,
406 gather_inputs: List[torch.Tensor],
407 param_meta: List[_GatherParamMeta],
408 current_tensors: List[torch.Tensor],
409 tensor_dim: int,
410 shard_size: int,
411 n_params: int,
412 is_last_axis: bool,
413 keep_indices: Optional[set],
414) -> List[Optional[torch.Tensor]]:
415 """Slice gathered buffer back to per-parameter full tensors."""
416 new_tensors: List[Optional[torch.Tensor]] = []
417 for i in range(n_params):
418 if is_last_axis and keep_indices is not None and i not in keep_indices:
419 new_tensors.append(None)
420 continue
422 meta = param_meta[i]
423 rest_shape = meta.rest_shape
424 rest_numel = math.prod(rest_shape) if rest_shape else 1
425 param_slice = gathered_view[:, meta.offset:meta.offset + meta.padded_numel]
427 param_chunks = []
428 for rank in range(shard_size):
429 valid_numel = meta.split_sizes[rank] * rest_numel
430 if valid_numel == 0:
431 continue
432 param_chunks.append(param_slice[rank, :valid_numel].contiguous().view(-1))
434 if param_chunks:
435 param_data = torch.cat(param_chunks, dim=0)
436 else:
437 param_data = gather_inputs[i].new_empty((0,))
439 result = param_data.view(sum(meta.split_sizes), *rest_shape)
440 tensor_dim_norm = tensor_dim % current_tensors[i].dim()
441 if tensor_dim_norm == 0:
442 new_tensors.append(result)
443 else:
444 new_tensors.append(result.movedim(0, tensor_dim_norm))
446 return new_tensors
449def fused_allgather_dtensor_params(
450 local_tensors: List[torch.Tensor],
451 shard_pgs: Sequence[dist.ProcessGroup],
452 layout_spec: ParamLayoutSpec,
453 param_shard_metadata: Optional[List[Optional[ParamShardMeta]]] = None,
454 buffer_cache: Optional[Dict] = None,
455 keep_indices: Optional[set] = None,
456) -> List[Optional[torch.Tensor]]:
457 """Fuse many parameter shards into one all-gather per shard axis."""
458 if not shard_pgs or not local_tensors:
459 return _early_return_tensors(local_tensors, keep_indices)
461 n_params = len(local_tensors)
462 device = local_tensors[0].device
463 dtype = local_tensors[0].dtype
464 alignment_bytes = 512
465 element_size = local_tensors[0].element_size()
466 alignment_elements = max(1, alignment_bytes // element_size)
468 active_axes = []
469 for axis_idx, ((_, tensor_dim), shard_pg) in enumerate(zip(layout_spec.shard_axes, shard_pgs)):
470 if shard_pg is None:
471 continue
472 shard_size = dist.get_world_size(shard_pg)
473 if shard_size <= 1:
474 continue
475 active_axes.append((axis_idx, tensor_dim, shard_pg, shard_size))
477 if not active_axes:
478 return _early_return_tensors(local_tensors, keep_indices)
480 current_tensors: List[torch.Tensor] = list(local_tensors)
482 for active_pos, (axis_idx, tensor_dim, shard_pg, shard_size) in enumerate(active_axes):
483 is_last_axis = active_pos == len(active_axes) - 1
484 stage_metas: Optional[List[Optional[ParamShardStageMeta]]] = None
485 if param_shard_metadata is not None:
486 stage_metas = []
487 for meta in param_shard_metadata:
488 if meta is None or axis_idx >= len(meta.stage_metas):
489 stage_metas.append(None)
490 else:
491 stage_metas.append(meta.stage_metas[axis_idx])
493 gather_inputs, param_meta, total_padded_numel = _prepare_gather_inputs(
494 current_tensors, tensor_dim, alignment_elements, shard_size, stage_metas,
495 )
497 gathered_view = _pack_and_allgather(
498 gather_inputs, param_meta, total_padded_numel,
499 axis_idx, dtype, device, shard_pg, shard_size, buffer_cache,
500 )
502 new_tensors = _unpack_gathered_results(
503 gathered_view, gather_inputs, param_meta,
504 current_tensors, tensor_dim, shard_size,
505 n_params, is_last_axis, keep_indices,
506 )
508 if is_last_axis:
509 return new_tensors
511 current_tensors = new_tensors # type: ignore[assignment]
513 return current_tensors