Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / platform.py: 63%
839 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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 Any, Optional, Sequence, 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
381class _TorchDifferentiableVariableAllGather(torch.autograd.Function):
382 """Variable dim-zero all-gather with an uneven reduce-scatter backward."""
384 @staticmethod
385 def forward(ctx, input_tensor, output_splits, group): # pylint: disable=arguments-differ
386 """Gather each rank's true row count without replicating inputs for A2A."""
387 if input_tensor.ndim == 0:
388 raise ValueError("variable all-gather input must have at least one dimension")
389 splits = tuple(output_splits)
390 if not splits:
391 raise ValueError("output_splits must contain at least one group rank")
392 if any(not isinstance(rows, int) or isinstance(rows, bool) or rows < 0 for rows in splits):
393 raise ValueError(f"output_splits must contain non-negative integers, got {splits!r}")
395 group_rank = dist.get_rank(group=group)
396 if group_rank < 0 or group_rank >= len(splits):
397 raise ValueError(f"group rank must be in [0, {len(splits)}), got {group_rank}")
398 if input_tensor.shape[0] != splits[group_rank]:
399 raise ValueError(
400 "variable all-gather local rows must match output_splits at the group rank, "
401 f"got local_rows={input_tensor.shape[0]}, group_rank={group_rank}, "
402 f"output_splits={splits!r}"
403 )
405 input_tensor = input_tensor.contiguous()
406 feature_shape = tuple(input_tensor.shape[1:])
407 if input_tensor.device.type == "npu":
408 gathered = [input_tensor.new_empty((rows, *feature_shape)) for rows in splits]
409 dist.all_gather(gathered, input_tensor, group=group)
410 else:
411 max_rows = max(splits)
412 if max_rows == 0:
413 gathered = [input_tensor.new_empty((0, *feature_shape)) for _ in splits]
414 else:
415 padded = input_tensor.new_zeros((max_rows, *feature_shape))
416 if input_tensor.shape[0] > 0:
417 padded[:input_tensor.shape[0]].copy_(input_tensor)
418 padded_outputs = [torch.empty_like(padded) for _ in splits]
419 dist.all_gather(padded_outputs, padded, group=group)
420 gathered = [
421 output[:rows].contiguous()
422 for output, rows in zip(padded_outputs, splits)
423 ]
425 ctx.output_splits = splits
426 ctx.group = group
427 ctx.group_rank = group_rank
428 return torch.cat(gathered, dim=0)
430 @staticmethod
431 def backward(ctx, grad_output):
432 """Sum replicated output gradients and return this rank's uneven shard."""
433 output_rows = ctx.output_splits[ctx.group_rank]
434 output = grad_output.new_empty((output_rows, *grad_output.shape[1:]))
435 if sum(ctx.output_splits) == 0:
436 return output, None, None
438 grad_output = grad_output.contiguous()
439 if grad_output.device.type == "npu":
440 from torch_npu.distributed import reduce_scatter_tensor_uneven # pylint: disable=C0415
441 reduce_scatter_tensor_uneven(
442 output,
443 grad_output,
444 input_split_sizes=list(ctx.output_splits),
445 op=dist.ReduceOp.SUM,
446 group=ctx.group,
447 )
448 else:
449 reduced = grad_output.clone()
450 dist.all_reduce(reduced, op=dist.ReduceOp.SUM, group=ctx.group)
451 start = sum(ctx.output_splits[:ctx.group_rank])
452 output.copy_(reduced.narrow(0, start, output_rows))
453 return output, None, None
456# Mapping from string op names to torch.distributed.ReduceOp
457_OP_MAP = {
458 'sum': dist.ReduceOp.SUM,
459 'prod': dist.ReduceOp.PRODUCT,
460 'max': dist.ReduceOp.MAX,
461 'min': dist.ReduceOp.MIN,
462 # convert tensor elements to int32 and use MIN
463 'all': dist.ReduceOp.MIN,
464 # 'avg' is typically handled by SUM followed by division in current implementation logic
465 'avg': dist.ReduceOp.SUM,
466}
468# Try to add AVG for 'mean' if supported by current torch version
469if hasattr(dist.ReduceOp, "AVG"):
470 _OP_MAP['mean'] = dist.ReduceOp.AVG
471else:
472 # Fallback for older torch versions if necessary, though this might require manual division upstream
473 # Assuming standard behavior where 'mean' implies native AVG support or upstream handling
474 _OP_MAP['mean'] = dist.ReduceOp.SUM
477def _ensure_contiguous(x):
478 """Return a contiguous copy of *x* if not already contiguous."""
479 if not x.is_contiguous() or x.storage_offset() != 0:
480 x = x.contiguous()
481 return x
484class _TorchBatchP2PWork:
485 """Single ``.wait()`` handle wrapping the per-op works returned by
486 ``torch.distributed.batch_isend_irecv``.
488 Torch returns one ``Work`` per op in the batch (the ops are coalesced
489 onto one comm stream), whereas the platform contract — and the scheduler
490 that consumes it — expects a single handle covering the whole batch so
491 the wait can be deferred to one consumption point (mirroring MindSpore's
492 single packaging ``CommHandle``). Waiting this handle waits every
493 underlying op.
494 """
496 __slots__ = ("_works",)
498 def __init__(self, works):
499 self._works = works
501 def wait(self):
502 for work in self._works:
503 if work is not None:
504 work.wait()
507# pylint: disable=C0103
508class TorchPlatform(Platform):
509 """Torch platform api"""
510 Tensor = Tensor
511 tensor = torch.tensor
512 Parameter = Parameter
513 Module = Module
514 DTensorBase = DTensorBase
515 PipelineStageBase = PipelineStageBase
516 platform_type = PlatformType.PYTORCH
517 tensor_dtype = torch
518 dtype = torch.dtype
519 Function = torch.autograd.Function
521 _custom_ops_cls = None
523 @property
524 def custom_ops(self):
525 """Return the Torch platform custom ops instance.
527 .. warning::
528 This is an experimental API that subject to change or deletion.
530 Returns:
531 TorchCustomOps: Custom ops class that raises NotImplementedError
532 for all operators (MindSpore-only at this time).
533 """
534 if self._custom_ops_cls is None:
535 from hyper_parallel.platform.torch.custom_ops import TorchCustomOps # pylint: disable=import-outside-toplevel
536 self._custom_ops_cls = TorchCustomOps
537 return self._custom_ops_cls
539 @staticmethod
540 def get_swap_optimizer():
541 """Return the Torch optimizer-state swap wrapper class."""
542 from hyper_parallel.platform.torch.swap_optimizer.swap_optimizer import ( # pylint: disable=import-outside-toplevel
543 get_swap_optimizer,
544 )
545 return get_swap_optimizer()
547 @staticmethod
548 def is_linear_module(module) -> bool:
549 """Check whether *module* is a ``torch.nn.Linear`` instance."""
550 return isinstance(module, nn.Linear)
552 @staticmethod
553 def is_embedding_module(module) -> bool:
554 """Check whether *module* is a ``torch.nn.Embedding`` instance."""
555 return isinstance(module, nn.Embedding)
557 @staticmethod
558 def device_count(device_handle):
559 """
560 Get the number of available devices.
562 Args:
563 device_handle: The device handle (e.g., torch.cuda, torch.npu).
565 Returns:
566 int: The number of available devices.
567 """
568 return device_handle.device_count()
570 def device_type(self):
571 """
572 Get the current device type.
574 Returns:
575 str: The device type string ("npu" for NPU, "cuda" for GPU).
576 """
577 device_handle = self.get_device_handle()
578 if device_handle == torch.npu:
579 return "npu"
580 return "cuda"
582 def device(self, device_idx=None):
583 """
584 Get a torch.device object for the specified device index.
586 Args:
587 device_idx (Optional[int]): The device index. If None, returns device without index.
589 Returns:
590 torch.device: A torch device object.
591 """
592 device_type = self.device_type()
593 if device_idx is None:
594 return torch.device(device_type)
595 return torch.device(f"{device_type}:{device_idx:d}")
597 @staticmethod
598 def get_rng_state(device=None, device_handle=None):
599 """
600 Get the random number generator state.
602 Args:
603 device (Optional): The device to get RNG state from.
604 device_handle (Optional): The device handle (torch.cuda, torch.npu, etc.).
606 Returns:
607 Tensor: The RNG state as a byte tensor.
608 """
609 if device_handle is None:
610 return torch.get_rng_state()
611 if device is None:
612 return device_handle.get_rng_state()
613 return device_handle.get_rng_state(device)
615 @staticmethod
616 def set_rng_state(state, device=None, device_handle=None):
617 """
618 Set the random number generator state.
620 Args:
621 state (Tensor): The RNG state to set.
622 device (Optional): The device to set RNG state for.
623 device_handle (Optional): The device handle (torch.cuda, torch.npu, etc.).
624 """
625 if device_handle is None:
626 return torch.set_rng_state(state)
627 if device is None:
628 return device_handle.set_rng_state(state)
629 return device_handle.set_rng_state(state, device)
631 @staticmethod
632 def manual_seed(seed):
633 """
634 Set the random seed for reproducibility.
636 Args:
637 seed (int): The random seed value.
639 Returns:
640 torch.Generator: The random number generator.
641 """
642 return torch.manual_seed(seed)
644 @staticmethod
645 def ones(size, dtype=None):
646 """
647 Create a tensor filled with ones.
649 Args:
650 size (tuple): The shape of the output tensor.
651 dtype (Optional[torch.dtype]): The desired data type.
653 Returns:
654 Tensor: A tensor filled with ones.
655 """
656 return torch.ones(size, dtype=dtype)
658 @staticmethod
659 def zeros(size, dtype=None, device=None):
660 """
661 Create a tensor filled with zeros.
663 Args:
664 size (tuple): The shape of the output tensor.
665 dtype (Optional[torch.dtype]): The desired data type.
666 device (Optional[torch.device]): The device to create the tensor on.
668 Returns:
669 Tensor: A tensor filled with zeros.
670 """
671 return torch.zeros(size, dtype=dtype, device=device)
673 @staticmethod
674 def full(size, fill_value, dtype=None):
675 """
676 Create a tensor filled with a scalar value.
678 Args:
679 size (tuple): The shape of the output tensor.
680 fill_value (scalar): The value to fill the tensor with.
681 dtype (Optional[torch.dtype]): The desired data type.
683 Returns:
684 Tensor: A tensor filled with the specified value.
685 """
686 return torch.full(size, fill_value, dtype=dtype)
688 @staticmethod
689 def empty(size, dtype=None, device=None):
690 """
691 Create an uninitialized tensor.
693 Args:
694 size (tuple): The shape of the output tensor.
695 dtype (Optional[torch.dtype]): The desired data type.
696 device (Optional[torch.device or str]): Target device. When
697 ``None`` the tensor is allocated on the default device
698 (CPU under PyTorch defaults), matching the original
699 back-compat behavior.
701 Returns:
702 Tensor: An uninitialized tensor.
703 """
704 return torch.empty(size, dtype=dtype, device=device)
706 @staticmethod
707 def rand(size, dtype=None, device=None):
708 """Create a tensor filled with uniform random values in ``[0, 1)``."""
709 return torch.rand(size, dtype=dtype, device=device)
711 @staticmethod
712 def randn(size, dtype=None, device=None):
713 """Create a tensor filled with standard-normal random values."""
714 return torch.randn(size, dtype=dtype, device=device)
716 @staticmethod
717 def get_rank():
718 """
719 Get the rank of the current process in the distributed group.
721 Returns:
722 int: The rank of the current process.
723 """
724 return dist.get_rank()
726 @staticmethod
727 def get_global_rank(group, group_rank):
728 """
729 Get the global rank from a group rank.
731 Args:
732 group (ProcessGroup): The process group.
733 group_rank (int): The rank within the group.
735 Returns:
736 int: The global rank.
737 """
738 return dist.get_global_rank(group, group_rank)
740 @staticmethod
741 def get_group_rank(group):
742 """Return this process's rank within *group*."""
743 return dist.get_group_rank(group, dist.get_rank())
745 @staticmethod
746 def get_world_size():
747 """
748 Get the total number of processes in the distributed group.
750 Returns:
751 int: The world size.
752 """
753 return dist.get_world_size()
755 @staticmethod
756 def get_param_local_shape(param):
757 """
758 Get the local shape of a parameter, handling both regular and distributed tensors.
760 Args:
761 param (Union[Tensor, DTensorBase]): The parameter tensor.
763 Returns:
764 torch.Size: The local shape of the parameter.
765 """
766 if isinstance(param, DTensorBase):
767 return param.local_shape
768 return param.shape
770 @staticmethod
771 def get_param_local_data(param):
772 """
773 Get the local data of a parameter, handling both regular and distributed tensors.
775 Args:
776 param (Union[Tensor, DTensorBase]): The parameter tensor.
778 Returns:
779 Tensor: The local tensor data.
780 """
781 if isinstance(param, DTensorBase):
782 return param.to_local()
783 return param
785 @staticmethod
786 def update_param_data(param, data):
787 """
788 Update the data of a parameter.
790 Args:
791 param (Parameter): The parameter to update.
792 data (Tensor): The new data tensor.
793 """
794 param.data = data
796 @staticmethod
797 def load_into_param(param, data):
798 """Load tensor *data* into *param* (plain tensor or DTensor)."""
799 if isinstance(param, DTensorBase):
800 local = param._local_tensor # pylint: disable=W0212
801 if local.is_meta:
802 # Meta tensor materialisation: replace the placeholder.
803 orig_requires_grad = param.requires_grad
804 param._local_tensor = data # pylint: disable=W0212
805 if data.requires_grad != orig_requires_grad:
806 param.requires_grad_(orig_requires_grad)
807 else:
808 local.copy_(data)
809 else:
810 param.copy_(data)
812 @staticmethod
813 def get_op_name(func):
814 """
815 Extract the operation name from various function types.
817 Args:
818 func: The function or operation to extract the name from.
820 Returns:
821 str: The operation name.
822 """
823 if hasattr(func, "__name__"):
824 return func.__name__
825 if isinstance(func, OpOverload):
826 full_name = func.name
827 core_name = full_name.split("::")[-1].split(".")[0]
828 return core_name
829 if isinstance(func, OpOverloadPacket):
830 return func.name.split("::")[-1]
831 func_str = str(func)
832 if "built-in function" in func_str:
833 return func_str.split()[-1].strip(">")
834 if "function" in func_str:
835 return func_str.split()[1]
836 return "unknown_op"
838 @staticmethod
839 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
840 data = _ensure_contiguous(data)
841 output = list(dist_func.all_gather(data, group=group))
842 if rank_list is not None:
843 group_ranks = dist.get_process_group_ranks(group)
844 if tuple(rank_list) != tuple(group_ranks):
845 rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}
846 output = [output[rank_to_idx[int(rank)]] for rank in rank_list]
847 return torch.cat(output, dim=concat_dim)
849 @staticmethod
850 def chunk(data, split_dim, split_size, index):
851 return torch.chunk(data, split_size, dim=split_dim)[index]
853 @staticmethod
854 def differentiable_all_to_all(input_data, output_shape, group):
855 input_data = _ensure_contiguous(input_data)
856 output_tensor = torch.empty(output_shape, device=input_data.device, dtype=input_data.dtype)
857 output_tensor = dist_func.all_to_all_single(
858 output_tensor,
859 input_data,
860 group=group
861 )
862 return output_tensor
864 @staticmethod
865 def tensor_type_cast(input_data, cast_type):
866 """Cast tensor to specified data type."""
867 type_mapping = {
868 'float32': torch.float32,
869 'float16': torch.float16,
870 'int64': torch.int64,
871 'int32': torch.int32
872 }
873 if cast_type not in type_mapping:
874 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")
875 return input_data.to(type_mapping[cast_type])
877 @staticmethod
878 def differentiable_all_reduce(data, op, group):
879 data = _ensure_contiguous(data)
880 # Resolve the op from string to ReduceOp enum if necessary
881 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
882 return dist_func.all_reduce(data, op=reduce_op, group=group)
884 @staticmethod
885 def get_cell_construct(cell):
886 return cell.forward
888 @staticmethod
889 def get_cells_and_names(cell):
890 return cell.named_modules()
892 @staticmethod
893 def get_modules(module):
894 return module.modules()
896 @staticmethod
897 def search_parameter_by_name(cell, param_name: str):
898 """
899 Find the parent Module of the parameter, the parameter's name in the parent Module, and the parameter.
900 Return value: (parent Module instance, parameter's name in parent Module, parameter object).
901 Returns None if not found.
902 """
903 # Remove the "self." prefix from param_name
904 param_name = param_name.replace("self.", "")
905 # Case 1: The parameter is a direct parameter of the current Module
906 if param_name in cell._parameters: # pylint: disable=protected-access
907 return (cell, param_name, cell._parameters[param_name]) # pylint: disable=protected-access
909 # Case 2: The parameter is in a sub-Module
910 if "." in param_name:
911 cell_path, param_key = param_name.rsplit(".", 1)
912 try:
913 # Locate the sub-Module where the parameter resides (supports multi-level paths)
914 target_cell = cell.get_submodule(cell_path)
915 # Check if the sub-Module directly contains this parameter
916 if param_key in target_cell._parameters: # pylint: disable=protected-access
917 return target_cell, param_key, target_cell._parameters[param_key] # pylint: disable=protected-access
918 except AttributeError:
919 pass
921 # Traverse all sub-Modules (recursively) to search for the parameter
922 for _, child_cell in cell.named_children():
923 if isinstance(child_cell, Module):
924 result = TorchPlatform.search_parameter_by_name(child_cell, param_name)
925 if result is not None:
926 return result
928 return None
930 @staticmethod
931 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
932 """
933 Modify the original parameter in a Module or sub-Module using the search result
934 """
935 parent_cell, param_key, _ = result
936 # Key operation: directly modify the _parameters dictionary.
937 if param_key in parent_cell._parameters: # pylint: disable=protected-access
938 parent_cell._parameters[param_key] = new_param # pylint: disable=protected-access
939 else:
940 parent_cell.register_parameter(param_key, new_param)
941 return True
943 @staticmethod
944 def set_layout_into_parameter(param, layout):
945 """Set layout into parameter"""
946 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel
947 from hyper_parallel.core.dtensor.layout import _get_slice_tensor_by_layout # pylint: disable=import-outside-toplevel
948 if isinstance(param, DTensor):
949 raise ValueError(f"Parameter {param} has been configured layout, cannot be set repeatedly.")
950 requires_grad = param.requires_grad
951 param_dtensor = DTensor.from_local(
952 _get_slice_tensor_by_layout(param, layout),
953 layout.mesh, layout.alias_placements)
954 new_param = Parameter(param_dtensor, requires_grad=requires_grad)
955 return new_param
957 @staticmethod
958 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
959 data = _ensure_contiguous(data)
960 input_tuple = torch.chunk(data, dev_num, dim=axis)
961 output_tensor = torch.empty(input_tuple[0].shape, device=data.device, dtype=data.dtype)
963 # Resolve the op from string to ReduceOp enum
964 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
966 output_tensor = dist_func.reduce_scatter(output_tensor, input_tuple, op=reduce_op, group=group)
968 # Keep manual handling for 'avg' string as it maps to SUM in _OP_MAP
969 if op == 'avg':
970 output_tensor = output_tensor / dev_num
971 return output_tensor
973 @staticmethod
974 def get_device_handle(device_type: str = "npu"):
975 """Return the torch device module (e.g. ``torch.npu`` or ``torch.cuda``) for the given device type."""
976 try:
977 handle = getattr(torch, device_type)
978 except AttributeError as e:
979 raise RuntimeError(f"TorchPlatform expect got device handle: 'torch.{device_type}' failed.") from e
980 return handle
982 @staticmethod
983 def get_param_type_size(param):
984 # pylint: disable=W0212
985 return torch._utils._element_size(param.dtype)
987 @staticmethod
988 def is_tensor(obj: Any) -> bool:
989 """Return True if ``obj`` is a ``torch.Tensor``."""
990 return isinstance(obj, Tensor)
992 @staticmethod
993 def get_tensor_storage_size(tensor: Any) -> int:
994 """Return serialized byte size (numel * element size) for a PyTorch tensor."""
995 if not TorchPlatform.is_tensor(tensor):
996 raise TypeError(
997 f"TorchPlatform.get_tensor_storage_size expects torch.Tensor, got {type(tensor)!r}"
998 )
999 return int(tensor.numel()) * int(tensor.element_size())
1001 @staticmethod
1002 def parameters_dict(cell: Module):
1003 return cell.named_parameters()
1005 @staticmethod
1006 def buffers_dict(cell: Module) -> Any:
1007 """Return all named buffers registered by the module tree."""
1008 return cell.named_buffers()
1010 @staticmethod
1011 def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
1012 """Get the state dictionary of a model.
1014 Delegates to torch-specific implementation that handles DTensor
1015 gathering, CPU offloading and frozen-parameter filtering.
1016 """
1017 # pylint: disable=C0415
1018 from hyper_parallel.platform.torch.fully_shard.state_dict_utils import (
1019 get_model_state_dict as _get_model_state_dict,
1020 )
1021 return _get_model_state_dict(model, options=options)
1023 @staticmethod
1024 def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
1025 """Set the state dictionary of a model.
1027 Delegates to torch-specific implementation that scatters full tensors
1028 into DTensor shards and performs an in-place load.
1029 """
1030 # pylint: disable=C0415
1031 from hyper_parallel.platform.torch.fully_shard.state_dict_utils import (
1032 set_model_state_dict as _set_model_state_dict,
1033 )
1034 return _set_model_state_dict(model, model_state_dict, options=options)
1036 @staticmethod
1037 def save_checkpoint(cell: Module, file_path: str, ckpt_format: str = "safetensors") -> None:
1038 if ckpt_format == "safetensors":
1039 save_file(tensors=cell, filename=file_path)
1040 else:
1041 torch.save(obj=cell, f=file_path)
1043 @staticmethod
1044 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
1045 if ckpt_format == "safetensors":
1046 return load_file(filename=file_path)
1047 return torch.load(f=file_path)
1049 @staticmethod
1050 def new_zero_parameter(param_shape, param_type, requires_grad, device):
1051 return nn.Parameter(torch.zeros(param_shape, dtype=param_type, device=device), requires_grad=requires_grad)
1053 @staticmethod
1054 def new_tensor(tensor_shape, tensor_type, device):
1055 return torch.empty(size=tensor_shape, dtype=tensor_type, device=device)
1057 @staticmethod
1058 def full_like(tensor, fill_value, dtype=None):
1059 return torch.full_like(tensor, fill_value, dtype=dtype)
1061 @staticmethod
1062 def set_tensor_requires_grad(input_tensor):
1063 """
1064 set requires grad flag for input tensor, only effective for leaf node
1065 """
1066 if input_tensor.is_leaf:
1067 input_tensor.requires_grad = True
1069 def _create_group(self, rank_list):
1070 normalized_rank_list = tuple(sorted(rank_list))
1071 world_rank_list = tuple(range(self.get_world_size()))
1072 if normalized_rank_list == world_rank_list:
1073 group = _get_default_group()
1074 EXISTING_COMM_GROUPS[str(normalized_rank_list)] = group
1075 return group
1076 group_dict = create_sub_groups(rank_list)
1077 return group_dict[normalized_rank_list]
1079 @staticmethod
1080 def all_gather_into_tensor(data, group_info, async_op=False):
1081 output_shape = list(data.shape)
1082 output_shape[0] = output_shape[0] * group_info.rank_size
1083 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
1084 handle = dist.all_gather_into_tensor(output, data, group=group_info.group, async_op=async_op)
1085 return output, handle
1087 @staticmethod
1088 def all_gather_single(input_tensor, output_shape, group, async_op=False):
1089 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
1090 handle = dist.all_gather_into_tensor(output, input_tensor, group=group, async_op=async_op)
1091 return output, handle
1093 @staticmethod
1094 def all_reduce(data, group_info, async_op=False):
1095 if not data.is_contiguous():
1096 data = data.contiguous()
1097 handle = dist.all_reduce(data, group=group_info.group, async_op=async_op)
1098 return data, handle
1100 @staticmethod
1101 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
1102 if group_src is not None:
1103 src = dist.get_global_rank(group, group_src)
1104 handle = dist.broadcast(data, src, group, async_op)
1105 if async_op and handle is not None:
1106 handle.wait()
1108 @staticmethod
1109 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
1110 if group_src is not None:
1111 src = dist.get_global_rank(group, group_src)
1112 handle = dist.scatter(output, scatter_list, src=src, group=group, async_op=async_op)
1113 if async_op and handle is not None:
1114 handle.wait()
1115 return output
1117 @staticmethod
1118 def isend(tensor, dst=None, group=None, tag=0):
1119 return dist.isend(tensor, dst, group, tag)
1121 @staticmethod
1122 def irecv(tensor, src=None, group=None, tag=0):
1123 return dist.irecv(tensor, src, group, tag)
1125 @staticmethod
1126 def p2p_op(op_type, tensor, peer, group=None):
1127 # torch's P2POp takes the op callable (dist.isend / dist.irecv), not
1128 # the "isend"/"irecv" string the stage specs builders emit.
1129 if op_type == "isend":
1130 op = dist.isend
1131 elif op_type == "irecv":
1132 op = dist.irecv
1133 else:
1134 raise ValueError(
1135 f"p2p_op op_type must be 'isend' or 'irecv', but got {op_type!r}."
1136 )
1137 return dist.P2POp(op, tensor, peer, group)
1139 @staticmethod
1140 def batch_isend_irecv(p2p_ops):
1141 """Launch a peer-batched P2P group as one coalesced op.
1143 ``torch.distributed.batch_isend_irecv`` coalesces the ops onto one
1144 comm stream and returns one ``Work`` per op; we wrap them in a single
1145 ``.wait()`` handle so a send and a recv to the same peer overlap on
1146 the duplex link and the caller can defer the whole batch's wait to one
1147 consumption point.
1148 """
1149 if not p2p_ops:
1150 return None
1151 works = dist.batch_isend_irecv(p2p_ops)
1152 return _TorchBatchP2PWork(works) if works else None
1154 @staticmethod
1155 def prepare_batch_p2p_group(group: Any = None) -> None:
1156 """Synchronize a group before its first subset batched P2P call.
1158 PyTorch requires every rank in a process group to participate when
1159 ``batch_isend_irecv`` is the first collective on that group. A barrier
1160 at the common pipeline run boundary initializes the communicator
1161 before ranks reach peer operations at different times.
1163 Args:
1164 group: The process group used by the batched P2P operations.
1165 ``None`` uses the default group.
1166 """
1167 dist.barrier(group=group)
1169 @staticmethod
1170 def p2p_exchange(tensor, peer_rank: int, group=None):
1171 if peer_rank == dist.get_rank(group):
1172 return tensor
1173 return _TorchP2PExchangeFunction.apply(tensor, peer_rank, group)
1175 @staticmethod
1176 def send_object_list(obj_list, dst=None, group=None):
1177 dist.send_object_list(obj_list, dst, group)
1179 @staticmethod
1180 def recv_object_list(obj_list, src=None, group=None):
1181 dist.recv_object_list(obj_list, src, group)
1183 @staticmethod
1184 def reduce_scatter_tensor(data, group_info, async_op=False):
1185 output_shape = list(data.shape)
1186 output_shape[0] = output_shape[0] // group_info.rank_size
1187 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
1188 handle = dist.reduce_scatter_tensor(output, data, group=group_info.group, async_op=async_op)
1189 return output, handle
1191 @staticmethod
1192 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1193 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
1194 handle = dist.reduce_scatter_tensor(output, input_tensor, group=group, async_op=async_op)
1195 return output, handle
1197 @staticmethod
1198 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1199 output = torch.empty(output_shape, device=input_tensor.device, dtype=input_tensor.dtype)
1200 work = dist.all_to_all_single(output, input_tensor, group=group, async_op=async_op)
1201 return output, work
1203 @staticmethod
1204 def differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group):
1205 """Variable-split all-to-all with autograd support for EP token dispatch/combine."""
1206 out_total = sum(output_splits)
1207 output = torch.empty(
1208 out_total, *input_tensor.shape[1:],
1209 dtype=input_tensor.dtype, device=input_tensor.device,
1210 )
1211 output = dist_func.all_to_all_single(
1212 output, input_tensor,
1213 output_split_sizes=output_splits,
1214 input_split_sizes=input_splits,
1215 group=group,
1216 )
1217 return output
1219 @staticmethod
1220 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1221 """Truly-async variant of :meth:`differentiable_all_to_all_single`.
1223 Both forward AND backward return :class:`AsyncCollectiveTensor`,
1224 so the ``wait_tensor`` op is queued lazily — only when a downstream
1225 kernel actually reads the result.
1227 Why both directions need lazy wait:
1229 * FWD: ACT lazy wait lets host return immediately and the paired
1230 BWD thread's compute kernel slip into the queue before the wait.
1231 * BWD: PyTorch's stock backward issues ``wait_tensor`` eagerly,
1232 and the autograd engine binds backward stream to the forward
1233 stream — so even running BWD inside a ``with torch.npu.stream
1234 (side_stream)`` context does not move that wait off the main
1235 stream. Returning ACT from backward defers the wait to the
1236 next backward op's first consumption, opening a small window
1237 during which FWD's Attention kernels can be queued onto the
1238 main stream **before** the wait lands.
1240 Args:
1241 input_tensor: Input tensor, split along dim 0 by ``input_splits``.
1242 input_splits: ``list[int]`` — rows sent to each rank.
1243 output_splits: ``list[int]`` — rows received from each rank.
1244 group: Process group.
1246 Returns:
1247 ``AsyncCollectiveTensor`` of shape
1248 ``[sum(output_splits), *input_tensor.shape[1:]]``.
1249 """
1250 return _AsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1252 @staticmethod
1253 def differentiable_variable_all_gather(
1254 input_tensor: Tensor, output_splits: Sequence[int], group: Any) -> Tensor:
1255 """Gather variable dim-zero shards on HCCL or Gloo with autograd support."""
1256 return _TorchDifferentiableVariableAllGather.apply(
1257 input_tensor, tuple(output_splits), group
1258 )
1260 @staticmethod
1261 def wait_async_tensor(tensor):
1262 """Wait for an async collective tensor to become materialised.
1264 Idempotent — calling on an already-waited tensor is a no-op.
1266 Args:
1267 tensor: ``AsyncCollectiveTensor`` whose device-side values may
1268 not yet be ready.
1270 Returns:
1271 The same *tensor*, now fully materialised.
1272 """
1273 from torch.distributed._functional_collectives import wait_tensor # pylint: disable=C0415
1274 wait_tensor(tensor)
1275 return tensor
1277 @staticmethod
1278 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1279 handle_box=None):
1280 """Wait async all-gather handle and reconstruct result (differentiable)."""
1281 return _TorchAsyncAllGatherFunction.apply(
1282 x, work, out_perm, group, world_size, gather_dim, handle_box
1283 )
1285 @staticmethod
1286 def arange(start, end=None, step=1, dtype=None, device=None):
1287 """Create a 1-D tensor with evenly spaced values."""
1288 if end is None:
1289 return torch.arange(start, dtype=dtype, device=device)
1290 return torch.arange(start, end, step, dtype=dtype, device=device)
1292 @staticmethod
1293 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim,
1294 handle_box=None):
1295 """Wait async A2A handle and reconstruct result (differentiable).
1297 Args:
1298 x: Input tensor.
1299 work: Async work handle from all_to_all.
1300 out_perm: Output buffer from all_to_all.
1301 group: Process group.
1302 world_size: World size.
1303 concat_dim: Dimension for concatenation.
1304 split_dim: Dimension for split.
1305 handle_box: Optional mutable list; backward appends (work, out_perm) here.
1306 """
1307 return _TorchAsyncA2AFunction.apply(
1308 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1309 )
1311 @staticmethod
1312 def differentiable_sync_hook(x, hook_name: str, coordinator):
1313 """Identity op that fires coordinator rendezvous on forward and backward.
1315 Always goes through ``_TorchSyncHookFunction.apply`` so that the
1316 autograd graph **records a SyncHook node regardless of whether the
1317 coordinator is currently enabled**. Skipping ``apply`` when
1318 disabled would leave warmup-forwarded graphs without the hook
1319 nodes, and a later ``overlap.run`` — whose BWD thread back-props
1320 such a graph — would then traverse zero hooks while the paired FWD
1321 thread (whose current forward DOES record hooks) waits at a
1322 barrier for a partner that never arrives.
1324 Args:
1325 x: Input tensor.
1326 hook_name: One of:
1327 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1328 full rendezvous on both directions.
1329 * ``"D_LAST"`` — closing D of the last MoE
1330 layer in a chunk. Forward: ``notify_dispatched``
1331 only (no Attention follows so rendezvous is
1332 skipped). Backward: pure skip (first BWD
1333 hook to fire; combine.bwd has already
1334 dispatched freely).
1335 coordinator: A :class:`HookCoordinator` instance.
1336 """
1337 return _TorchSyncHookFunction.apply(x, hook_name, coordinator)
1339 @staticmethod
1340 def get_tensor_transform():
1341 raise NotImplementedError("Unsupported get_tensor_transform for torch platform")
1343 @staticmethod
1344 def construct_strided_slice(x, begin, end, stride):
1345 raise NotImplementedError("Unsupported construct_strided_slice for torch platform")
1347 @staticmethod
1348 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1349 # pylint: disable=C0415
1350 from hyper_parallel.platform.torch.pipeline_parallel._utils import _MicroBatch
1351 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1353 @staticmethod
1354 def get_symmetric_memory_handler():
1355 # pylint: disable=C0415
1356 from hyper_parallel.platform.torch.symmetric_memory import TorchSymmetricMemoryHandler
1357 symmetric_memory = TorchSymmetricMemoryHandler()
1358 return symmetric_memory
1360 @staticmethod
1361 def get_multicore_handler():
1362 """Return a TorchMulticoreHandler instance for multi-core device management."""
1363 # pylint: disable=C0415
1364 from hyper_parallel.platform.torch.multicore import TorchMulticoreHandler
1365 return TorchMulticoreHandler()
1367 def new_stream(self):
1368 device = self.get_device_handle()
1369 return device.Stream()
1371 def get_stream_context(self):
1372 device = self.get_device_handle()
1373 return device.stream
1375 @staticmethod
1376 def all_gather_object(object_list, obj, group=None) -> None:
1377 """
1378 Gathers objects from the given group into object list.
1380 Args:
1381 object_list (list[Any]): Define the output list, which size equal to the size of group.
1382 obj (Any): The object on current rank and in given process group.
1383 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1384 global group.
1386 Returns:
1387 None. Objs are gathered into ``object_list``.
1388 """
1389 dist.all_gather_object(object_list, obj, group)
1391 @staticmethod
1392 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1393 """
1394 Synchronize all processes in the given process group.
1396 Args:
1397 group (ProcessGroup, optional): The process group to work on. Default is ``None``,
1398 meaning the default process group.
1399 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1400 device_ids (list[int], optional): Device ids for backends that require a device for
1401 barrier (e.g. NCCL). Default: ``None``.
1403 Returns:
1404 Async work handle if ``async_op`` is True; otherwise ``None``.
1405 """
1406 return dist.barrier(group, async_op, device_ids)
1408 @staticmethod
1409 def init_process_group(
1410 backend: Optional[str] = None,
1411 *,
1412 init_method: Optional[str] = None,
1413 timeout: Optional[timedelta] = None,
1414 world_size: int = -1,
1415 rank: int = -1,
1416 store: Optional[Store] = None,
1417 pg_options: Optional[Any] = None,
1418 device_id: Optional[Union[torch.device, int]] = None,
1419 ) -> None:
1420 """
1421 Initialize global process group.
1423 Args:
1424 backend (str or Backend, optional): The backend to use for distributed communication.
1425 init_method (str, optional): URL specifying how to initialize the process group. Default is "env://",
1426 can not be specified at the same time with ``store``.
1427 timeout (timedelta, optional): Timeout for process group. Default 10 minutes for NCCL and for other
1428 backends 30 minutes.
1429 world_size (int, optional): Number of processes. If ``store`` is specified, world_size is required.
1430 rank (int, optional): Rank of the current process, which value must between 0 and ``world_size``-1. If
1431 ``store`` is specified, rank is required.
1432 store (Store, optional): Key/value store accessible to all workers, used to exchange connection/address
1433 information. Can not be specified at the same time with ``init_method``.
1434 pg_options (ProcessGroupOptions, optional): Extra options to pass during constructing process groups.
1435 device_id (torch.device | int, optional): Specific device this process will work on.
1436 """
1437 try:
1438 _get_default_group()
1439 # except multi version error
1440 except (ValueError, RuntimeError):
1441 if backend is None:
1442 backend = "hccl"
1443 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1444 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1446 @staticmethod
1447 def destroy_process_group(group: Optional[ProcessGroup] = None) -> None:
1448 """
1449 Destroy given process group.
1451 Args:
1452 group (ProcessGroup, optional): Given process group will be destroyed, if not given, all process groups
1453 will be destroyed.
1454 """
1455 group = group or _get_default_group()
1456 if group in EXISTING_COMM_GROUPS.values():
1457 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1458 for k in keys_to_destroy:
1459 del EXISTING_COMM_GROUPS[k]
1460 dist.destroy_process_group(group)
1462 @staticmethod
1463 def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> list[int]:
1464 """
1465 Get all ranks relative to given process group.
1467 Args:
1468 group (Optional[ProcessGroup]): Process group worked on. Default is ``None``, and ``None`` means global
1469 group.
1471 Returns:
1472 Rank list.
1473 """
1474 group = group or _get_default_group()
1475 return dist.get_process_group_ranks(group)
1477 @staticmethod
1478 def get_backend(group: Optional[ProcessGroup] = None) -> Backend:
1479 """
1480 Get the backend of the given process group.
1482 Args:
1483 group (ProcessGroup, optional): Process group worked on. Default is ``None``, and ``None`` means global
1484 group.
1486 Returns:
1487 The backend object of the given process group.
1488 """
1489 group = group or _get_default_group()
1490 return dist.get_backend(group)
1492 @staticmethod
1493 def split_group(parent_pg: Optional[ProcessGroup] = None,
1494 split_ranks: Optional[list] = None,
1495 timeout: Optional[timedelta] = None,
1496 pg_options: Optional[Any] = None,
1497 group_desc: Optional[str] = None,
1498 ) -> Optional[ProcessGroup]:
1499 """
1500 Create split groups for every group rank in split_ranks, and return the split process group which relative to
1501 current rank id.
1503 Args:
1504 parent_pg (Optional[ProcessGroup]): A process group which the goal group split from.
1505 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1506 timeout (Optional[timedelta]): Timeout for process group. Default 10 minutes for NCCL and for other
1507 backend 30 minutes.
1508 pg_options (Optional[Any]): Extra options to pass during constructing process groups.
1509 group_desc (Optional[str]): Description of process group.
1511 Return:
1512 Optional[ProcessGroup]: One of split process group which relative to current rank id
1513 """
1514 if split_ranks is None or len(split_ranks) == 0:
1515 raise ValueError("split_ranks cannot be None or empty")
1517 split_group = None
1518 for split_rank in split_ranks:
1519 dist_group = TorchPlatform.get_created_group(split_rank)
1520 if dist_group is None:
1521 dist_group = dist.new_group(ranks=split_rank)
1522 EXISTING_COMM_GROUPS[str(tuple(sorted(split_rank)))] = dist_group
1523 if TorchPlatform.get_rank() in split_rank:
1524 split_group = dist_group
1526 return split_group
1528 @staticmethod
1529 def get_group_local_rank(group: ProcessGroup = None) -> int:
1530 """get group local rank id."""
1531 group = group or _get_default_group()
1532 return group.rank()
1534 @staticmethod
1535 def no_grad():
1536 return torch.no_grad()
1538 @staticmethod
1539 def preserve_version_counter(tensor):
1540 return torch.autograd._unsafe_preserve_version_counter(tensor) # pylint: disable=W0212
1542 @staticmethod
1543 def relu(tensor):
1544 return torch.relu(tensor)
1546 @staticmethod
1547 def cat(tensors, dim=0):
1548 return torch.cat(tensors, dim=dim)
1550 @staticmethod
1551 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1552 return torch.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1554 def get_current_stream(self):
1555 device = self.get_device_handle()
1556 return device.current_stream()
1558 def new_event(self):
1559 device = self.get_device_handle()
1560 return device.Event()
1562 def tree_map(self, fn, tree):
1563 return torch.utils._pytree.tree_map(fn, tree) # pylint: disable=protected-access
1565 @property
1566 def checkpoint(self):
1567 # pylint: disable=C0415
1568 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import checkpoint
1569 return checkpoint
1571 @staticmethod
1572 def recompute_handle_collector_ctx():
1573 # pylint: disable=C0415
1574 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import recompute_handle_collector_ctx
1575 return recompute_handle_collector_ctx()
1577 @staticmethod
1578 def recompute_handle(handle, session_id):
1579 # pylint: disable=C0415
1580 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import recompute_handle
1581 return recompute_handle(handle, session_id)
1583 @staticmethod
1584 def recompute_session_ctx(session_id, retain_on_unpack=False):
1585 # pylint: disable=C0415
1586 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import recompute_session_ctx
1587 return recompute_session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack)
1589 @staticmethod
1590 def clear_recompute_session(session_id):
1591 # pylint: disable=C0415
1592 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import clear_recompute_session
1593 return clear_recompute_session(session_id)
1595 @staticmethod
1596 def checkpoint_wrapper(module, **checkpoint_kwargs):
1597 # pylint: disable=C0415
1598 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1599 return ckpt_wrapper(module, **checkpoint_kwargs)
1601 @staticmethod
1602 def checkpoint_exclude_wrapper(module: Any, *, save_output: bool = True) -> Any:
1603 """Wrap a module or callable whose activations should not be recomputed.
1605 Args:
1606 module: PyTorch Module or callable to exclude from checkpoint replay.
1607 save_output: Whether to retain the excluded region output for replay.
1609 Returns:
1610 The platform-specific checkpoint exclusion wrapper.
1611 """
1612 # pylint: disable=C0415
1613 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint_exclude_wrapper import (
1614 checkpoint_exclude_wrapper,
1615 )
1616 return checkpoint_exclude_wrapper(module, save_output=save_output)
1618 @staticmethod
1619 def swap_wrapper(module, policy_fn=None, group_swap=False):
1620 # pylint: disable=C0415
1621 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_wrapper
1622 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1624 @staticmethod
1625 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1626 # pylint: disable=C0415
1627 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_tensor_wrapper
1628 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1630 @staticmethod
1631 def get_class_activation_wrapper():
1632 # pylint: disable=C0415
1633 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import ActivationWrapper
1634 return ActivationWrapper
1636 @property
1637 def noop_context_fn(self):
1638 return noop_context_fn
1640 @staticmethod
1641 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1642 # pylint: disable=C0415
1643 from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts
1644 return create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation, group_swap)
1646 @staticmethod
1647 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1648 # pylint: disable=C0415
1649 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import AsyncSaveOnCpu
1650 return AsyncSaveOnCpu(policy_fn, group_swap=group_swap)
1652 @staticmethod
1653 def get_element_size(tensor):
1654 """Get Tensor Element Size"""
1655 return tensor.element_size()
1657 @staticmethod
1658 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1659 """Allocate an uninitialized 1-D tensor buffer."""
1660 if pin_memory:
1661 return torch.empty(numel, dtype=dtype, device='cpu', pin_memory=True)
1662 return torch.empty(numel, dtype=dtype, device=device)
1664 @staticmethod
1665 def tensor_to_numpy(tensor) -> np.ndarray:
1666 """Convert PyTorch tensor to numpy array."""
1667 return tensor.cpu().numpy()
1669 @staticmethod
1670 def from_numpy(np_array):
1671 """Create a host (CPU) PyTorch tensor from a numpy array."""
1672 return torch.from_numpy(np_array)
1674 @staticmethod
1675 def clip_grad_norm_(
1676 parameters, max_norm, norm_type=2.0,
1677 error_if_nonfinite=False, foreach=None,
1678 ):
1679 # pylint: disable=C0415
1680 from hyper_parallel.platform.torch.clip_grad import (
1681 clip_grad_norm_ as _clip_grad_norm,
1682 )
1683 return _clip_grad_norm(
1684 parameters, max_norm, norm_type,
1685 error_if_nonfinite=error_if_nonfinite, foreach=foreach,
1686 )
1688 @staticmethod
1689 def profiler_record(name):
1690 """Profiler context manager for recording operations using torch.profiler."""
1691 return torch.profiler.record_function(name)
1693 def cast_fp_tensor(self, dtype, x):
1694 """
1695 Cast floating-point tensor to target dtype if applicable.
1696 """
1697 if (
1698 not isinstance(x, torch.Tensor)
1699 or not torch.is_floating_point(x)
1700 or x.dtype == dtype
1701 ):
1702 return x
1703 return x.to(dtype)
1705 def apply_to_tensors(self, fn, container):
1706 """Recursively apply to all tensor in different kinds of container types."""
1708 def apply(x):
1710 if isinstance(x, torch.Tensor):
1711 return fn(x)
1712 if hasattr(x, "__dataclass_fields__"):
1713 dc = dataclasses.replace(x)
1714 changes = {
1715 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
1716 }
1717 return dataclasses.replace(dc, **changes)
1718 if isinstance(x, OrderedDict):
1719 od = x.__class__()
1720 for key, value in x.items():
1721 od[key] = apply(value)
1722 return od
1723 if isinstance(x, PackedSequence):
1724 apply(x.data)
1725 return x
1726 if isinstance(x, dict):
1727 return {key: apply(value) for key, value in x.items()}
1728 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
1729 res = (apply(el) for el in x)
1730 return type(x)(*res)
1731 if isinstance(x, (list, tuple, set)):
1732 return type(x)(apply(el) for el in x)
1733 return x
1735 return apply(container)
1738 @property
1739 def meta_device(self):
1740 return torch.device("meta")
1742 def init_on_device(self, device, include_buffers=False):
1743 return _init_on_device(device, include_buffers=include_buffers)
1745 def str_to_dtype(self, dtype_str: str) -> torch.dtype:
1746 """Map ``torch.<type>`` strings from checkpoint metadata to ``torch.dtype``."""
1747 parts = dtype_str.split(".", 1)
1748 if len(parts) != 2:
1749 raise ValueError(
1750 f"Expected dtype string like 'torch.float32', got {dtype_str!r}."
1751 )
1752 prefix, name = parts
1753 if prefix != "torch":
1754 raise ValueError(
1755 f"Expected PyTorch dtype string with prefix 'torch', got {dtype_str!r}."
1756 )
1757 dtype = getattr(torch, name)
1758 if isinstance(dtype, torch.dtype):
1759 return dtype
1760 raise ValueError(f"{dtype_str!r} does not resolve to a torch.dtype.")
1762 def list_to_size(self, size_list: list[int]) -> torch.Size:
1763 return torch.Size(size_list)