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