Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / platform.py: 61%
747 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +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 get_rank():
625 """
626 Get the rank of the current process in the distributed group.
628 Returns:
629 int: The rank of the current process.
630 """
631 return dist.get_rank()
633 @staticmethod
634 def get_global_rank(group, group_rank):
635 """
636 Get the global rank from a group rank.
638 Args:
639 group (ProcessGroup): The process group.
640 group_rank (int): The rank within the group.
642 Returns:
643 int: The global rank.
644 """
645 return dist.get_global_rank(group, group_rank)
647 @staticmethod
648 def get_group_rank(group):
649 """Return this process's rank within *group*."""
650 return dist.get_group_rank(group, dist.get_rank())
652 @staticmethod
653 def get_world_size():
654 """
655 Get the total number of processes in the distributed group.
657 Returns:
658 int: The world size.
659 """
660 return dist.get_world_size()
662 @staticmethod
663 def get_param_local_shape(param):
664 """
665 Get the local shape of a parameter, handling both regular and distributed tensors.
667 Args:
668 param (Union[Tensor, DTensorBase]): The parameter tensor.
670 Returns:
671 torch.Size: The local shape of the parameter.
672 """
673 if isinstance(param, DTensorBase):
674 return param.local_shape
675 return param.shape
677 @staticmethod
678 def get_param_local_data(param):
679 """
680 Get the local data of a parameter, handling both regular and distributed tensors.
682 Args:
683 param (Union[Tensor, DTensorBase]): The parameter tensor.
685 Returns:
686 Tensor: The local tensor data.
687 """
688 if isinstance(param, DTensorBase):
689 return param.to_local()
690 return param
692 @staticmethod
693 def update_param_data(param, data):
694 """
695 Update the data of a parameter.
697 Args:
698 param (Parameter): The parameter to update.
699 data (Tensor): The new data tensor.
700 """
701 param.data = data
703 @staticmethod
704 def load_into_param(param, data):
705 """Load tensor *data* into *param* (plain tensor or DTensor)."""
706 if isinstance(param, DTensorBase):
707 local = param._local_tensor # pylint: disable=W0212
708 if local.is_meta:
709 # Meta tensor materialisation: replace the placeholder.
710 orig_requires_grad = param.requires_grad
711 param._local_tensor = data # pylint: disable=W0212
712 if data.requires_grad != orig_requires_grad:
713 param.requires_grad_(orig_requires_grad)
714 else:
715 local.copy_(data)
716 else:
717 param.copy_(data)
719 @staticmethod
720 def get_op_name(func):
721 """
722 Extract the operation name from various function types.
724 Args:
725 func: The function or operation to extract the name from.
727 Returns:
728 str: The operation name.
729 """
730 if hasattr(func, "__name__"):
731 return func.__name__
732 if isinstance(func, OpOverload):
733 full_name = func.name
734 core_name = full_name.split("::")[-1].split(".")[0]
735 return core_name
736 if isinstance(func, OpOverloadPacket):
737 return func.name.split("::")[-1]
738 func_str = str(func)
739 if "built-in function" in func_str:
740 return func_str.split()[-1].strip(">")
741 if "function" in func_str:
742 return func_str.split()[1]
743 return "unknown_op"
745 @staticmethod
746 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
747 data = _ensure_contiguous(data)
748 output = list(dist_func.all_gather(data, group=group))
749 if rank_list is not None:
750 group_ranks = dist.get_process_group_ranks(group)
751 if tuple(rank_list) != tuple(group_ranks):
752 rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}
753 output = [output[rank_to_idx[int(rank)]] for rank in rank_list]
754 return torch.cat(output, dim=concat_dim)
756 @staticmethod
757 def chunk(data, split_dim, split_size, index):
758 return torch.chunk(data, split_size, dim=split_dim)[index]
760 @staticmethod
761 def differentiable_all_to_all(input_data, output_shape, group):
762 input_data = _ensure_contiguous(input_data)
763 output_tensor = torch.empty(output_shape, device=input_data.device, dtype=input_data.dtype)
764 output_tensor = dist_func.all_to_all_single(
765 output_tensor,
766 input_data,
767 group=group
768 )
769 return output_tensor
771 @staticmethod
772 def tensor_type_cast(input_data, cast_type):
773 """Cast tensor to specified data type."""
774 type_mapping = {
775 'float32': torch.float32,
776 'float16': torch.float16,
777 'int64': torch.int64,
778 'int32': torch.int32
779 }
780 if cast_type not in type_mapping:
781 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")
782 return input_data.to(type_mapping[cast_type])
784 @staticmethod
785 def differentiable_all_reduce(data, op, group):
786 data = _ensure_contiguous(data)
787 # Resolve the op from string to ReduceOp enum if necessary
788 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
789 return dist_func.all_reduce(data, op=reduce_op, group=group)
791 @staticmethod
792 def get_cell_construct(cell):
793 return cell.forward
795 @staticmethod
796 def get_cells_and_names(cell):
797 return cell.named_modules()
799 @staticmethod
800 def get_modules(module):
801 return module.modules()
803 @staticmethod
804 def search_parameter_by_name(cell, param_name: str):
805 """
806 Find the parent Module of the parameter, the parameter's name in the parent Module, and the parameter.
807 Return value: (parent Module instance, parameter's name in parent Module, parameter object).
808 Returns None if not found.
809 """
810 # Remove the "self." prefix from param_name
811 param_name = param_name.replace("self.", "")
812 # Case 1: The parameter is a direct parameter of the current Module
813 if param_name in cell._parameters: # pylint: disable=protected-access
814 return (cell, param_name, cell._parameters[param_name]) # pylint: disable=protected-access
816 # Case 2: The parameter is in a sub-Module
817 if "." in param_name:
818 cell_path, param_key = param_name.rsplit(".", 1)
819 try:
820 # Locate the sub-Module where the parameter resides (supports multi-level paths)
821 target_cell = cell.get_submodule(cell_path)
822 # Check if the sub-Module directly contains this parameter
823 if param_key in target_cell._parameters: # pylint: disable=protected-access
824 return target_cell, param_key, target_cell._parameters[param_key] # pylint: disable=protected-access
825 except AttributeError:
826 pass
828 # Traverse all sub-Modules (recursively) to search for the parameter
829 for _, child_cell in cell.named_children():
830 if isinstance(child_cell, Module):
831 result = TorchPlatform.search_parameter_by_name(child_cell, param_name)
832 if result is not None:
833 return result
835 return None
837 @staticmethod
838 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
839 """
840 Modify the original parameter in a Module or sub-Module using the search result
841 """
842 parent_cell, param_key, _ = result
843 # Key operation: directly modify the _parameters dictionary.
844 if param_key in parent_cell._parameters: # pylint: disable=protected-access
845 parent_cell._parameters[param_key] = new_param # pylint: disable=protected-access
846 else:
847 parent_cell.register_parameter(param_key, new_param)
848 return True
850 @staticmethod
851 def set_layout_into_parameter(param, layout):
852 """Set layout into parameter"""
853 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel
854 from hyper_parallel.core.dtensor.layout import _get_slice_tensor_by_layout # pylint: disable=import-outside-toplevel
855 if isinstance(param, DTensor):
856 raise ValueError(f"Parameter {param} has been configured layout, cannot be set repeatedly.")
857 requires_grad = param.requires_grad
858 param_dtensor = DTensor.from_local(
859 _get_slice_tensor_by_layout(param, layout),
860 layout.mesh, layout.alias_placements)
861 new_param = Parameter(param_dtensor, requires_grad=requires_grad)
862 return new_param
864 @staticmethod
865 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
866 data = _ensure_contiguous(data)
867 input_tuple = torch.chunk(data, dev_num, dim=axis)
868 output_tensor = torch.empty(input_tuple[0].shape, device=data.device, dtype=data.dtype)
870 # Resolve the op from string to ReduceOp enum
871 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
873 output_tensor = dist_func.reduce_scatter(output_tensor, input_tuple, op=reduce_op, group=group)
875 # Keep manual handling for 'avg' string as it maps to SUM in _OP_MAP
876 if op == 'avg':
877 output_tensor = output_tensor / dev_num
878 return output_tensor
880 @staticmethod
881 def get_device_handle(device_type: str = "npu"):
882 """Return the torch device module (e.g. ``torch.npu`` or ``torch.cuda``) for the given device type."""
883 try:
884 handle = getattr(torch, device_type)
885 except AttributeError as e:
886 raise RuntimeError(f"TorchPlatform expect got device handle: 'torch.{device_type}' failed.") from e
887 return handle
889 @staticmethod
890 def get_param_type_size(param):
891 # pylint: disable=W0212
892 return torch._utils._element_size(param.dtype)
894 @staticmethod
895 def is_tensor(obj: Any) -> bool:
896 """Return True if ``obj`` is a ``torch.Tensor``."""
897 return isinstance(obj, Tensor)
899 @staticmethod
900 def get_tensor_storage_size(tensor: Any) -> int:
901 """Return serialized byte size (numel * element size) for a PyTorch tensor."""
902 if not TorchPlatform.is_tensor(tensor):
903 raise TypeError(
904 f"TorchPlatform.get_tensor_storage_size expects torch.Tensor, got {type(tensor)!r}"
905 )
906 return int(tensor.numel()) * int(tensor.element_size())
908 @staticmethod
909 def parameters_dict(cell: Module):
910 return cell.named_parameters()
912 @staticmethod
913 def get_model_state_dict(model, *, options=None):
914 # pylint: disable=C0415
915 from hyper_parallel.platform.torch.fully_shard.state_dict_utils import (
916 get_model_state_dict as _get_model_state_dict,
917 )
918 return _get_model_state_dict(model, options=options)
920 @staticmethod
921 def save_checkpoint(cell: Module, file_path: str, ckpt_format: str = "safetensors") -> None:
922 if ckpt_format == "safetensors":
923 save_file(tensors=cell, filename=file_path)
924 else:
925 torch.save(obj=cell, f=file_path)
927 @staticmethod
928 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
929 if ckpt_format == "safetensors":
930 return load_file(filename=file_path)
931 return torch.load(f=file_path)
933 @staticmethod
934 def new_zero_parameter(param_shape, param_type, requires_grad, device):
935 return nn.Parameter(torch.zeros(param_shape, dtype=param_type, device=device), requires_grad=requires_grad)
937 @staticmethod
938 def new_tensor(tensor_shape, tensor_type, device):
939 return torch.empty(size=tensor_shape, dtype=tensor_type, device=device)
941 @staticmethod
942 def full_like(tensor, fill_value, dtype=None):
943 return torch.full_like(tensor, fill_value, dtype=dtype)
945 @staticmethod
946 def set_tensor_requires_grad(input_tensor):
947 """
948 set requires grad flag for input tensor, only effective for leaf node
949 """
950 if input_tensor.is_leaf:
951 input_tensor.requires_grad = True
953 def _create_group(self, rank_list):
954 normalized_rank_list = tuple(sorted(rank_list))
955 world_rank_list = tuple(range(self.get_world_size()))
956 if normalized_rank_list == world_rank_list:
957 group = _get_default_group()
958 EXISTING_COMM_GROUPS[str(normalized_rank_list)] = group
959 return group
960 group_dict = create_sub_groups(rank_list)
961 return group_dict[normalized_rank_list]
963 @staticmethod
964 def all_gather_into_tensor(data, group_info, async_op=False):
965 output_shape = list(data.shape)
966 output_shape[0] = output_shape[0] * group_info.rank_size
967 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
968 handle = dist.all_gather_into_tensor(output, data, group=group_info.group, async_op=async_op)
969 return output, handle
971 @staticmethod
972 def all_gather_single(input_tensor, output_shape, group, async_op=False):
973 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
974 handle = dist.all_gather_into_tensor(output, input_tensor, group=group, async_op=async_op)
975 return output, handle
977 @staticmethod
978 def all_reduce(data, group_info, async_op=False):
979 if not data.is_contiguous():
980 data = data.contiguous()
981 handle = dist.all_reduce(data, group=group_info.group, async_op=async_op)
982 return data, handle
984 @staticmethod
985 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
986 if group_src is not None:
987 src = dist.get_global_rank(group, group_src)
988 handle = dist.broadcast(data, src, group, async_op)
989 if async_op and handle is not None:
990 handle.wait()
992 @staticmethod
993 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
994 if group_src is not None:
995 src = dist.get_global_rank(group, group_src)
996 handle = dist.scatter(output, scatter_list, src=src, group=group, async_op=async_op)
997 if async_op and handle is not None:
998 handle.wait()
999 return output
1001 @staticmethod
1002 def isend(tensor, dst=None, group=None, tag=0):
1003 return dist.isend(tensor, dst, group, tag)
1005 @staticmethod
1006 def irecv(tensor, src=None, group=None, tag=0):
1007 return dist.irecv(tensor, src, group, tag)
1009 @staticmethod
1010 def p2p_op(op_type, tensor, peer, group=None):
1011 # torch's P2POp takes the op callable (dist.isend / dist.irecv), not
1012 # the "isend"/"irecv" string the stage specs builders emit.
1013 if op_type == "isend":
1014 op = dist.isend
1015 elif op_type == "irecv":
1016 op = dist.irecv
1017 else:
1018 raise ValueError(
1019 f"p2p_op op_type must be 'isend' or 'irecv', but got {op_type!r}."
1020 )
1021 return dist.P2POp(op, tensor, peer, group)
1023 @staticmethod
1024 def batch_isend_irecv(p2p_ops):
1025 """Launch a peer-batched P2P group as one coalesced op.
1027 ``torch.distributed.batch_isend_irecv`` coalesces the ops onto one
1028 comm stream and returns one ``Work`` per op; we wrap them in a single
1029 ``.wait()`` handle so a send and a recv to the same peer overlap on
1030 the duplex link and the caller can defer the whole batch's wait to one
1031 consumption point.
1032 """
1033 if not p2p_ops:
1034 return None
1035 works = dist.batch_isend_irecv(p2p_ops)
1036 return _TorchBatchP2PWork(works) if works else None
1038 @staticmethod
1039 def p2p_exchange(tensor, peer_rank: int, group=None):
1040 if peer_rank == dist.get_rank(group):
1041 return tensor
1042 return _TorchP2PExchangeFunction.apply(tensor, peer_rank, group)
1044 @staticmethod
1045 def send_object_list(obj_list, dst=None, group=None):
1046 dist.send_object_list(obj_list, dst, group)
1048 @staticmethod
1049 def recv_object_list(obj_list, src=None, group=None):
1050 dist.recv_object_list(obj_list, src, group)
1052 @staticmethod
1053 def reduce_scatter_tensor(data, group_info, async_op=False):
1054 output_shape = list(data.shape)
1055 output_shape[0] = output_shape[0] // group_info.rank_size
1056 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
1057 handle = dist.reduce_scatter_tensor(output, data, group=group_info.group, async_op=async_op)
1058 return output, handle
1060 @staticmethod
1061 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1062 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
1063 handle = dist.reduce_scatter_tensor(output, input_tensor, group=group, async_op=async_op)
1064 return output, handle
1066 @staticmethod
1067 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1068 output = torch.empty(output_shape, device=input_tensor.device, dtype=input_tensor.dtype)
1069 work = dist.all_to_all_single(output, input_tensor, group=group, async_op=async_op)
1070 return output, work
1072 @staticmethod
1073 def differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group):
1074 """Variable-split all-to-all with autograd support for EP token dispatch/combine."""
1075 out_total = sum(output_splits)
1076 output = torch.empty(
1077 out_total, *input_tensor.shape[1:],
1078 dtype=input_tensor.dtype, device=input_tensor.device,
1079 )
1080 output = dist_func.all_to_all_single(
1081 output, input_tensor,
1082 output_split_sizes=output_splits,
1083 input_split_sizes=input_splits,
1084 group=group,
1085 )
1086 return output
1088 @staticmethod
1089 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1090 """Truly-async variant of :meth:`differentiable_all_to_all_single`.
1092 Both forward AND backward return :class:`AsyncCollectiveTensor`,
1093 so the ``wait_tensor`` op is queued lazily — only when a downstream
1094 kernel actually reads the result.
1096 Why both directions need lazy wait:
1098 * FWD: ACT lazy wait lets host return immediately and the paired
1099 BWD thread's compute kernel slip into the queue before the wait.
1100 * BWD: PyTorch's stock backward issues ``wait_tensor`` eagerly,
1101 and the autograd engine binds backward stream to the forward
1102 stream — so even running BWD inside a ``with torch.npu.stream
1103 (side_stream)`` context does not move that wait off the main
1104 stream. Returning ACT from backward defers the wait to the
1105 next backward op's first consumption, opening a small window
1106 during which FWD's Attention kernels can be queued onto the
1107 main stream **before** the wait lands.
1109 Args:
1110 input_tensor: Input tensor, split along dim 0 by ``input_splits``.
1111 input_splits: ``list[int]`` — rows sent to each rank.
1112 output_splits: ``list[int]`` — rows received from each rank.
1113 group: Process group.
1115 Returns:
1116 ``AsyncCollectiveTensor`` of shape
1117 ``[sum(output_splits), *input_tensor.shape[1:]]``.
1118 """
1119 return _AsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1121 @staticmethod
1122 def wait_async_tensor(tensor):
1123 """Wait for an async collective tensor to become materialised.
1125 Idempotent — calling on an already-waited tensor is a no-op.
1127 Args:
1128 tensor: ``AsyncCollectiveTensor`` whose device-side values may
1129 not yet be ready.
1131 Returns:
1132 The same *tensor*, now fully materialised.
1133 """
1134 from torch.distributed._functional_collectives import wait_tensor # pylint: disable=C0415
1135 wait_tensor(tensor)
1136 return tensor
1138 @staticmethod
1139 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1140 handle_box=None):
1141 """Wait async all-gather handle and reconstruct result (differentiable)."""
1142 return _TorchAsyncAllGatherFunction.apply(
1143 x, work, out_perm, group, world_size, gather_dim, handle_box
1144 )
1146 @staticmethod
1147 def arange(start, end=None, step=1, dtype=None, device=None):
1148 """Create a 1-D tensor with evenly spaced values."""
1149 if end is None:
1150 return torch.arange(start, dtype=dtype, device=device)
1151 return torch.arange(start, end, step, dtype=dtype, device=device)
1153 @staticmethod
1154 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim,
1155 handle_box=None):
1156 """Wait async A2A handle and reconstruct result (differentiable).
1158 Args:
1159 x: Input tensor.
1160 work: Async work handle from all_to_all.
1161 out_perm: Output buffer from all_to_all.
1162 group: Process group.
1163 world_size: World size.
1164 concat_dim: Dimension for concatenation.
1165 split_dim: Dimension for split.
1166 handle_box: Optional mutable list; backward appends (work, out_perm) here.
1167 """
1168 return _TorchAsyncA2AFunction.apply(
1169 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1170 )
1172 @staticmethod
1173 def differentiable_sync_hook(x, hook_name: str, coordinator):
1174 """Identity op that fires coordinator rendezvous on forward and backward.
1176 Always goes through ``_TorchSyncHookFunction.apply`` so that the
1177 autograd graph **records a SyncHook node regardless of whether the
1178 coordinator is currently enabled**. Skipping ``apply`` when
1179 disabled would leave warmup-forwarded graphs without the hook
1180 nodes, and a later ``overlap.run`` — whose BWD thread back-props
1181 such a graph — would then traverse zero hooks while the paired FWD
1182 thread (whose current forward DOES record hooks) waits at a
1183 barrier for a partner that never arrives.
1185 Args:
1186 x: Input tensor.
1187 hook_name: One of:
1188 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1189 full rendezvous on both directions.
1190 * ``"D_LAST"`` — closing D of the last MoE
1191 layer in a chunk. Forward: ``notify_dispatched``
1192 only (no Attention follows so rendezvous is
1193 skipped). Backward: pure skip (first BWD
1194 hook to fire; combine.bwd has already
1195 dispatched freely).
1196 coordinator: A :class:`HookCoordinator` instance.
1197 """
1198 return _TorchSyncHookFunction.apply(x, hook_name, coordinator)
1200 @staticmethod
1201 def get_tensor_transform():
1202 raise NotImplementedError("Unsupported get_tensor_transform for torch platform")
1204 @staticmethod
1205 def construct_strided_slice(x, begin, end, stride):
1206 raise NotImplementedError("Unsupported construct_strided_slice for torch platform")
1208 @staticmethod
1209 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1210 # pylint: disable=C0415
1211 from hyper_parallel.platform.torch.pipeline_parallel._utils import _MicroBatch
1212 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1214 @staticmethod
1215 def get_symmetric_memory_handler():
1216 # pylint: disable=C0415
1217 from hyper_parallel.platform.torch.symmetric_memory import TorchSymmetricMemoryHandler
1218 symmetric_memory = TorchSymmetricMemoryHandler()
1219 return symmetric_memory
1221 @staticmethod
1222 def get_multicore_handler():
1223 """Return a TorchMulticoreHandler instance for multi-core device management."""
1224 # pylint: disable=C0415
1225 from hyper_parallel.platform.torch.multicore import TorchMulticoreHandler
1226 return TorchMulticoreHandler()
1228 def new_stream(self):
1229 device = self.get_device_handle()
1230 return device.Stream()
1232 def get_stream_context(self):
1233 device = self.get_device_handle()
1234 return device.stream
1236 @staticmethod
1237 def all_gather_object(object_list, obj, group=None) -> None:
1238 """
1239 Gathers objects from the given group into object list.
1241 Args:
1242 object_list (list[Any]): Define the output list, which size equal to the size of group.
1243 obj (Any): The object on current rank and in given process group.
1244 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1245 global group.
1247 Returns:
1248 None. Objs are gathered into ``object_list``.
1249 """
1250 dist.all_gather_object(object_list, obj, group)
1252 @staticmethod
1253 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1254 """
1255 Synchronize all processes in the given process group.
1257 Args:
1258 group (ProcessGroup, optional): The process group to work on. Default is ``None``,
1259 meaning the default process group.
1260 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1261 device_ids (list[int], optional): Device ids for backends that require a device for
1262 barrier (e.g. NCCL). Default: ``None``.
1264 Returns:
1265 Async work handle if ``async_op`` is True; otherwise ``None``.
1266 """
1267 return dist.barrier(group, async_op, device_ids)
1269 @staticmethod
1270 def init_process_group(
1271 backend: Optional[str] = None,
1272 *,
1273 init_method: Optional[str] = None,
1274 timeout: Optional[timedelta] = None,
1275 world_size: int = -1,
1276 rank: int = -1,
1277 store: Optional[Store] = None,
1278 pg_options: Optional[Any] = None,
1279 device_id: Optional[Union[torch.device, int]] = None,
1280 ) -> None:
1281 """
1282 Initialize global process group.
1284 Args:
1285 backend (str or Backend, optional): The backend to use for distributed communication.
1286 init_method (str, optional): URL specifying how to initialize the process group. Default is "env://",
1287 can not be specified at the same time with ``store``.
1288 timeout (timedelta, optional): Timeout for process group. Default 10 minutes for NCCL and for other
1289 backends 30 minutes.
1290 world_size (int, optional): Number of processes. If ``store`` is specified, world_size is required.
1291 rank (int, optional): Rank of the current process, which value must between 0 and ``world_size``-1. If
1292 ``store`` is specified, rank is required.
1293 store (Store, optional): Key/value store accessible to all workers, used to exchange connection/address
1294 information. Can not be specified at the same time with ``init_method``.
1295 pg_options (ProcessGroupOptions, optional): Extra options to pass during constructing process groups.
1296 device_id (torch.device | int, optional): Specific device this process will work on.
1297 """
1298 try:
1299 _get_default_group()
1300 # except multi version error
1301 except (ValueError, RuntimeError):
1302 if backend is None:
1303 backend = "hccl"
1304 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1305 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1307 @staticmethod
1308 def destroy_process_group(group: Optional[ProcessGroup] = None) -> None:
1309 """
1310 Destroy given process group.
1312 Args:
1313 group (ProcessGroup, optional): Given process group will be destroyed, if not given, all process groups
1314 will be destroyed.
1315 """
1316 group = group or _get_default_group()
1317 if group in EXISTING_COMM_GROUPS.values():
1318 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1319 for k in keys_to_destroy:
1320 del EXISTING_COMM_GROUPS[k]
1321 dist.destroy_process_group(group)
1323 @staticmethod
1324 def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> list[int]:
1325 """
1326 Get all ranks relative to given process group.
1328 Args:
1329 group (Optional[ProcessGroup]): Process group worked on. Default is ``None``, and ``None`` means global
1330 group.
1332 Returns:
1333 Rank list.
1334 """
1335 group = group or _get_default_group()
1336 return dist.get_process_group_ranks(group)
1338 @staticmethod
1339 def get_backend(group: Optional[ProcessGroup] = None) -> Backend:
1340 """
1341 Get the backend of the given process group.
1343 Args:
1344 group (ProcessGroup, optional): Process group worked on. Default is ``None``, and ``None`` means global
1345 group.
1347 Returns:
1348 The backend object of the given process group.
1349 """
1350 group = group or _get_default_group()
1351 return dist.get_backend(group)
1353 @staticmethod
1354 def split_group(parent_pg: Optional[ProcessGroup] = None,
1355 split_ranks: Optional[list] = None,
1356 timeout: Optional[timedelta] = None,
1357 pg_options: Optional[Any] = None,
1358 group_desc: Optional[str] = None,
1359 ) -> Optional[ProcessGroup]:
1360 """
1361 Create split groups for every group rank in split_ranks, and return the split process group which relative to
1362 current rank id.
1364 Args:
1365 parent_pg (Optional[ProcessGroup]): A process group which the goal group split from.
1366 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1367 timeout (Optional[timedelta]): Timeout for process group. Default 10 minutes for NCCL and for other
1368 backend 30 minutes.
1369 pg_options (Optional[Any]): Extra options to pass during constructing process groups.
1370 group_desc (Optional[str]): Description of process group.
1372 Return:
1373 Optional[ProcessGroup]: One of split process group which relative to current rank id
1374 """
1375 if split_ranks is None or len(split_ranks) == 0:
1376 raise ValueError("split_ranks cannot be None or empty")
1378 split_group = None
1379 for split_rank in split_ranks:
1380 dist_group = TorchPlatform.get_created_group(split_rank)
1381 if dist_group is None:
1382 dist_group = dist.new_group(ranks=split_rank)
1383 EXISTING_COMM_GROUPS[str(tuple(sorted(split_rank)))] = dist_group
1384 if TorchPlatform.get_rank() in split_rank:
1385 split_group = dist_group
1387 return split_group
1389 @staticmethod
1390 def get_group_local_rank(group: ProcessGroup = None) -> int:
1391 """get group local rank id."""
1392 group = group or _get_default_group()
1393 return group.rank()
1395 @staticmethod
1396 def no_grad():
1397 return torch.no_grad()
1399 @staticmethod
1400 def preserve_version_counter(tensor):
1401 return torch.autograd._unsafe_preserve_version_counter(tensor) # pylint: disable=W0212
1403 @staticmethod
1404 def relu(tensor):
1405 return torch.relu(tensor)
1407 @staticmethod
1408 def cat(tensors, dim=0):
1409 return torch.cat(tensors, dim=dim)
1411 @staticmethod
1412 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1413 return torch.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1415 def get_current_stream(self):
1416 device = self.get_device_handle()
1417 return device.current_stream()
1419 def new_event(self):
1420 device = self.get_device_handle()
1421 return device.Event()
1423 def tree_map(self, fn, tree):
1424 return torch.utils._pytree.tree_map(fn, tree) # pylint: disable=protected-access
1426 @property
1427 def checkpoint(self):
1428 return torch.utils.checkpoint.checkpoint
1430 @staticmethod
1431 def checkpoint_wrapper(module, **checkpoint_kwargs):
1432 # pylint: disable=C0415
1433 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1434 return ckpt_wrapper(module, **checkpoint_kwargs)
1436 @staticmethod
1437 def swap_wrapper(module, policy_fn=None, group_swap=False):
1438 # pylint: disable=C0415
1439 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_wrapper
1440 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1442 @staticmethod
1443 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1444 # pylint: disable=C0415
1445 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_tensor_wrapper
1446 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1448 @staticmethod
1449 def get_class_activation_wrapper():
1450 # pylint: disable=C0415
1451 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import ActivationWrapper
1452 return ActivationWrapper
1454 @property
1455 def noop_context_fn(self):
1456 return noop_context_fn
1458 @staticmethod
1459 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1460 # pylint: disable=C0415
1461 from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts
1462 return create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation, group_swap)
1464 @staticmethod
1465 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1466 # pylint: disable=C0415
1467 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import AsyncSaveOnCpu
1468 return AsyncSaveOnCpu(policy_fn, group_swap=group_swap)
1470 @staticmethod
1471 def get_element_size(tensor):
1472 """Get Tensor Element Size"""
1473 return tensor.element_size()
1475 @staticmethod
1476 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1477 """Allocate an uninitialized 1-D tensor buffer."""
1478 if pin_memory:
1479 return torch.empty(numel, dtype=dtype, device='cpu', pin_memory=True)
1480 return torch.empty(numel, dtype=dtype, device=device)
1482 @staticmethod
1483 def tensor_to_numpy(tensor) -> np.ndarray:
1484 """Convert PyTorch tensor to numpy array."""
1485 return tensor.cpu().numpy()
1487 @staticmethod
1488 def from_numpy(np_array):
1489 """Create a host (CPU) PyTorch tensor from a numpy array."""
1490 return torch.from_numpy(np_array)
1492 @staticmethod
1493 def clip_grad_norm_(
1494 parameters, max_norm, norm_type=2.0,
1495 error_if_nonfinite=False, foreach=None,
1496 ):
1497 # pylint: disable=C0415
1498 from hyper_parallel.platform.torch.clip_grad import (
1499 clip_grad_norm_ as _clip_grad_norm,
1500 )
1501 return _clip_grad_norm(
1502 parameters, max_norm, norm_type,
1503 error_if_nonfinite=error_if_nonfinite, foreach=foreach,
1504 )
1506 @staticmethod
1507 def profiler_record(name):
1508 """Profiler context manager for recording operations using torch.profiler."""
1509 return torch.profiler.record_function(name)
1511 def cast_fp_tensor(self, dtype, x):
1512 """
1513 Cast floating-point tensor to target dtype if applicable.
1514 """
1515 if (
1516 not isinstance(x, torch.Tensor)
1517 or not torch.is_floating_point(x)
1518 or x.dtype == dtype
1519 ):
1520 return x
1521 return x.to(dtype)
1523 def apply_to_tensors(self, fn, container):
1524 """Recursively apply to all tensor in different kinds of container types."""
1526 def apply(x):
1528 if isinstance(x, torch.Tensor):
1529 return fn(x)
1530 if hasattr(x, "__dataclass_fields__"):
1531 dc = dataclasses.replace(x)
1532 changes = {
1533 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
1534 }
1535 return dataclasses.replace(dc, **changes)
1536 if isinstance(x, OrderedDict):
1537 od = x.__class__()
1538 for key, value in x.items():
1539 od[key] = apply(value)
1540 return od
1541 if isinstance(x, PackedSequence):
1542 apply(x.data)
1543 return x
1544 if isinstance(x, dict):
1545 return {key: apply(value) for key, value in x.items()}
1546 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
1547 res = (apply(el) for el in x)
1548 return type(x)(*res)
1549 if isinstance(x, (list, tuple, set)):
1550 return type(x)(apply(el) for el in x)
1551 return x
1553 return apply(container)
1556 @property
1557 def meta_device(self):
1558 return torch.device("meta")
1560 def init_on_device(self, device, include_buffers=False):
1561 return _init_on_device(device, include_buffers=include_buffers)
1563 def str_to_dtype(self, dtype_str: str) -> torch.dtype:
1564 """Map ``torch.<type>`` strings from checkpoint metadata to ``torch.dtype``."""
1565 parts = dtype_str.split(".", 1)
1566 if len(parts) != 2:
1567 raise ValueError(
1568 f"Expected dtype string like 'torch.float32', got {dtype_str!r}."
1569 )
1570 prefix, name = parts
1571 if prefix != "torch":
1572 raise ValueError(
1573 f"Expected PyTorch dtype string with prefix 'torch', got {dtype_str!r}."
1574 )
1575 dtype = getattr(torch, name)
1576 if isinstance(dtype, torch.dtype):
1577 return dtype
1578 raise ValueError(f"{dtype_str!r} does not resolve to a torch.dtype.")
1580 def list_to_size(self, size_list: list[int]) -> torch.Size:
1581 return torch.Size(size_list)