Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / platform.py: 61%
763 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2025-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"""Torch platform api"""
16from datetime import timedelta
17from typing import Optional, Any, Union
18import dataclasses
19from collections import OrderedDict
21import numpy as np
22from safetensors.torch import save_file, load_file
23import torch
24from torch import nn
25from torch import Tensor
26from torch._C._distributed_c10d import Store, ProcessGroup
27from torch.distributed import Backend
28from torch.distributed.distributed_c10d import _get_default_group
29from torch.nn import Parameter, Module
30from torch.nn.utils.rnn import PackedSequence
31from torch._ops import OpOverload, OpOverloadPacket
32from torch.utils.checkpoint import noop_context_fn
34import torch.distributed.nn.functional as dist_func
35import torch.distributed as dist
36from hyper_parallel.platform.torch.dtensor import DTensorBase
37from hyper_parallel.platform.torch.pipeline_parallel.stage import PipelineStageBase
38from hyper_parallel.platform.torch.group_utils import create_sub_groups
39from hyper_parallel.platform.platform import Platform, PlatformType, EXISTING_COMM_GROUPS
40from hyper_parallel.platform.torch.function_override import override_functions
41from hyper_parallel.platform.torch.init_weights import init_on_device as _init_on_device
43override_functions()
46# ---------------------------------------------------------------------------
47# Module-level A2A reshape helpers
48# ---------------------------------------------------------------------------
50def _a2a_reconstruct(out_perm: torch.Tensor, concat_dim: int) -> torch.Tensor:
51 """Reconstruct A2A result from raw out_perm buffer.
53 ``out_perm`` has shape ``[ws, *rest_dims]``, chunk at ``concat_dim + 1``.
54 Returns tensor with merged chunk dimension.
55 """
56 new_ndim = out_perm.dim()
57 chunk_in_perm = concat_dim + 1
58 recon_perm = list(range(1, chunk_in_perm)) + [0] + list(range(chunk_in_perm, new_ndim))
59 x_recon = out_perm.permute(recon_perm).contiguous()
60 shape = list(x_recon.shape)
61 merged = shape[concat_dim] * shape[concat_dim + 1]
62 return x_recon.reshape(shape[:concat_dim] + [merged] + shape[concat_dim + 2:])
65def _normalize_dim(dim: int, ndim: int) -> int:
66 """Normalize a possibly negative dimension index."""
67 return dim + ndim if dim < 0 else dim
70def _move_dim_to_front(tensor: torch.Tensor, dim: int) -> torch.Tensor:
71 """Move ``dim`` to the front while keeping the other dimensions ordered."""
72 dim = _normalize_dim(dim, tensor.dim())
73 if dim == 0:
74 return tensor.contiguous()
75 perm = [dim] + [i for i in range(tensor.dim()) if i != dim]
76 return tensor.permute(perm).contiguous()
79def _move_dim_from_front(tensor: torch.Tensor, dim: int) -> torch.Tensor:
80 """Inverse of :func:`_move_dim_to_front`."""
81 dim = _normalize_dim(dim, tensor.dim())
82 if dim == 0:
83 return tensor.contiguous()
84 perm = [dim] + [i for i in range(tensor.dim()) if i != dim]
85 inverse = [0] * len(perm)
86 for idx, value in enumerate(perm):
87 inverse[value] = idx
88 return tensor.permute(inverse).contiguous()
91class _TorchAsyncA2AFunction(torch.autograd.Function):
92 """Differentiable wrapper for pre-launched async all-to-all.
94 Forward: wait async handle, reconstruct A2A result.
95 Backward: launch async head→seq A2A and store handle in ``handle_box``
96 for the projection pre-hook to wait, achieving GEMM–A2A overlap.
97 """
99 @staticmethod
100 def forward(ctx, x, work, out_perm, group, world_size, concat_dim, split_dim, # pylint: disable=arguments-differ
101 handle_box):
102 """Wait for pre-launched async A2A and return reconstructed output."""
103 ctx.group = group
104 ctx.world_size = world_size
105 ctx.concat_dim = concat_dim
106 ctx.split_dim = split_dim
107 ctx.handle_box = handle_box
108 ctx.x_shape = x.shape
109 work.wait()
110 return _a2a_reconstruct(out_perm, concat_dim)
112 @staticmethod
113 def backward(ctx, grad_output):
114 """Launch async head→seq A2A for backward overlap, or return zero grad."""
115 if ctx.handle_box is not None:
116 # Launch async head→seq A2A (reverse of forward seq→head)
117 g = grad_output.contiguous()
118 shape = list(g.shape)
119 seq_dim = ctx.concat_dim
120 s_full = shape[seq_dim]
121 ndim = len(shape) + 1
122 x_perm = g.reshape(
123 shape[:seq_dim] + [ctx.world_size, s_full // ctx.world_size] + shape[seq_dim + 1:]
124 ).permute(
125 [seq_dim] + list(range(seq_dim)) + list(range(seq_dim + 1, ndim))
126 ).contiguous()
127 out_perm = torch.empty_like(x_perm)
128 work = dist.all_to_all_single(out_perm, x_perm, group=ctx.group, async_op=True)
129 ctx.handle_box.append((work, out_perm))
130 return grad_output.new_zeros(ctx.x_shape), None, None, None, None, None, None, None
133class _TorchAsyncAllGatherFunction(torch.autograd.Function):
134 """Differentiable wrapper for pre-launched async all-gather."""
136 @staticmethod
137 def forward(ctx, x, work, out_perm, group, world_size, gather_dim, handle_box): # pylint: disable=arguments-differ
138 """Wait for pre-launched all-gather and reconstruct the gathered tensor."""
139 ctx.group = group
140 ctx.world_size = world_size
141 ctx.gather_dim = gather_dim
142 ctx.handle_box = handle_box
143 ctx.x_shape = x.shape
144 work.wait()
145 return _move_dim_from_front(out_perm, gather_dim)
147 @staticmethod
148 def backward(ctx, grad_output):
149 """Launch reverse reduce-scatter for the all-gather."""
150 grad_perm = _move_dim_to_front(grad_output.contiguous(), ctx.gather_dim)
151 output_shape = list(grad_perm.shape)
152 if output_shape[0] % ctx.world_size != 0:
153 raise ValueError(
154 "all_gather backward expected gathered dimension to be divisible by world_size, "
155 f"got {output_shape[0]} and {ctx.world_size}."
156 )
157 output_shape[0] //= ctx.world_size
158 output = torch.empty(output_shape, dtype=grad_perm.dtype, device=grad_perm.device)
159 work = dist.reduce_scatter_tensor(output, grad_perm, group=ctx.group, async_op=True)
160 if ctx.handle_box is not None:
161 ctx.handle_box.append((work, output, ctx.gather_dim))
162 return grad_output.new_zeros(ctx.x_shape), None, None, None, None, None, None
163 work.wait()
164 return _move_dim_from_front(output, ctx.gather_dim), None, None, None, None, None, None
167class _AsyncA2ALazyBwd(torch.autograd.Function):
168 """All-to-all whose forward AND backward return ``AsyncCollectiveTensor``.
170 PyTorch's stock ``all_to_all_single_autograd`` calls ``wait_tensor`` in
171 its backward eagerly, and the autograd engine binds backward stream
172 context to the forward stream — so even if the BWD thread is wrapped
173 in a side-stream context, that wait still lands on the FWD main
174 stream and blocks Attention launches.
176 This Function bypasses the engine's binding by calling the
177 non-autograd functional op in both directions and returning ACT.
178 The wait is deferred to the next consumer's first non-view access
179 (e.g. the indexing backward of ``_unpermute``), giving the FWD
180 thread a small Python window to enqueue its Attention kernels onto
181 the main stream **before** the wait lands there.
182 """
184 @staticmethod
185 def forward(ctx, input_tensor, output_splits, input_splits, group): # pylint: disable=arguments-differ
186 """Perform the forward all-to-all single collective, saving splits and group for backward."""
187 ctx.input_splits = input_splits
188 ctx.output_splits = output_splits
189 ctx.group = group
190 # pylint: disable=C0415
191 from torch.distributed._functional_collectives import all_to_all_single
192 return all_to_all_single(
193 input_tensor, output_splits, input_splits, group,
194 )
196 @staticmethod
197 def backward(ctx, grad_output):
198 """Compute the backward pass by performing the inverse all-to-all with swapped splits."""
199 # pylint: disable=C0415
200 from torch.distributed._functional_collectives import all_to_all_single
201 grad_input = all_to_all_single(
202 grad_output, ctx.input_splits, ctx.output_splits, ctx.group,
203 )
204 return grad_input, None, None, None
207class _TorchSyncHookFunction(torch.autograd.Function):
208 """Autograd identity that fires HookCoordinator rendezvous on fwd/bwd.
210 Uses a **4-hook** design (``A``, ``B``, ``C``, ``D``) with pure
211 COMM / COMPUTE roles — no NONE role. Every rendezvous is a strict
212 COMM + COMPUTE pair, guaranteeing NCCL-first dispatch ordering at
213 **all** points including layer boundaries.
215 Hook placement per MoE layer::
217 [A] → dispatch → [B] → module → [C] → combine → [D] → (Attention) → [A_next]
219 At layer boundaries (D / A hooks), the Attention that runs between
220 layers is treated as COMPUTE, and the combine / combine.bwd is treated
221 as COMM, so the coordinator enforces comm-first ordering even across
222 layer transitions.
223 """
225 # 4-hook role tables: (prev_role_idx, next_role_idx).
226 # Index encoding: 1 = COMM, 2 = COMPUTE.
227 #
228 # Torch only uses the four core hooks A/B/C/D + D_LAST sentinel.
229 # The MS backend adds ``CHUNK_START`` / ``CHUNK_END`` because of
230 # MS-specific issues (stream binding follows the calling thread;
231 # autograd cannot have FWD-record + BWD-replay concurrently).
232 # Torch has neither problem — CUDA streams are process-wide and
233 # Torch autograd is thread-safe — so we keep the original
234 # 4-hook design here. Do not add CHUNK_START / CHUNK_END to
235 # the Torch tables; if a future test does need them, copy the
236 # MS implementation and add the matching skip rules in
237 # ``forward`` / ``backward``.
238 _FWD_ROLES = {
239 # (prev, next) prev op next op
240 "A": (2, 1), # COMPUTE, COMM Attention | dispatch
241 "B": (1, 2), # COMM, COMPUTE dispatch | module
242 "C": (2, 1), # COMPUTE, COMM module | combine
243 "D": (1, 2), # COMM, COMPUTE combine | Attention
244 }
245 _BWD_ROLES = {
246 "D": (2, 1), # COMPUTE, COMM Attn.bwd | combine.bwd
247 "C": (1, 2), # COMM, COMPUTE combine.bwd | module.bwd
248 "B": (2, 1), # COMPUTE, COMM module.bwd | dispatch.bwd
249 "A": (1, 2), # COMM, COMPUTE dispatch.bwd| Attn.bwd
250 }
252 _ROLE_CACHE = None
254 @staticmethod
255 def _role_enum(idx: int):
256 if _TorchSyncHookFunction._ROLE_CACHE is None:
257 from hyper_parallel.core.pipeline_parallel.hook_coordinator import HookRole # pylint: disable=C0415
258 _TorchSyncHookFunction._ROLE_CACHE = (None, HookRole.COMM, HookRole.COMPUTE)
259 return _TorchSyncHookFunction._ROLE_CACHE[idx]
261 @staticmethod
262 def forward(ctx, x, hook_name, coordinator): # pylint: disable=arguments-differ
263 """Identity forward that fires a HookCoordinator rendezvous.
265 Notifies the previous op's role and rendezvouses for the next op's
266 role per the ``_FWD_ROLES`` table. ``"D_LAST"`` is a sentinel
267 meaning "skip this rendezvous" (last layer's closing D — no
268 Attention follows).
270 Args:
271 ctx: Autograd context, stores ``hook_name`` and
272 ``coordinator`` for the backward pass.
273 x: Input tensor, returned unchanged.
274 hook_name: One of ``"A"``, ``"B"``, ``"C"``, ``"D"``,
275 ``"D_LAST"``.
276 coordinator: The :class:`HookCoordinator` driving the rendezvous.
278 Returns:
279 ``x`` unchanged.
280 """
281 ctx.hook_name = hook_name
282 ctx.coordinator = coordinator
284 if not coordinator.is_enabled():
285 return x
287 if hook_name == "D_LAST":
288 # ``D_LAST`` marks the last layer's closing D hook — no
289 # Attention follows in this chunk, so the rendezvous is
290 # meaningless and is skipped. We still
291 # ``notify_dispatched(COMM)`` so the COMPUTE side of the
292 # preceding ``C`` rendezvous unblocks early, letting
293 # BWD's Attn.bwd_last overlap with FWD's post-combine
294 # work — Torch autograd is thread-safe so this concurrent
295 # FWD-record + BWD-replay is fine.
296 prev_idx, _ = _TorchSyncHookFunction._FWD_ROLES["D"]
297 role_of = _TorchSyncHookFunction._role_enum
298 coordinator.notify_dispatched(role_of(prev_idx))
299 return x
301 prev_idx, next_idx = _TorchSyncHookFunction._FWD_ROLES[hook_name]
302 role_of = _TorchSyncHookFunction._role_enum
303 coordinator.notify_dispatched(role_of(prev_idx))
304 coordinator.rendezvous(role_of(next_idx))
305 return x
307 @staticmethod
308 def backward(ctx, grad_output):
309 """Identity backward that fires a HookCoordinator rendezvous.
311 Mirror of :meth:`forward` using the ``_BWD_ROLES`` table.
312 ``"D_LAST"`` skips the rendezvous because this is the first BWD
313 hook to fire and ``combine.bwd`` has already dispatched freely
314 before any rendezvous can happen.
316 Args:
317 ctx: Autograd context with ``hook_name`` and
318 ``coordinator`` saved during forward.
319 grad_output: Gradient w.r.t. the forward output, returned
320 unchanged.
322 Returns:
323 ``(grad_output, None, None)`` — gradients only flow back to
324 the tensor input, ``hook_name`` and ``coordinator`` are
325 non-tensor inputs.
326 """
327 hook_name = ctx.hook_name
328 coordinator = ctx.coordinator
330 if not coordinator.is_enabled():
331 return grad_output, None, None
333 if hook_name == "D_LAST":
334 # First BWD hook to fire; combine.bwd has already
335 # dispatched freely before any rendezvous can happen.
336 # Skipping here is safe on Torch because CUDA streams
337 # are process-wide and the NCCL FIFO order is consistent
338 # across ranks regardless of which thread launched
339 # combine.bwd.
340 return grad_output, None, None
342 prev_idx, next_idx = _TorchSyncHookFunction._BWD_ROLES[hook_name]
343 role_of = _TorchSyncHookFunction._role_enum
344 coordinator.notify_dispatched(role_of(prev_idx))
345 coordinator.rendezvous(role_of(next_idx))
346 return grad_output, None, None
349class _TorchP2PExchangeFunction(torch.autograd.Function):
350 """Symmetric bidirectional P2P: send local tensor to peer, receive peer's tensor."""
352 @staticmethod
353 def forward(ctx, tensor: torch.Tensor, peer_rank: int, group) -> torch.Tensor: # pylint: disable=arguments-differ
354 """Perform symmetric bidirectional P2P exchange with peer_rank."""
355 ctx.peer_rank = peer_rank
356 ctx.group = group
357 send_buf = tensor.contiguous()
358 recv_buf = torch.empty_like(send_buf)
359 reqs = dist.batch_isend_irecv([
360 dist.P2POp(dist.isend, send_buf, peer_rank, group),
361 dist.P2POp(dist.irecv, recv_buf, peer_rank, group),
362 ])
363 for req in reqs:
364 req.wait()
365 return recv_buf
367 @staticmethod
368 def backward(ctx, grad_output: torch.Tensor):
369 """Perform symmetric P2P exchange for the backward gradient pass."""
370 send_buf = grad_output.contiguous()
371 recv_buf = torch.empty_like(send_buf)
372 reqs = dist.batch_isend_irecv([
373 dist.P2POp(dist.isend, send_buf, ctx.peer_rank, ctx.group),
374 dist.P2POp(dist.irecv, recv_buf, ctx.peer_rank, ctx.group),
375 ])
376 for req in reqs:
377 req.wait()
378 return recv_buf, None, None
381# Mapping from string op names to torch.distributed.ReduceOp
382_OP_MAP = {
383 'sum': dist.ReduceOp.SUM,
384 'prod': dist.ReduceOp.PRODUCT,
385 'max': dist.ReduceOp.MAX,
386 'min': dist.ReduceOp.MIN,
387 # convert tensor elements to int32 and use MIN
388 'all': dist.ReduceOp.MIN,
389 # 'avg' is typically handled by SUM followed by division in current implementation logic
390 'avg': dist.ReduceOp.SUM,
391}
393# Try to add AVG for 'mean' if supported by current torch version
394if hasattr(dist.ReduceOp, "AVG"):
395 _OP_MAP['mean'] = dist.ReduceOp.AVG
396else:
397 # Fallback for older torch versions if necessary, though this might require manual division upstream
398 # Assuming standard behavior where 'mean' implies native AVG support or upstream handling
399 _OP_MAP['mean'] = dist.ReduceOp.SUM
402def _ensure_contiguous(x):
403 """Return a contiguous copy of *x* if not already contiguous."""
404 if not x.is_contiguous() or x.storage_offset() != 0:
405 x = x.contiguous()
406 return x
409class _TorchBatchP2PWork:
410 """Single ``.wait()`` handle wrapping the per-op works returned by
411 ``torch.distributed.batch_isend_irecv``.
413 Torch returns one ``Work`` per op in the batch (the ops are coalesced
414 onto one comm stream), whereas the platform contract — and the scheduler
415 that consumes it — expects a single handle covering the whole batch so
416 the wait can be deferred to one consumption point (mirroring MindSpore's
417 single packaging ``CommHandle``). Waiting this handle waits every
418 underlying op.
419 """
421 __slots__ = ("_works",)
423 def __init__(self, works):
424 self._works = works
426 def wait(self):
427 for work in self._works:
428 if work is not None:
429 work.wait()
432# pylint: disable=C0103
433class TorchPlatform(Platform):
434 """Torch platform api"""
435 Tensor = Tensor
436 tensor = torch.tensor
437 Parameter = Parameter
438 Module = Module
439 DTensorBase = DTensorBase
440 PipelineStageBase = PipelineStageBase
441 platform_type = PlatformType.PYTORCH
442 tensor_dtype = torch
443 dtype = torch.dtype
444 Function = torch.autograd.Function
446 _custom_ops_cls = None
448 @property
449 def custom_ops(self):
450 """Return the Torch platform custom ops instance.
452 .. warning::
453 This is an experimental API that subject to change or deletion.
455 Returns:
456 TorchCustomOps: Custom ops class that raises NotImplementedError
457 for all operators (MindSpore-only at this time).
458 """
459 if self._custom_ops_cls is None:
460 from hyper_parallel.platform.torch.custom_ops import TorchCustomOps # pylint: disable=import-outside-toplevel
461 self._custom_ops_cls = TorchCustomOps
462 return self._custom_ops_cls
464 @staticmethod
465 def is_linear_module(module) -> bool:
466 """Check whether *module* is a ``torch.nn.Linear`` instance."""
467 return isinstance(module, nn.Linear)
469 @staticmethod
470 def is_embedding_module(module) -> bool:
471 """Check whether *module* is a ``torch.nn.Embedding`` instance."""
472 return isinstance(module, nn.Embedding)
474 @staticmethod
475 def device_count(device_handle):
476 """
477 Get the number of available devices.
479 Args:
480 device_handle: The device handle (e.g., torch.cuda, torch.npu).
482 Returns:
483 int: The number of available devices.
484 """
485 return device_handle.device_count()
487 def device_type(self):
488 """
489 Get the current device type.
491 Returns:
492 str: The device type string ("npu" for NPU, "cuda" for GPU).
493 """
494 device_handle = self.get_device_handle()
495 if device_handle == torch.npu:
496 return "npu"
497 return "cuda"
499 def device(self, device_idx=None):
500 """
501 Get a torch.device object for the specified device index.
503 Args:
504 device_idx (Optional[int]): The device index. If None, returns device without index.
506 Returns:
507 torch.device: A torch device object.
508 """
509 device_type = self.device_type()
510 if device_idx is None:
511 return torch.device(device_type)
512 return torch.device(f"{device_type}:{device_idx:d}")
514 @staticmethod
515 def get_rng_state(device=None, device_handle=None):
516 """
517 Get the random number generator state.
519 Args:
520 device (Optional): The device to get RNG state from.
521 device_handle (Optional): The device handle (torch.cuda, torch.npu, etc.).
523 Returns:
524 Tensor: The RNG state as a byte tensor.
525 """
526 if device_handle is None:
527 return torch.get_rng_state()
528 if device is None:
529 return device_handle.get_rng_state()
530 return device_handle.get_rng_state(device)
532 @staticmethod
533 def set_rng_state(state, device=None, device_handle=None):
534 """
535 Set the random number generator state.
537 Args:
538 state (Tensor): The RNG state to set.
539 device (Optional): The device to set RNG state for.
540 device_handle (Optional): The device handle (torch.cuda, torch.npu, etc.).
541 """
542 if device_handle is None:
543 return torch.set_rng_state(state)
544 if device is None:
545 return device_handle.set_rng_state(state)
546 return device_handle.set_rng_state(state, device)
548 @staticmethod
549 def manual_seed(seed):
550 """
551 Set the random seed for reproducibility.
553 Args:
554 seed (int): The random seed value.
556 Returns:
557 torch.Generator: The random number generator.
558 """
559 return torch.manual_seed(seed)
561 @staticmethod
562 def ones(size, dtype=None):
563 """
564 Create a tensor filled with ones.
566 Args:
567 size (tuple): The shape of the output tensor.
568 dtype (Optional[torch.dtype]): The desired data type.
570 Returns:
571 Tensor: A tensor filled with ones.
572 """
573 return torch.ones(size, dtype=dtype)
575 @staticmethod
576 def zeros(size, dtype=None, device=None):
577 """
578 Create a tensor filled with zeros.
580 Args:
581 size (tuple): The shape of the output tensor.
582 dtype (Optional[torch.dtype]): The desired data type.
583 device (Optional[torch.device]): The device to create the tensor on.
585 Returns:
586 Tensor: A tensor filled with zeros.
587 """
588 return torch.zeros(size, dtype=dtype, device=device)
590 @staticmethod
591 def full(size, fill_value, dtype=None):
592 """
593 Create a tensor filled with a scalar value.
595 Args:
596 size (tuple): The shape of the output tensor.
597 fill_value (scalar): The value to fill the tensor with.
598 dtype (Optional[torch.dtype]): The desired data type.
600 Returns:
601 Tensor: A tensor filled with the specified value.
602 """
603 return torch.full(size, fill_value, dtype=dtype)
605 @staticmethod
606 def empty(size, dtype=None, device=None):
607 """
608 Create an uninitialized tensor.
610 Args:
611 size (tuple): The shape of the output tensor.
612 dtype (Optional[torch.dtype]): The desired data type.
613 device (Optional[torch.device or str]): Target device. When
614 ``None`` the tensor is allocated on the default device
615 (CPU under PyTorch defaults), matching the original
616 back-compat behavior.
618 Returns:
619 Tensor: An uninitialized tensor.
620 """
621 return torch.empty(size, dtype=dtype, device=device)
623 @staticmethod
624 def rand(size, dtype=None, device=None):
625 """Create a tensor filled with uniform random values in ``[0, 1)``."""
626 return torch.rand(size, dtype=dtype, device=device)
628 @staticmethod
629 def randn(size, dtype=None, device=None):
630 """Create a tensor filled with standard-normal random values."""
631 return torch.randn(size, dtype=dtype, device=device)
633 @staticmethod
634 def get_rank():
635 """
636 Get the rank of the current process in the distributed group.
638 Returns:
639 int: The rank of the current process.
640 """
641 return dist.get_rank()
643 @staticmethod
644 def get_global_rank(group, group_rank):
645 """
646 Get the global rank from a group rank.
648 Args:
649 group (ProcessGroup): The process group.
650 group_rank (int): The rank within the group.
652 Returns:
653 int: The global rank.
654 """
655 return dist.get_global_rank(group, group_rank)
657 @staticmethod
658 def get_group_rank(group):
659 """Return this process's rank within *group*."""
660 return dist.get_group_rank(group, dist.get_rank())
662 @staticmethod
663 def get_world_size():
664 """
665 Get the total number of processes in the distributed group.
667 Returns:
668 int: The world size.
669 """
670 return dist.get_world_size()
672 @staticmethod
673 def get_param_local_shape(param):
674 """
675 Get the local shape of a parameter, handling both regular and distributed tensors.
677 Args:
678 param (Union[Tensor, DTensorBase]): The parameter tensor.
680 Returns:
681 torch.Size: The local shape of the parameter.
682 """
683 if isinstance(param, DTensorBase):
684 return param.local_shape
685 return param.shape
687 @staticmethod
688 def get_param_local_data(param):
689 """
690 Get the local data of a parameter, handling both regular and distributed tensors.
692 Args:
693 param (Union[Tensor, DTensorBase]): The parameter tensor.
695 Returns:
696 Tensor: The local tensor data.
697 """
698 if isinstance(param, DTensorBase):
699 return param.to_local()
700 return param
702 @staticmethod
703 def update_param_data(param, data):
704 """
705 Update the data of a parameter.
707 Args:
708 param (Parameter): The parameter to update.
709 data (Tensor): The new data tensor.
710 """
711 param.data = data
713 @staticmethod
714 def load_into_param(param, data):
715 """Load tensor *data* into *param* (plain tensor or DTensor)."""
716 if isinstance(param, DTensorBase):
717 local = param._local_tensor # pylint: disable=W0212
718 if local.is_meta:
719 # Meta tensor materialisation: replace the placeholder.
720 orig_requires_grad = param.requires_grad
721 param._local_tensor = data # pylint: disable=W0212
722 if data.requires_grad != orig_requires_grad:
723 param.requires_grad_(orig_requires_grad)
724 else:
725 local.copy_(data)
726 else:
727 param.copy_(data)
729 @staticmethod
730 def get_op_name(func):
731 """
732 Extract the operation name from various function types.
734 Args:
735 func: The function or operation to extract the name from.
737 Returns:
738 str: The operation name.
739 """
740 if hasattr(func, "__name__"):
741 return func.__name__
742 if isinstance(func, OpOverload):
743 full_name = func.name
744 core_name = full_name.split("::")[-1].split(".")[0]
745 return core_name
746 if isinstance(func, OpOverloadPacket):
747 return func.name.split("::")[-1]
748 func_str = str(func)
749 if "built-in function" in func_str:
750 return func_str.split()[-1].strip(">")
751 if "function" in func_str:
752 return func_str.split()[1]
753 return "unknown_op"
755 @staticmethod
756 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
757 data = _ensure_contiguous(data)
758 output = list(dist_func.all_gather(data, group=group))
759 if rank_list is not None:
760 group_ranks = dist.get_process_group_ranks(group)
761 if tuple(rank_list) != tuple(group_ranks):
762 rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}
763 output = [output[rank_to_idx[int(rank)]] for rank in rank_list]
764 return torch.cat(output, dim=concat_dim)
766 @staticmethod
767 def chunk(data, split_dim, split_size, index):
768 return torch.chunk(data, split_size, dim=split_dim)[index]
770 @staticmethod
771 def differentiable_all_to_all(input_data, output_shape, group):
772 input_data = _ensure_contiguous(input_data)
773 output_tensor = torch.empty(output_shape, device=input_data.device, dtype=input_data.dtype)
774 output_tensor = dist_func.all_to_all_single(
775 output_tensor,
776 input_data,
777 group=group
778 )
779 return output_tensor
781 @staticmethod
782 def tensor_type_cast(input_data, cast_type):
783 """Cast tensor to specified data type."""
784 type_mapping = {
785 'float32': torch.float32,
786 'float16': torch.float16,
787 'int64': torch.int64,
788 'int32': torch.int32
789 }
790 if cast_type not in type_mapping:
791 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")
792 return input_data.to(type_mapping[cast_type])
794 @staticmethod
795 def differentiable_all_reduce(data, op, group):
796 data = _ensure_contiguous(data)
797 # Resolve the op from string to ReduceOp enum if necessary
798 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
799 return dist_func.all_reduce(data, op=reduce_op, group=group)
801 @staticmethod
802 def get_cell_construct(cell):
803 return cell.forward
805 @staticmethod
806 def get_cells_and_names(cell):
807 return cell.named_modules()
809 @staticmethod
810 def get_modules(module):
811 return module.modules()
813 @staticmethod
814 def search_parameter_by_name(cell, param_name: str):
815 """
816 Find the parent Module of the parameter, the parameter's name in the parent Module, and the parameter.
817 Return value: (parent Module instance, parameter's name in parent Module, parameter object).
818 Returns None if not found.
819 """
820 # Remove the "self." prefix from param_name
821 param_name = param_name.replace("self.", "")
822 # Case 1: The parameter is a direct parameter of the current Module
823 if param_name in cell._parameters: # pylint: disable=protected-access
824 return (cell, param_name, cell._parameters[param_name]) # pylint: disable=protected-access
826 # Case 2: The parameter is in a sub-Module
827 if "." in param_name:
828 cell_path, param_key = param_name.rsplit(".", 1)
829 try:
830 # Locate the sub-Module where the parameter resides (supports multi-level paths)
831 target_cell = cell.get_submodule(cell_path)
832 # Check if the sub-Module directly contains this parameter
833 if param_key in target_cell._parameters: # pylint: disable=protected-access
834 return target_cell, param_key, target_cell._parameters[param_key] # pylint: disable=protected-access
835 except AttributeError:
836 pass
838 # Traverse all sub-Modules (recursively) to search for the parameter
839 for _, child_cell in cell.named_children():
840 if isinstance(child_cell, Module):
841 result = TorchPlatform.search_parameter_by_name(child_cell, param_name)
842 if result is not None:
843 return result
845 return None
847 @staticmethod
848 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
849 """
850 Modify the original parameter in a Module or sub-Module using the search result
851 """
852 parent_cell, param_key, _ = result
853 # Key operation: directly modify the _parameters dictionary.
854 if param_key in parent_cell._parameters: # pylint: disable=protected-access
855 parent_cell._parameters[param_key] = new_param # pylint: disable=protected-access
856 else:
857 parent_cell.register_parameter(param_key, new_param)
858 return True
860 @staticmethod
861 def set_layout_into_parameter(param, layout):
862 """Set layout into parameter"""
863 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel
864 from hyper_parallel.core.dtensor.layout import _get_slice_tensor_by_layout # pylint: disable=import-outside-toplevel
865 if isinstance(param, DTensor):
866 raise ValueError(f"Parameter {param} has been configured layout, cannot be set repeatedly.")
867 requires_grad = param.requires_grad
868 param_dtensor = DTensor.from_local(
869 _get_slice_tensor_by_layout(param, layout),
870 layout.mesh, layout.alias_placements)
871 new_param = Parameter(param_dtensor, requires_grad=requires_grad)
872 return new_param
874 @staticmethod
875 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
876 data = _ensure_contiguous(data)
877 input_tuple = torch.chunk(data, dev_num, dim=axis)
878 output_tensor = torch.empty(input_tuple[0].shape, device=data.device, dtype=data.dtype)
880 # Resolve the op from string to ReduceOp enum
881 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
883 output_tensor = dist_func.reduce_scatter(output_tensor, input_tuple, op=reduce_op, group=group)
885 # Keep manual handling for 'avg' string as it maps to SUM in _OP_MAP
886 if op == 'avg':
887 output_tensor = output_tensor / dev_num
888 return output_tensor
890 @staticmethod
891 def get_device_handle(device_type: str = "npu"):
892 """Return the torch device module (e.g. ``torch.npu`` or ``torch.cuda``) for the given device type."""
893 try:
894 handle = getattr(torch, device_type)
895 except AttributeError as e:
896 raise RuntimeError(f"TorchPlatform expect got device handle: 'torch.{device_type}' failed.") from e
897 return handle
899 @staticmethod
900 def get_param_type_size(param):
901 # pylint: disable=W0212
902 return torch._utils._element_size(param.dtype)
904 @staticmethod
905 def is_tensor(obj: Any) -> bool:
906 """Return True if ``obj`` is a ``torch.Tensor``."""
907 return isinstance(obj, Tensor)
909 @staticmethod
910 def get_tensor_storage_size(tensor: Any) -> int:
911 """Return serialized byte size (numel * element size) for a PyTorch tensor."""
912 if not TorchPlatform.is_tensor(tensor):
913 raise TypeError(
914 f"TorchPlatform.get_tensor_storage_size expects torch.Tensor, got {type(tensor)!r}"
915 )
916 return int(tensor.numel()) * int(tensor.element_size())
918 @staticmethod
919 def parameters_dict(cell: Module):
920 return cell.named_parameters()
922 @staticmethod
923 def buffers_dict(cell: Module) -> Any:
924 """Return all named buffers registered by the module tree."""
925 return cell.named_buffers()
927 @staticmethod
928 def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
929 """Get the state dictionary of a model.
931 Delegates to torch-specific implementation that handles DTensor
932 gathering, CPU offloading and frozen-parameter filtering.
933 """
934 # pylint: disable=C0415
935 from hyper_parallel.platform.torch.fully_shard.state_dict_utils import (
936 get_model_state_dict as _get_model_state_dict,
937 )
938 return _get_model_state_dict(model, options=options)
940 @staticmethod
941 def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
942 """Set the state dictionary of a model.
944 Delegates to torch-specific implementation that scatters full tensors
945 into DTensor shards and performs an in-place load.
946 """
947 # pylint: disable=C0415
948 from hyper_parallel.platform.torch.fully_shard.state_dict_utils import (
949 set_model_state_dict as _set_model_state_dict,
950 )
951 return _set_model_state_dict(model, model_state_dict, options=options)
953 @staticmethod
954 def save_checkpoint(cell: Module, file_path: str, ckpt_format: str = "safetensors") -> None:
955 if ckpt_format == "safetensors":
956 save_file(tensors=cell, filename=file_path)
957 else:
958 torch.save(obj=cell, f=file_path)
960 @staticmethod
961 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
962 if ckpt_format == "safetensors":
963 return load_file(filename=file_path)
964 return torch.load(f=file_path)
966 @staticmethod
967 def new_zero_parameter(param_shape, param_type, requires_grad, device):
968 return nn.Parameter(torch.zeros(param_shape, dtype=param_type, device=device), requires_grad=requires_grad)
970 @staticmethod
971 def new_tensor(tensor_shape, tensor_type, device):
972 return torch.empty(size=tensor_shape, dtype=tensor_type, device=device)
974 @staticmethod
975 def full_like(tensor, fill_value, dtype=None):
976 return torch.full_like(tensor, fill_value, dtype=dtype)
978 @staticmethod
979 def set_tensor_requires_grad(input_tensor):
980 """
981 set requires grad flag for input tensor, only effective for leaf node
982 """
983 if input_tensor.is_leaf:
984 input_tensor.requires_grad = True
986 def _create_group(self, rank_list):
987 normalized_rank_list = tuple(sorted(rank_list))
988 world_rank_list = tuple(range(self.get_world_size()))
989 if normalized_rank_list == world_rank_list:
990 group = _get_default_group()
991 EXISTING_COMM_GROUPS[str(normalized_rank_list)] = group
992 return group
993 group_dict = create_sub_groups(rank_list)
994 return group_dict[normalized_rank_list]
996 @staticmethod
997 def all_gather_into_tensor(data, group_info, async_op=False):
998 output_shape = list(data.shape)
999 output_shape[0] = output_shape[0] * group_info.rank_size
1000 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
1001 handle = dist.all_gather_into_tensor(output, data, group=group_info.group, async_op=async_op)
1002 return output, handle
1004 @staticmethod
1005 def all_gather_single(input_tensor, output_shape, group, async_op=False):
1006 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
1007 handle = dist.all_gather_into_tensor(output, input_tensor, group=group, async_op=async_op)
1008 return output, handle
1010 @staticmethod
1011 def all_reduce(data, group_info, async_op=False):
1012 if not data.is_contiguous():
1013 data = data.contiguous()
1014 handle = dist.all_reduce(data, group=group_info.group, async_op=async_op)
1015 return data, handle
1017 @staticmethod
1018 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
1019 if group_src is not None:
1020 src = dist.get_global_rank(group, group_src)
1021 handle = dist.broadcast(data, src, group, async_op)
1022 if async_op and handle is not None:
1023 handle.wait()
1025 @staticmethod
1026 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
1027 if group_src is not None:
1028 src = dist.get_global_rank(group, group_src)
1029 handle = dist.scatter(output, scatter_list, src=src, group=group, async_op=async_op)
1030 if async_op and handle is not None:
1031 handle.wait()
1032 return output
1034 @staticmethod
1035 def isend(tensor, dst=None, group=None, tag=0):
1036 return dist.isend(tensor, dst, group, tag)
1038 @staticmethod
1039 def irecv(tensor, src=None, group=None, tag=0):
1040 return dist.irecv(tensor, src, group, tag)
1042 @staticmethod
1043 def p2p_op(op_type, tensor, peer, group=None):
1044 # torch's P2POp takes the op callable (dist.isend / dist.irecv), not
1045 # the "isend"/"irecv" string the stage specs builders emit.
1046 if op_type == "isend":
1047 op = dist.isend
1048 elif op_type == "irecv":
1049 op = dist.irecv
1050 else:
1051 raise ValueError(
1052 f"p2p_op op_type must be 'isend' or 'irecv', but got {op_type!r}."
1053 )
1054 return dist.P2POp(op, tensor, peer, group)
1056 @staticmethod
1057 def batch_isend_irecv(p2p_ops):
1058 """Launch a peer-batched P2P group as one coalesced op.
1060 ``torch.distributed.batch_isend_irecv`` coalesces the ops onto one
1061 comm stream and returns one ``Work`` per op; we wrap them in a single
1062 ``.wait()`` handle so a send and a recv to the same peer overlap on
1063 the duplex link and the caller can defer the whole batch's wait to one
1064 consumption point.
1065 """
1066 if not p2p_ops:
1067 return None
1068 works = dist.batch_isend_irecv(p2p_ops)
1069 return _TorchBatchP2PWork(works) if works else None
1071 @staticmethod
1072 def prepare_batch_p2p_group(group: Any = None) -> None:
1073 """Synchronize a group before its first subset batched P2P call.
1075 PyTorch requires every rank in a process group to participate when
1076 ``batch_isend_irecv`` is the first collective on that group. A barrier
1077 at the common pipeline run boundary initializes the communicator
1078 before ranks reach peer operations at different times.
1080 Args:
1081 group: The process group used by the batched P2P operations.
1082 ``None`` uses the default group.
1083 """
1084 dist.barrier(group=group)
1086 @staticmethod
1087 def p2p_exchange(tensor, peer_rank: int, group=None):
1088 if peer_rank == dist.get_rank(group):
1089 return tensor
1090 return _TorchP2PExchangeFunction.apply(tensor, peer_rank, group)
1092 @staticmethod
1093 def send_object_list(obj_list, dst=None, group=None):
1094 dist.send_object_list(obj_list, dst, group)
1096 @staticmethod
1097 def recv_object_list(obj_list, src=None, group=None):
1098 dist.recv_object_list(obj_list, src, group)
1100 @staticmethod
1101 def reduce_scatter_tensor(data, group_info, async_op=False):
1102 output_shape = list(data.shape)
1103 output_shape[0] = output_shape[0] // group_info.rank_size
1104 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
1105 handle = dist.reduce_scatter_tensor(output, data, group=group_info.group, async_op=async_op)
1106 return output, handle
1108 @staticmethod
1109 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1110 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
1111 handle = dist.reduce_scatter_tensor(output, input_tensor, group=group, async_op=async_op)
1112 return output, handle
1114 @staticmethod
1115 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1116 output = torch.empty(output_shape, device=input_tensor.device, dtype=input_tensor.dtype)
1117 work = dist.all_to_all_single(output, input_tensor, group=group, async_op=async_op)
1118 return output, work
1120 @staticmethod
1121 def differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group):
1122 """Variable-split all-to-all with autograd support for EP token dispatch/combine."""
1123 out_total = sum(output_splits)
1124 output = torch.empty(
1125 out_total, *input_tensor.shape[1:],
1126 dtype=input_tensor.dtype, device=input_tensor.device,
1127 )
1128 output = dist_func.all_to_all_single(
1129 output, input_tensor,
1130 output_split_sizes=output_splits,
1131 input_split_sizes=input_splits,
1132 group=group,
1133 )
1134 return output
1136 @staticmethod
1137 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1138 """Truly-async variant of :meth:`differentiable_all_to_all_single`.
1140 Both forward AND backward return :class:`AsyncCollectiveTensor`,
1141 so the ``wait_tensor`` op is queued lazily — only when a downstream
1142 kernel actually reads the result.
1144 Why both directions need lazy wait:
1146 * FWD: ACT lazy wait lets host return immediately and the paired
1147 BWD thread's compute kernel slip into the queue before the wait.
1148 * BWD: PyTorch's stock backward issues ``wait_tensor`` eagerly,
1149 and the autograd engine binds backward stream to the forward
1150 stream — so even running BWD inside a ``with torch.npu.stream
1151 (side_stream)`` context does not move that wait off the main
1152 stream. Returning ACT from backward defers the wait to the
1153 next backward op's first consumption, opening a small window
1154 during which FWD's Attention kernels can be queued onto the
1155 main stream **before** the wait lands.
1157 Args:
1158 input_tensor: Input tensor, split along dim 0 by ``input_splits``.
1159 input_splits: ``list[int]`` — rows sent to each rank.
1160 output_splits: ``list[int]`` — rows received from each rank.
1161 group: Process group.
1163 Returns:
1164 ``AsyncCollectiveTensor`` of shape
1165 ``[sum(output_splits), *input_tensor.shape[1:]]``.
1166 """
1167 return _AsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1169 @staticmethod
1170 def wait_async_tensor(tensor):
1171 """Wait for an async collective tensor to become materialised.
1173 Idempotent — calling on an already-waited tensor is a no-op.
1175 Args:
1176 tensor: ``AsyncCollectiveTensor`` whose device-side values may
1177 not yet be ready.
1179 Returns:
1180 The same *tensor*, now fully materialised.
1181 """
1182 from torch.distributed._functional_collectives import wait_tensor # pylint: disable=C0415
1183 wait_tensor(tensor)
1184 return tensor
1186 @staticmethod
1187 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1188 handle_box=None):
1189 """Wait async all-gather handle and reconstruct result (differentiable)."""
1190 return _TorchAsyncAllGatherFunction.apply(
1191 x, work, out_perm, group, world_size, gather_dim, handle_box
1192 )
1194 @staticmethod
1195 def arange(start, end=None, step=1, dtype=None, device=None):
1196 """Create a 1-D tensor with evenly spaced values."""
1197 if end is None:
1198 return torch.arange(start, dtype=dtype, device=device)
1199 return torch.arange(start, end, step, dtype=dtype, device=device)
1201 @staticmethod
1202 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim,
1203 handle_box=None):
1204 """Wait async A2A handle and reconstruct result (differentiable).
1206 Args:
1207 x: Input tensor.
1208 work: Async work handle from all_to_all.
1209 out_perm: Output buffer from all_to_all.
1210 group: Process group.
1211 world_size: World size.
1212 concat_dim: Dimension for concatenation.
1213 split_dim: Dimension for split.
1214 handle_box: Optional mutable list; backward appends (work, out_perm) here.
1215 """
1216 return _TorchAsyncA2AFunction.apply(
1217 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1218 )
1220 @staticmethod
1221 def differentiable_sync_hook(x, hook_name: str, coordinator):
1222 """Identity op that fires coordinator rendezvous on forward and backward.
1224 Always goes through ``_TorchSyncHookFunction.apply`` so that the
1225 autograd graph **records a SyncHook node regardless of whether the
1226 coordinator is currently enabled**. Skipping ``apply`` when
1227 disabled would leave warmup-forwarded graphs without the hook
1228 nodes, and a later ``overlap.run`` — whose BWD thread back-props
1229 such a graph — would then traverse zero hooks while the paired FWD
1230 thread (whose current forward DOES record hooks) waits at a
1231 barrier for a partner that never arrives.
1233 Args:
1234 x: Input tensor.
1235 hook_name: One of:
1236 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1237 full rendezvous on both directions.
1238 * ``"D_LAST"`` — closing D of the last MoE
1239 layer in a chunk. Forward: ``notify_dispatched``
1240 only (no Attention follows so rendezvous is
1241 skipped). Backward: pure skip (first BWD
1242 hook to fire; combine.bwd has already
1243 dispatched freely).
1244 coordinator: A :class:`HookCoordinator` instance.
1245 """
1246 return _TorchSyncHookFunction.apply(x, hook_name, coordinator)
1248 @staticmethod
1249 def get_tensor_transform():
1250 raise NotImplementedError("Unsupported get_tensor_transform for torch platform")
1252 @staticmethod
1253 def construct_strided_slice(x, begin, end, stride):
1254 raise NotImplementedError("Unsupported construct_strided_slice for torch platform")
1256 @staticmethod
1257 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1258 # pylint: disable=C0415
1259 from hyper_parallel.platform.torch.pipeline_parallel._utils import _MicroBatch
1260 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1262 @staticmethod
1263 def get_symmetric_memory_handler():
1264 # pylint: disable=C0415
1265 from hyper_parallel.platform.torch.symmetric_memory import TorchSymmetricMemoryHandler
1266 symmetric_memory = TorchSymmetricMemoryHandler()
1267 return symmetric_memory
1269 @staticmethod
1270 def get_multicore_handler():
1271 """Return a TorchMulticoreHandler instance for multi-core device management."""
1272 # pylint: disable=C0415
1273 from hyper_parallel.platform.torch.multicore import TorchMulticoreHandler
1274 return TorchMulticoreHandler()
1276 def new_stream(self):
1277 device = self.get_device_handle()
1278 return device.Stream()
1280 def get_stream_context(self):
1281 device = self.get_device_handle()
1282 return device.stream
1284 @staticmethod
1285 def all_gather_object(object_list, obj, group=None) -> None:
1286 """
1287 Gathers objects from the given group into object list.
1289 Args:
1290 object_list (list[Any]): Define the output list, which size equal to the size of group.
1291 obj (Any): The object on current rank and in given process group.
1292 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1293 global group.
1295 Returns:
1296 None. Objs are gathered into ``object_list``.
1297 """
1298 dist.all_gather_object(object_list, obj, group)
1300 @staticmethod
1301 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1302 """
1303 Synchronize all processes in the given process group.
1305 Args:
1306 group (ProcessGroup, optional): The process group to work on. Default is ``None``,
1307 meaning the default process group.
1308 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1309 device_ids (list[int], optional): Device ids for backends that require a device for
1310 barrier (e.g. NCCL). Default: ``None``.
1312 Returns:
1313 Async work handle if ``async_op`` is True; otherwise ``None``.
1314 """
1315 return dist.barrier(group, async_op, device_ids)
1317 @staticmethod
1318 def init_process_group(
1319 backend: Optional[str] = None,
1320 *,
1321 init_method: Optional[str] = None,
1322 timeout: Optional[timedelta] = None,
1323 world_size: int = -1,
1324 rank: int = -1,
1325 store: Optional[Store] = None,
1326 pg_options: Optional[Any] = None,
1327 device_id: Optional[Union[torch.device, int]] = None,
1328 ) -> None:
1329 """
1330 Initialize global process group.
1332 Args:
1333 backend (str or Backend, optional): The backend to use for distributed communication.
1334 init_method (str, optional): URL specifying how to initialize the process group. Default is "env://",
1335 can not be specified at the same time with ``store``.
1336 timeout (timedelta, optional): Timeout for process group. Default 10 minutes for NCCL and for other
1337 backends 30 minutes.
1338 world_size (int, optional): Number of processes. If ``store`` is specified, world_size is required.
1339 rank (int, optional): Rank of the current process, which value must between 0 and ``world_size``-1. If
1340 ``store`` is specified, rank is required.
1341 store (Store, optional): Key/value store accessible to all workers, used to exchange connection/address
1342 information. Can not be specified at the same time with ``init_method``.
1343 pg_options (ProcessGroupOptions, optional): Extra options to pass during constructing process groups.
1344 device_id (torch.device | int, optional): Specific device this process will work on.
1345 """
1346 try:
1347 _get_default_group()
1348 # except multi version error
1349 except (ValueError, RuntimeError):
1350 if backend is None:
1351 backend = "hccl"
1352 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1353 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1355 @staticmethod
1356 def destroy_process_group(group: Optional[ProcessGroup] = None) -> None:
1357 """
1358 Destroy given process group.
1360 Args:
1361 group (ProcessGroup, optional): Given process group will be destroyed, if not given, all process groups
1362 will be destroyed.
1363 """
1364 group = group or _get_default_group()
1365 if group in EXISTING_COMM_GROUPS.values():
1366 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1367 for k in keys_to_destroy:
1368 del EXISTING_COMM_GROUPS[k]
1369 dist.destroy_process_group(group)
1371 @staticmethod
1372 def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> list[int]:
1373 """
1374 Get all ranks relative to given process group.
1376 Args:
1377 group (Optional[ProcessGroup]): Process group worked on. Default is ``None``, and ``None`` means global
1378 group.
1380 Returns:
1381 Rank list.
1382 """
1383 group = group or _get_default_group()
1384 return dist.get_process_group_ranks(group)
1386 @staticmethod
1387 def get_backend(group: Optional[ProcessGroup] = None) -> Backend:
1388 """
1389 Get the backend of the given process group.
1391 Args:
1392 group (ProcessGroup, optional): Process group worked on. Default is ``None``, and ``None`` means global
1393 group.
1395 Returns:
1396 The backend object of the given process group.
1397 """
1398 group = group or _get_default_group()
1399 return dist.get_backend(group)
1401 @staticmethod
1402 def split_group(parent_pg: Optional[ProcessGroup] = None,
1403 split_ranks: Optional[list] = None,
1404 timeout: Optional[timedelta] = None,
1405 pg_options: Optional[Any] = None,
1406 group_desc: Optional[str] = None,
1407 ) -> Optional[ProcessGroup]:
1408 """
1409 Create split groups for every group rank in split_ranks, and return the split process group which relative to
1410 current rank id.
1412 Args:
1413 parent_pg (Optional[ProcessGroup]): A process group which the goal group split from.
1414 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1415 timeout (Optional[timedelta]): Timeout for process group. Default 10 minutes for NCCL and for other
1416 backend 30 minutes.
1417 pg_options (Optional[Any]): Extra options to pass during constructing process groups.
1418 group_desc (Optional[str]): Description of process group.
1420 Return:
1421 Optional[ProcessGroup]: One of split process group which relative to current rank id
1422 """
1423 if split_ranks is None or len(split_ranks) == 0:
1424 raise ValueError("split_ranks cannot be None or empty")
1426 split_group = None
1427 for split_rank in split_ranks:
1428 dist_group = TorchPlatform.get_created_group(split_rank)
1429 if dist_group is None:
1430 dist_group = dist.new_group(ranks=split_rank)
1431 EXISTING_COMM_GROUPS[str(tuple(sorted(split_rank)))] = dist_group
1432 if TorchPlatform.get_rank() in split_rank:
1433 split_group = dist_group
1435 return split_group
1437 @staticmethod
1438 def get_group_local_rank(group: ProcessGroup = None) -> int:
1439 """get group local rank id."""
1440 group = group or _get_default_group()
1441 return group.rank()
1443 @staticmethod
1444 def no_grad():
1445 return torch.no_grad()
1447 @staticmethod
1448 def preserve_version_counter(tensor):
1449 return torch.autograd._unsafe_preserve_version_counter(tensor) # pylint: disable=W0212
1451 @staticmethod
1452 def relu(tensor):
1453 return torch.relu(tensor)
1455 @staticmethod
1456 def cat(tensors, dim=0):
1457 return torch.cat(tensors, dim=dim)
1459 @staticmethod
1460 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1461 return torch.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1463 def get_current_stream(self):
1464 device = self.get_device_handle()
1465 return device.current_stream()
1467 def new_event(self):
1468 device = self.get_device_handle()
1469 return device.Event()
1471 def tree_map(self, fn, tree):
1472 return torch.utils._pytree.tree_map(fn, tree) # pylint: disable=protected-access
1474 @property
1475 def checkpoint(self):
1476 return torch.utils.checkpoint.checkpoint
1478 @staticmethod
1479 def checkpoint_wrapper(module, **checkpoint_kwargs):
1480 # pylint: disable=C0415
1481 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1482 return ckpt_wrapper(module, **checkpoint_kwargs)
1484 @staticmethod
1485 def swap_wrapper(module, policy_fn=None, group_swap=False):
1486 # pylint: disable=C0415
1487 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_wrapper
1488 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1490 @staticmethod
1491 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1492 # pylint: disable=C0415
1493 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_tensor_wrapper
1494 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1496 @staticmethod
1497 def get_class_activation_wrapper():
1498 # pylint: disable=C0415
1499 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import ActivationWrapper
1500 return ActivationWrapper
1502 @property
1503 def noop_context_fn(self):
1504 return noop_context_fn
1506 @staticmethod
1507 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1508 # pylint: disable=C0415
1509 from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts
1510 return create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation, group_swap)
1512 @staticmethod
1513 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1514 # pylint: disable=C0415
1515 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import AsyncSaveOnCpu
1516 return AsyncSaveOnCpu(policy_fn, group_swap=group_swap)
1518 @staticmethod
1519 def get_element_size(tensor):
1520 """Get Tensor Element Size"""
1521 return tensor.element_size()
1523 @staticmethod
1524 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1525 """Allocate an uninitialized 1-D tensor buffer."""
1526 if pin_memory:
1527 return torch.empty(numel, dtype=dtype, device='cpu', pin_memory=True)
1528 return torch.empty(numel, dtype=dtype, device=device)
1530 @staticmethod
1531 def tensor_to_numpy(tensor) -> np.ndarray:
1532 """Convert PyTorch tensor to numpy array."""
1533 return tensor.cpu().numpy()
1535 @staticmethod
1536 def from_numpy(np_array):
1537 """Create a host (CPU) PyTorch tensor from a numpy array."""
1538 return torch.from_numpy(np_array)
1540 @staticmethod
1541 def clip_grad_norm_(
1542 parameters, max_norm, norm_type=2.0,
1543 error_if_nonfinite=False, foreach=None,
1544 ):
1545 # pylint: disable=C0415
1546 from hyper_parallel.platform.torch.clip_grad import (
1547 clip_grad_norm_ as _clip_grad_norm,
1548 )
1549 return _clip_grad_norm(
1550 parameters, max_norm, norm_type,
1551 error_if_nonfinite=error_if_nonfinite, foreach=foreach,
1552 )
1554 @staticmethod
1555 def profiler_record(name):
1556 """Profiler context manager for recording operations using torch.profiler."""
1557 return torch.profiler.record_function(name)
1559 def cast_fp_tensor(self, dtype, x):
1560 """
1561 Cast floating-point tensor to target dtype if applicable.
1562 """
1563 if (
1564 not isinstance(x, torch.Tensor)
1565 or not torch.is_floating_point(x)
1566 or x.dtype == dtype
1567 ):
1568 return x
1569 return x.to(dtype)
1571 def apply_to_tensors(self, fn, container):
1572 """Recursively apply to all tensor in different kinds of container types."""
1574 def apply(x):
1576 if isinstance(x, torch.Tensor):
1577 return fn(x)
1578 if hasattr(x, "__dataclass_fields__"):
1579 dc = dataclasses.replace(x)
1580 changes = {
1581 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
1582 }
1583 return dataclasses.replace(dc, **changes)
1584 if isinstance(x, OrderedDict):
1585 od = x.__class__()
1586 for key, value in x.items():
1587 od[key] = apply(value)
1588 return od
1589 if isinstance(x, PackedSequence):
1590 apply(x.data)
1591 return x
1592 if isinstance(x, dict):
1593 return {key: apply(value) for key, value in x.items()}
1594 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
1595 res = (apply(el) for el in x)
1596 return type(x)(*res)
1597 if isinstance(x, (list, tuple, set)):
1598 return type(x)(apply(el) for el in x)
1599 return x
1601 return apply(container)
1604 @property
1605 def meta_device(self):
1606 return torch.device("meta")
1608 def init_on_device(self, device, include_buffers=False):
1609 return _init_on_device(device, include_buffers=include_buffers)
1611 def str_to_dtype(self, dtype_str: str) -> torch.dtype:
1612 """Map ``torch.<type>`` strings from checkpoint metadata to ``torch.dtype``."""
1613 parts = dtype_str.split(".", 1)
1614 if len(parts) != 2:
1615 raise ValueError(
1616 f"Expected dtype string like 'torch.float32', got {dtype_str!r}."
1617 )
1618 prefix, name = parts
1619 if prefix != "torch":
1620 raise ValueError(
1621 f"Expected PyTorch dtype string with prefix 'torch', got {dtype_str!r}."
1622 )
1623 dtype = getattr(torch, name)
1624 if isinstance(dtype, torch.dtype):
1625 return dtype
1626 raise ValueError(f"{dtype_str!r} does not resolve to a torch.dtype.")
1628 def list_to_size(self, size_list: list[int]) -> torch.Size:
1629 return torch.Size(size_list)