Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / platform.py: 63%
852 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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, Callable, 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 torch.compiler.is_compiling():
480 return x.contiguous()
481 if not x.is_contiguous() or x.storage_offset() != 0:
482 return x.contiguous()
483 return x
486class _TorchBatchP2PWork:
487 """Single ``.wait()`` handle wrapping the per-op works returned by
488 ``torch.distributed.batch_isend_irecv``.
490 Torch returns one ``Work`` per op in the batch (the ops are coalesced
491 onto one comm stream), whereas the platform contract — and the scheduler
492 that consumes it — expects a single handle covering the whole batch so
493 the wait can be deferred to one consumption point (mirroring MindSpore's
494 single packaging ``CommHandle``). Waiting this handle waits every
495 underlying op.
496 """
498 __slots__ = ("_works",)
500 def __init__(self, works):
501 self._works = works
503 def wait(self):
504 for work in self._works:
505 if work is not None:
506 work.wait()
509# pylint: disable=C0103
510class TorchPlatform(Platform):
511 """Torch platform api"""
512 Tensor = Tensor
513 tensor = torch.tensor
514 Parameter = Parameter
515 Module = Module
516 DTensorBase = DTensorBase
517 PipelineStageBase = PipelineStageBase
518 platform_type = PlatformType.PYTORCH
519 tensor_dtype = torch
520 dtype = torch.dtype
521 Function = torch.autograd.Function
523 _custom_ops_cls = None
525 @property
526 def custom_ops(self):
527 """Return the Torch platform custom ops instance.
529 .. warning::
530 This is an experimental API that subject to change or deletion.
532 Returns:
533 TorchCustomOps: Custom ops class that raises NotImplementedError
534 for all operators (MindSpore-only at this time).
535 """
536 if self._custom_ops_cls is None:
537 from hyper_parallel.platform.torch.custom_ops import TorchCustomOps # pylint: disable=import-outside-toplevel
538 self._custom_ops_cls = TorchCustomOps
539 return self._custom_ops_cls
541 @staticmethod
542 def get_swap_optimizer():
543 """Return the Torch optimizer-state swap wrapper class."""
544 from hyper_parallel.platform.torch.swap_optimizer.swap_optimizer import ( # pylint: disable=import-outside-toplevel
545 get_swap_optimizer,
546 )
547 return get_swap_optimizer()
549 @staticmethod
550 def is_linear_module(module) -> bool:
551 """Check whether *module* is a ``torch.nn.Linear`` instance."""
552 return isinstance(module, nn.Linear)
554 @staticmethod
555 def is_embedding_module(module) -> bool:
556 """Check whether *module* is a ``torch.nn.Embedding`` instance."""
557 return isinstance(module, nn.Embedding)
559 @staticmethod
560 def device_count(device_handle):
561 """
562 Get the number of available devices.
564 Args:
565 device_handle: The device handle (e.g., torch.cuda, torch.npu).
567 Returns:
568 int: The number of available devices.
569 """
570 return device_handle.device_count()
572 def device_type(self):
573 """
574 Get the current device type.
576 Returns:
577 str: The device type string ("npu" for NPU, "cuda" for GPU).
578 """
579 device_handle = self.get_device_handle()
580 if device_handle == torch.npu:
581 return "npu"
582 return "cuda"
584 def device(self, device_idx=None):
585 """
586 Get a torch.device object for the specified device index.
588 Args:
589 device_idx (Optional[int]): The device index. If None, returns device without index.
591 Returns:
592 torch.device: A torch device object.
593 """
594 device_type = self.device_type()
595 if device_idx is None:
596 return torch.device(device_type)
597 return torch.device(f"{device_type}:{device_idx:d}")
599 @staticmethod
600 def get_rng_state(device=None, device_handle=None):
601 """
602 Get the random number generator state.
604 Args:
605 device (Optional): The device to get RNG state from.
606 device_handle (Optional): The device handle (torch.cuda, torch.npu, etc.).
608 Returns:
609 Tensor: The RNG state as a byte tensor.
610 """
611 if device_handle is None:
612 return torch.get_rng_state()
613 if device is None:
614 return device_handle.get_rng_state()
615 return device_handle.get_rng_state(device)
617 @staticmethod
618 def set_rng_state(state, device=None, device_handle=None):
619 """
620 Set the random number generator state.
622 Args:
623 state (Tensor): The RNG state to set.
624 device (Optional): The device to set RNG state for.
625 device_handle (Optional): The device handle (torch.cuda, torch.npu, etc.).
626 """
627 if device_handle is None:
628 return torch.set_rng_state(state)
629 if device is None:
630 return device_handle.set_rng_state(state)
631 return device_handle.set_rng_state(state, device)
633 @staticmethod
634 def manual_seed(seed):
635 """
636 Set the random seed for reproducibility.
638 Args:
639 seed (int): The random seed value.
641 Returns:
642 torch.Generator: The random number generator.
643 """
644 return torch.manual_seed(seed)
646 @staticmethod
647 def ones(size, dtype=None):
648 """
649 Create a tensor filled with ones.
651 Args:
652 size (tuple): The shape of the output tensor.
653 dtype (Optional[torch.dtype]): The desired data type.
655 Returns:
656 Tensor: A tensor filled with ones.
657 """
658 return torch.ones(size, dtype=dtype)
660 @staticmethod
661 def zeros(size, dtype=None, device=None):
662 """
663 Create a tensor filled with zeros.
665 Args:
666 size (tuple): The shape of the output tensor.
667 dtype (Optional[torch.dtype]): The desired data type.
668 device (Optional[torch.device]): The device to create the tensor on.
670 Returns:
671 Tensor: A tensor filled with zeros.
672 """
673 return torch.zeros(size, dtype=dtype, device=device)
675 @staticmethod
676 def full(size, fill_value, dtype=None):
677 """
678 Create a tensor filled with a scalar value.
680 Args:
681 size (tuple): The shape of the output tensor.
682 fill_value (scalar): The value to fill the tensor with.
683 dtype (Optional[torch.dtype]): The desired data type.
685 Returns:
686 Tensor: A tensor filled with the specified value.
687 """
688 return torch.full(size, fill_value, dtype=dtype)
690 @staticmethod
691 def empty(size, dtype=None, device=None):
692 """
693 Create an uninitialized tensor.
695 Args:
696 size (tuple): The shape of the output tensor.
697 dtype (Optional[torch.dtype]): The desired data type.
698 device (Optional[torch.device or str]): Target device. When
699 ``None`` the tensor is allocated on the default device
700 (CPU under PyTorch defaults), matching the original
701 back-compat behavior.
703 Returns:
704 Tensor: An uninitialized tensor.
705 """
706 return torch.empty(size, dtype=dtype, device=device)
708 @staticmethod
709 def rand(size, dtype=None, device=None):
710 """Create a tensor filled with uniform random values in ``[0, 1)``."""
711 return torch.rand(size, dtype=dtype, device=device)
713 @staticmethod
714 def randn(size, dtype=None, device=None):
715 """Create a tensor filled with standard-normal random values."""
716 return torch.randn(size, dtype=dtype, device=device)
718 @staticmethod
719 def get_rank():
720 """
721 Get the rank of the current process in the distributed group.
723 Returns:
724 int: The rank of the current process.
725 """
726 return dist.get_rank()
728 @staticmethod
729 def get_global_rank(group, group_rank):
730 """
731 Get the global rank from a group rank.
733 Args:
734 group (ProcessGroup): The process group.
735 group_rank (int): The rank within the group.
737 Returns:
738 int: The global rank.
739 """
740 return dist.get_global_rank(group, group_rank)
742 @staticmethod
743 def get_group_rank(group):
744 """Return this process's rank within *group*."""
745 return dist.get_group_rank(group, dist.get_rank())
747 @staticmethod
748 def get_world_size():
749 """
750 Get the total number of processes in the distributed group.
752 Returns:
753 int: The world size.
754 """
755 return dist.get_world_size()
757 @staticmethod
758 def get_param_local_shape(param):
759 """
760 Get the local shape of a parameter, handling both regular and distributed tensors.
762 Args:
763 param (Union[Tensor, DTensorBase]): The parameter tensor.
765 Returns:
766 torch.Size: The local shape of the parameter.
767 """
768 if isinstance(param, DTensorBase):
769 return param.local_shape
770 return param.shape
772 @staticmethod
773 def get_param_local_data(param):
774 """
775 Get the local data of a parameter, handling both regular and distributed tensors.
777 Args:
778 param (Union[Tensor, DTensorBase]): The parameter tensor.
780 Returns:
781 Tensor: The local tensor data.
782 """
783 if isinstance(param, DTensorBase):
784 return param.to_local()
785 return param
787 @staticmethod
788 def update_param_data(param, data):
789 """
790 Update the data of a parameter.
792 Args:
793 param (Parameter): The parameter to update.
794 data (Tensor): The new data tensor.
795 """
796 param.data = data
798 @staticmethod
799 def load_into_param(param, data):
800 """Load tensor *data* into *param* (plain tensor or DTensor)."""
801 if isinstance(param, DTensorBase):
802 local = param._local_tensor # pylint: disable=W0212
803 if local.is_meta:
804 # Meta tensor materialisation: replace the placeholder.
805 orig_requires_grad = param.requires_grad
806 param._local_tensor = data # pylint: disable=W0212
807 if data.requires_grad != orig_requires_grad:
808 param.requires_grad_(orig_requires_grad)
809 else:
810 local.copy_(data)
811 else:
812 param.copy_(data)
814 @staticmethod
815 def get_op_name(func):
816 """
817 Extract the operation name from various function types.
819 Args:
820 func: The function or operation to extract the name from.
822 Returns:
823 str: The operation name.
824 """
825 if hasattr(func, "__name__"):
826 return func.__name__
827 if isinstance(func, OpOverload):
828 full_name = func.name
829 core_name = full_name.split("::")[-1].split(".")[0]
830 return core_name
831 if isinstance(func, OpOverloadPacket):
832 return func.name.split("::")[-1]
833 func_str = str(func)
834 if "built-in function" in func_str:
835 return func_str.split()[-1].strip(">")
836 if "function" in func_str:
837 return func_str.split()[1]
838 return "unknown_op"
840 @staticmethod
841 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
842 data = _ensure_contiguous(data)
843 output = list(dist_func.all_gather(data, group=group))
844 if rank_list is not None:
845 group_ranks = dist.get_process_group_ranks(group)
846 if tuple(rank_list) != tuple(group_ranks):
847 rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}
848 output = [output[rank_to_idx[int(rank)]] for rank in rank_list]
849 return torch.cat(output, dim=concat_dim)
851 @staticmethod
852 def chunk(data, split_dim, split_size, index):
853 return torch.chunk(data, split_size, dim=split_dim)[index]
855 @staticmethod
856 def differentiable_all_to_all(input_data, output_shape, group):
857 input_data = _ensure_contiguous(input_data)
858 output_tensor = torch.empty(output_shape, device=input_data.device, dtype=input_data.dtype)
859 output_tensor = dist_func.all_to_all_single(
860 output_tensor,
861 input_data,
862 group=group
863 )
864 return output_tensor
866 @staticmethod
867 def tensor_type_cast(input_data, cast_type):
868 """Cast tensor to specified data type."""
869 type_mapping = {
870 'float32': torch.float32,
871 'float16': torch.float16,
872 'int64': torch.int64,
873 'int32': torch.int32
874 }
875 if cast_type not in type_mapping:
876 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")
877 return input_data.to(type_mapping[cast_type])
879 @staticmethod
880 def differentiable_all_reduce(data, op, group):
881 data = _ensure_contiguous(data)
882 # Resolve the op from string to ReduceOp enum if necessary
883 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
884 return dist_func.all_reduce(data, op=reduce_op, group=group)
886 @staticmethod
887 def get_cell_construct(cell):
888 return cell.forward
890 @staticmethod
891 def get_cells_and_names(cell):
892 return cell.named_modules()
894 @staticmethod
895 def get_modules(module):
896 return module.modules()
898 @staticmethod
899 def search_parameter_by_name(cell, param_name: str):
900 """
901 Find the parent Module of the parameter, the parameter's name in the parent Module, and the parameter.
902 Return value: (parent Module instance, parameter's name in parent Module, parameter object).
903 Returns None if not found.
904 """
905 # Remove the "self." prefix from param_name
906 param_name = param_name.replace("self.", "")
907 # Case 1: The parameter is a direct parameter of the current Module
908 if param_name in cell._parameters: # pylint: disable=protected-access
909 return (cell, param_name, cell._parameters[param_name]) # pylint: disable=protected-access
911 # Case 2: The parameter is in a sub-Module
912 if "." in param_name:
913 cell_path, param_key = param_name.rsplit(".", 1)
914 try:
915 # Locate the sub-Module where the parameter resides (supports multi-level paths)
916 target_cell = cell.get_submodule(cell_path)
917 # Check if the sub-Module directly contains this parameter
918 if param_key in target_cell._parameters: # pylint: disable=protected-access
919 return target_cell, param_key, target_cell._parameters[param_key] # pylint: disable=protected-access
920 except AttributeError:
921 pass
923 # Traverse all sub-Modules (recursively) to search for the parameter
924 for _, child_cell in cell.named_children():
925 if isinstance(child_cell, Module):
926 result = TorchPlatform.search_parameter_by_name(child_cell, param_name)
927 if result is not None:
928 return result
930 return None
932 @staticmethod
933 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
934 """
935 Modify the original parameter in a Module or sub-Module using the search result
936 """
937 parent_cell, param_key, _ = result
938 # Key operation: directly modify the _parameters dictionary.
939 if param_key in parent_cell._parameters: # pylint: disable=protected-access
940 parent_cell._parameters[param_key] = new_param # pylint: disable=protected-access
941 else:
942 parent_cell.register_parameter(param_key, new_param)
943 return True
945 @staticmethod
946 def set_layout_into_parameter(param, layout):
947 """Set layout into parameter"""
948 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel
949 from hyper_parallel.core.dtensor.layout import _get_slice_tensor_by_layout # pylint: disable=import-outside-toplevel
950 if isinstance(param, DTensor):
951 raise ValueError(f"Parameter {param} has been configured layout, cannot be set repeatedly.")
952 requires_grad = param.requires_grad
953 param_dtensor = DTensor.from_local(
954 _get_slice_tensor_by_layout(param, layout),
955 layout.mesh, layout.alias_placements)
956 new_param = Parameter(param_dtensor, requires_grad=requires_grad)
957 return new_param
959 @staticmethod
960 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
961 data = _ensure_contiguous(data)
962 input_tuple = torch.chunk(data, dev_num, dim=axis)
963 output_tensor = torch.empty(input_tuple[0].shape, device=data.device, dtype=data.dtype)
965 # Resolve the op from string to ReduceOp enum
966 reduce_op = _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op
968 output_tensor = dist_func.reduce_scatter(output_tensor, input_tuple, op=reduce_op, group=group)
970 # Keep manual handling for 'avg' string as it maps to SUM in _OP_MAP
971 if op == 'avg':
972 output_tensor = output_tensor / dev_num
973 return output_tensor
975 @staticmethod
976 def get_device_handle(device_type: str = "npu"):
977 """Return the torch device module (e.g. ``torch.npu`` or ``torch.cuda``) for the given device type."""
978 try:
979 handle = getattr(torch, device_type)
980 except AttributeError as e:
981 raise RuntimeError(f"TorchPlatform expect got device handle: 'torch.{device_type}' failed.") from e
982 return handle
984 @staticmethod
985 def get_param_type_size(param):
986 # pylint: disable=W0212
987 return torch._utils._element_size(param.dtype)
989 @staticmethod
990 def is_tensor(obj: Any) -> bool:
991 """Return True if ``obj`` is a ``torch.Tensor``."""
992 return isinstance(obj, Tensor)
994 @staticmethod
995 def get_tensor_storage_size(tensor: Any) -> int:
996 """Return serialized byte size (numel * element size) for a PyTorch tensor."""
997 if not TorchPlatform.is_tensor(tensor):
998 raise TypeError(
999 f"TorchPlatform.get_tensor_storage_size expects torch.Tensor, got {type(tensor)!r}"
1000 )
1001 return int(tensor.numel()) * int(tensor.element_size())
1003 @staticmethod
1004 def parameters_dict(cell: Module):
1005 return cell.named_parameters()
1007 @staticmethod
1008 def buffers_dict(cell: Module) -> Any:
1009 """Return all named buffers registered by the module tree."""
1010 return cell.named_buffers()
1012 @staticmethod
1013 def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
1014 """Get the state dictionary of a model.
1016 Delegates to torch-specific implementation that handles DTensor
1017 gathering, CPU offloading and frozen-parameter filtering.
1018 """
1019 # pylint: disable=C0415
1020 from hyper_parallel.platform.torch.fully_shard.state_dict_utils import (
1021 get_model_state_dict as _get_model_state_dict,
1022 )
1023 return _get_model_state_dict(model, options=options)
1025 @staticmethod
1026 def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
1027 """Set the state dictionary of a model.
1029 Delegates to torch-specific implementation that scatters full tensors
1030 into DTensor shards and performs an in-place load.
1031 """
1032 # pylint: disable=C0415
1033 from hyper_parallel.platform.torch.fully_shard.state_dict_utils import (
1034 set_model_state_dict as _set_model_state_dict,
1035 )
1036 return _set_model_state_dict(model, model_state_dict, options=options)
1038 @staticmethod
1039 def save_checkpoint(cell: Module, file_path: str, ckpt_format: str = "safetensors") -> None:
1040 if ckpt_format == "safetensors":
1041 save_file(tensors=cell, filename=file_path)
1042 else:
1043 torch.save(obj=cell, f=file_path)
1045 @staticmethod
1046 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
1047 if ckpt_format == "safetensors":
1048 return load_file(filename=file_path)
1049 return torch.load(f=file_path)
1051 @staticmethod
1052 def new_zero_parameter(param_shape, param_type, requires_grad, device):
1053 return nn.Parameter(torch.zeros(param_shape, dtype=param_type, device=device), requires_grad=requires_grad)
1055 @staticmethod
1056 def new_tensor(tensor_shape, tensor_type, device):
1057 return torch.empty(size=tensor_shape, dtype=tensor_type, device=device)
1059 @staticmethod
1060 def full_like(tensor, fill_value, dtype=None):
1061 return torch.full_like(tensor, fill_value, dtype=dtype)
1063 @staticmethod
1064 def set_tensor_requires_grad(input_tensor):
1065 """
1066 set requires grad flag for input tensor, only effective for leaf node
1067 """
1068 if input_tensor.is_leaf:
1069 input_tensor.requires_grad = True
1071 def _create_group(self, rank_list):
1072 normalized_rank_list = tuple(sorted(rank_list))
1073 world_rank_list = tuple(range(self.get_world_size()))
1074 if normalized_rank_list == world_rank_list:
1075 group = _get_default_group()
1076 EXISTING_COMM_GROUPS[str(normalized_rank_list)] = group
1077 return group
1078 group_dict = create_sub_groups(rank_list)
1079 return group_dict[normalized_rank_list]
1081 @staticmethod
1082 def all_gather_into_tensor(data, group_info, async_op=False):
1083 output_shape = list(data.shape)
1084 output_shape[0] = output_shape[0] * group_info.rank_size
1085 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
1086 handle = dist.all_gather_into_tensor(output, data, group=group_info.group, async_op=async_op)
1087 return output, handle
1089 @staticmethod
1090 def all_gather_single(input_tensor, output_shape, group, async_op=False):
1091 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
1092 handle = dist.all_gather_into_tensor(output, input_tensor, group=group, async_op=async_op)
1093 return output, handle
1095 @staticmethod
1096 def all_reduce(data, group_info, async_op=False):
1097 if not data.is_contiguous():
1098 data = data.contiguous()
1099 handle = dist.all_reduce(data, group=group_info.group, async_op=async_op)
1100 return data, handle
1102 @staticmethod
1103 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
1104 if group_src is not None:
1105 src = dist.get_global_rank(group, group_src)
1106 handle = dist.broadcast(data, src, group, async_op)
1107 if async_op and handle is not None:
1108 handle.wait()
1110 @staticmethod
1111 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
1112 if group_src is not None:
1113 src = dist.get_global_rank(group, group_src)
1114 handle = dist.scatter(output, scatter_list, src=src, group=group, async_op=async_op)
1115 if async_op and handle is not None:
1116 handle.wait()
1117 return output
1119 @staticmethod
1120 def isend(tensor, dst=None, group=None, tag=0):
1121 return dist.isend(tensor, dst, group, tag)
1123 @staticmethod
1124 def irecv(tensor, src=None, group=None, tag=0):
1125 return dist.irecv(tensor, src, group, tag)
1127 @staticmethod
1128 def p2p_op(op_type, tensor, peer, group=None):
1129 # torch's P2POp takes the op callable (dist.isend / dist.irecv), not
1130 # the "isend"/"irecv" string the stage specs builders emit.
1131 if op_type == "isend":
1132 op = dist.isend
1133 elif op_type == "irecv":
1134 op = dist.irecv
1135 else:
1136 raise ValueError(
1137 f"p2p_op op_type must be 'isend' or 'irecv', but got {op_type!r}."
1138 )
1139 return dist.P2POp(op, tensor, peer, group)
1141 @staticmethod
1142 def batch_isend_irecv(p2p_ops):
1143 """Launch a peer-batched P2P group as one coalesced op.
1145 ``torch.distributed.batch_isend_irecv`` coalesces the ops onto one
1146 comm stream and returns one ``Work`` per op; we wrap them in a single
1147 ``.wait()`` handle so a send and a recv to the same peer overlap on
1148 the duplex link and the caller can defer the whole batch's wait to one
1149 consumption point.
1150 """
1151 if not p2p_ops:
1152 return None
1153 works = dist.batch_isend_irecv(p2p_ops)
1154 return _TorchBatchP2PWork(works) if works else None
1156 @staticmethod
1157 def prepare_batch_p2p_group(group: Any = None) -> None:
1158 """Synchronize a group before its first subset batched P2P call.
1160 PyTorch requires every rank in a process group to participate when
1161 ``batch_isend_irecv`` is the first collective on that group. A barrier
1162 at the common pipeline run boundary initializes the communicator
1163 before ranks reach peer operations at different times.
1165 Args:
1166 group: The process group used by the batched P2P operations.
1167 ``None`` uses the default group.
1168 """
1169 dist.barrier(group=group)
1171 @staticmethod
1172 def p2p_exchange(tensor, peer_rank: int, group=None):
1173 if peer_rank == dist.get_rank(group):
1174 return tensor
1175 return _TorchP2PExchangeFunction.apply(tensor, peer_rank, group)
1177 @staticmethod
1178 def send_object_list(obj_list, dst=None, group=None):
1179 dist.send_object_list(obj_list, dst, group)
1181 @staticmethod
1182 def recv_object_list(obj_list, src=None, group=None):
1183 dist.recv_object_list(obj_list, src, group)
1185 @staticmethod
1186 def reduce_scatter_tensor(data, group_info, async_op=False):
1187 output_shape = list(data.shape)
1188 output_shape[0] = output_shape[0] // group_info.rank_size
1189 output = torch.empty(output_shape, dtype=data.dtype, device=data.device)
1190 handle = dist.reduce_scatter_tensor(output, data, group=group_info.group, async_op=async_op)
1191 return output, handle
1193 @staticmethod
1194 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1195 output = torch.empty(output_shape, dtype=input_tensor.dtype, device=input_tensor.device)
1196 handle = dist.reduce_scatter_tensor(output, input_tensor, group=group, async_op=async_op)
1197 return output, handle
1199 @staticmethod
1200 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1201 output = torch.empty(output_shape, device=input_tensor.device, dtype=input_tensor.dtype)
1202 work = dist.all_to_all_single(output, input_tensor, group=group, async_op=async_op)
1203 return output, work
1205 @staticmethod
1206 def differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group):
1207 """Variable-split all-to-all with autograd support for EP token dispatch/combine."""
1208 out_total = sum(output_splits)
1209 output = torch.empty(
1210 out_total, *input_tensor.shape[1:],
1211 dtype=input_tensor.dtype, device=input_tensor.device,
1212 )
1213 output = dist_func.all_to_all_single(
1214 output, input_tensor,
1215 output_split_sizes=output_splits,
1216 input_split_sizes=input_splits,
1217 group=group,
1218 )
1219 return output
1221 @staticmethod
1222 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1223 """Truly-async variant of :meth:`differentiable_all_to_all_single`.
1225 Both forward AND backward return :class:`AsyncCollectiveTensor`,
1226 so the ``wait_tensor`` op is queued lazily — only when a downstream
1227 kernel actually reads the result.
1229 Why both directions need lazy wait:
1231 * FWD: ACT lazy wait lets host return immediately and the paired
1232 BWD thread's compute kernel slip into the queue before the wait.
1233 * BWD: PyTorch's stock backward issues ``wait_tensor`` eagerly,
1234 and the autograd engine binds backward stream to the forward
1235 stream — so even running BWD inside a ``with torch.npu.stream
1236 (side_stream)`` context does not move that wait off the main
1237 stream. Returning ACT from backward defers the wait to the
1238 next backward op's first consumption, opening a small window
1239 during which FWD's Attention kernels can be queued onto the
1240 main stream **before** the wait lands.
1242 Args:
1243 input_tensor: Input tensor, split along dim 0 by ``input_splits``.
1244 input_splits: ``list[int]`` — rows sent to each rank.
1245 output_splits: ``list[int]`` — rows received from each rank.
1246 group: Process group.
1248 Returns:
1249 ``AsyncCollectiveTensor`` of shape
1250 ``[sum(output_splits), *input_tensor.shape[1:]]``.
1251 """
1252 return _AsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1254 @staticmethod
1255 def differentiable_variable_all_gather(
1256 input_tensor: Tensor, output_splits: Sequence[int], group: Any) -> Tensor:
1257 """Gather variable dim-zero shards on HCCL or Gloo with autograd support."""
1258 return _TorchDifferentiableVariableAllGather.apply(
1259 input_tensor, tuple(output_splits), group
1260 )
1262 @staticmethod
1263 def wait_async_tensor(tensor):
1264 """Wait for an async collective tensor to become materialised.
1266 Idempotent — calling on an already-waited tensor is a no-op.
1268 Args:
1269 tensor: ``AsyncCollectiveTensor`` whose device-side values may
1270 not yet be ready.
1272 Returns:
1273 The same *tensor*, now fully materialised.
1274 """
1275 from torch.distributed._functional_collectives import wait_tensor # pylint: disable=C0415
1276 wait_tensor(tensor)
1277 return tensor
1279 @staticmethod
1280 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1281 handle_box=None):
1282 """Wait async all-gather handle and reconstruct result (differentiable)."""
1283 return _TorchAsyncAllGatherFunction.apply(
1284 x, work, out_perm, group, world_size, gather_dim, handle_box
1285 )
1287 @staticmethod
1288 def arange(start, end=None, step=1, dtype=None, device=None):
1289 """Create a 1-D tensor with evenly spaced values."""
1290 if end is None:
1291 return torch.arange(start, dtype=dtype, device=device)
1292 return torch.arange(start, end, step, dtype=dtype, device=device)
1294 @staticmethod
1295 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim,
1296 handle_box=None):
1297 """Wait async A2A handle and reconstruct result (differentiable).
1299 Args:
1300 x: Input tensor.
1301 work: Async work handle from all_to_all.
1302 out_perm: Output buffer from all_to_all.
1303 group: Process group.
1304 world_size: World size.
1305 concat_dim: Dimension for concatenation.
1306 split_dim: Dimension for split.
1307 handle_box: Optional mutable list; backward appends (work, out_perm) here.
1308 """
1309 return _TorchAsyncA2AFunction.apply(
1310 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1311 )
1313 @staticmethod
1314 def differentiable_sync_hook(x, hook_name: str, coordinator):
1315 """Identity op that fires coordinator rendezvous on forward and backward.
1317 Always goes through ``_TorchSyncHookFunction.apply`` so that the
1318 autograd graph **records a SyncHook node regardless of whether the
1319 coordinator is currently enabled**. Skipping ``apply`` when
1320 disabled would leave warmup-forwarded graphs without the hook
1321 nodes, and a later ``overlap.run`` — whose BWD thread back-props
1322 such a graph — would then traverse zero hooks while the paired FWD
1323 thread (whose current forward DOES record hooks) waits at a
1324 barrier for a partner that never arrives.
1326 Args:
1327 x: Input tensor.
1328 hook_name: One of:
1329 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1330 full rendezvous on both directions.
1331 * ``"D_LAST"`` — closing D of the last MoE
1332 layer in a chunk. Forward: ``notify_dispatched``
1333 only (no Attention follows so rendezvous is
1334 skipped). Backward: pure skip (first BWD
1335 hook to fire; combine.bwd has already
1336 dispatched freely).
1337 coordinator: A :class:`HookCoordinator` instance.
1338 """
1339 return _TorchSyncHookFunction.apply(x, hook_name, coordinator)
1341 @staticmethod
1342 def get_tensor_transform():
1343 raise NotImplementedError("Unsupported get_tensor_transform for torch platform")
1345 @staticmethod
1346 def construct_strided_slice(x, begin, end, stride):
1347 raise NotImplementedError("Unsupported construct_strided_slice for torch platform")
1349 @staticmethod
1350 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1351 # pylint: disable=C0415
1352 from hyper_parallel.platform.torch.pipeline_parallel._utils import _MicroBatch
1353 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1355 @staticmethod
1356 def get_symmetric_memory_handler():
1357 # pylint: disable=C0415
1358 from hyper_parallel.platform.torch.symmetric_memory import TorchSymmetricMemoryHandler
1359 symmetric_memory = TorchSymmetricMemoryHandler()
1360 return symmetric_memory
1362 @staticmethod
1363 def get_multicore_handler():
1364 """Return a TorchMulticoreHandler instance for multi-core device management."""
1365 # pylint: disable=C0415
1366 from hyper_parallel.platform.torch.multicore import TorchMulticoreHandler
1367 return TorchMulticoreHandler()
1369 def new_stream(self):
1370 device = self.get_device_handle()
1371 return device.Stream()
1373 def get_stream_context(self):
1374 device = self.get_device_handle()
1375 return device.stream
1377 @staticmethod
1378 def all_gather_object(object_list, obj, group=None) -> None:
1379 """
1380 Gathers objects from the given group into object list.
1382 Args:
1383 object_list (list[Any]): Define the output list, which size equal to the size of group.
1384 obj (Any): The object on current rank and in given process group.
1385 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1386 global group.
1388 Returns:
1389 None. Objs are gathered into ``object_list``.
1390 """
1391 dist.all_gather_object(object_list, obj, group)
1393 @staticmethod
1394 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1395 """
1396 Synchronize all processes in the given process group.
1398 Args:
1399 group (ProcessGroup, optional): The process group to work on. Default is ``None``,
1400 meaning the default process group.
1401 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1402 device_ids (list[int], optional): Device ids for backends that require a device for
1403 barrier (e.g. NCCL). Default: ``None``.
1405 Returns:
1406 Async work handle if ``async_op`` is True; otherwise ``None``.
1407 """
1408 return dist.barrier(group, async_op, device_ids)
1410 @staticmethod
1411 def init_process_group(
1412 backend: Optional[str] = None,
1413 *,
1414 init_method: Optional[str] = None,
1415 timeout: Optional[timedelta] = None,
1416 world_size: int = -1,
1417 rank: int = -1,
1418 store: Optional[Store] = None,
1419 pg_options: Optional[Any] = None,
1420 device_id: Optional[Union[torch.device, int]] = None,
1421 ) -> None:
1422 """
1423 Initialize global process group.
1425 Args:
1426 backend (str or Backend, optional): The backend to use for distributed communication.
1427 init_method (str, optional): URL specifying how to initialize the process group. Default is "env://",
1428 can not be specified at the same time with ``store``.
1429 timeout (timedelta, optional): Timeout for process group. Default 10 minutes for NCCL and for other
1430 backends 30 minutes.
1431 world_size (int, optional): Number of processes. If ``store`` is specified, world_size is required.
1432 rank (int, optional): Rank of the current process, which value must between 0 and ``world_size``-1. If
1433 ``store`` is specified, rank is required.
1434 store (Store, optional): Key/value store accessible to all workers, used to exchange connection/address
1435 information. Can not be specified at the same time with ``init_method``.
1436 pg_options (ProcessGroupOptions, optional): Extra options to pass during constructing process groups.
1437 device_id (torch.device | int, optional): Specific device this process will work on.
1438 """
1439 try:
1440 _get_default_group()
1441 # except multi version error
1442 except (ValueError, RuntimeError):
1443 if backend is None:
1444 backend = "hccl"
1445 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1446 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1448 @staticmethod
1449 def destroy_process_group(group: Optional[ProcessGroup] = None) -> None:
1450 """
1451 Destroy given process group.
1453 Args:
1454 group (ProcessGroup, optional): Given process group will be destroyed, if not given, all process groups
1455 will be destroyed.
1456 """
1457 group = group or _get_default_group()
1458 if group in EXISTING_COMM_GROUPS.values():
1459 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1460 for k in keys_to_destroy:
1461 del EXISTING_COMM_GROUPS[k]
1462 dist.destroy_process_group(group)
1464 @staticmethod
1465 def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> list[int]:
1466 """
1467 Get all ranks relative to given process group.
1469 Args:
1470 group (Optional[ProcessGroup]): Process group worked on. Default is ``None``, and ``None`` means global
1471 group.
1473 Returns:
1474 Rank list.
1475 """
1476 group = group or _get_default_group()
1477 return dist.get_process_group_ranks(group)
1479 @staticmethod
1480 def get_backend(group: Optional[ProcessGroup] = None) -> Backend:
1481 """
1482 Get the backend of the given process group.
1484 Args:
1485 group (ProcessGroup, optional): Process group worked on. Default is ``None``, and ``None`` means global
1486 group.
1488 Returns:
1489 The backend object of the given process group.
1490 """
1491 group = group or _get_default_group()
1492 return dist.get_backend(group)
1494 @staticmethod
1495 def split_group(parent_pg: Optional[ProcessGroup] = None,
1496 split_ranks: Optional[list] = None,
1497 timeout: Optional[timedelta] = None,
1498 pg_options: Optional[Any] = None,
1499 group_desc: Optional[str] = None,
1500 ) -> Optional[ProcessGroup]:
1501 """
1502 Create split groups for every group rank in split_ranks, and return the split process group which relative to
1503 current rank id.
1505 Args:
1506 parent_pg (Optional[ProcessGroup]): A process group which the goal group split from.
1507 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1508 timeout (Optional[timedelta]): Timeout for process group. Default 10 minutes for NCCL and for other
1509 backend 30 minutes.
1510 pg_options (Optional[Any]): Extra options to pass during constructing process groups.
1511 group_desc (Optional[str]): Description of process group.
1513 Return:
1514 Optional[ProcessGroup]: One of split process group which relative to current rank id
1515 """
1516 if split_ranks is None or len(split_ranks) == 0:
1517 raise ValueError("split_ranks cannot be None or empty")
1519 split_group = None
1520 for split_rank in split_ranks:
1521 dist_group = TorchPlatform.get_created_group(split_rank)
1522 if dist_group is None:
1523 dist_group = dist.new_group(ranks=split_rank)
1524 EXISTING_COMM_GROUPS[str(tuple(sorted(split_rank)))] = dist_group
1525 if TorchPlatform.get_rank() in split_rank:
1526 split_group = dist_group
1528 return split_group
1530 @staticmethod
1531 def get_group_local_rank(group: ProcessGroup = None) -> int:
1532 """get group local rank id."""
1533 group = group or _get_default_group()
1534 return group.rank()
1536 @staticmethod
1537 def no_grad():
1538 return torch.no_grad()
1540 @staticmethod
1541 def preserve_version_counter(tensor):
1542 return torch.autograd._unsafe_preserve_version_counter(tensor) # pylint: disable=W0212
1544 @staticmethod
1545 def relu(tensor):
1546 return torch.relu(tensor)
1548 @staticmethod
1549 def cat(tensors, dim=0):
1550 return torch.cat(tensors, dim=dim)
1552 @staticmethod
1553 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1554 return torch.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1556 def get_current_stream(self):
1557 device = self.get_device_handle()
1558 return device.current_stream()
1560 def new_event(self):
1561 device = self.get_device_handle()
1562 return device.Event()
1564 def tree_map(self, fn, tree):
1565 return torch.utils._pytree.tree_map(fn, tree) # pylint: disable=protected-access
1567 @property
1568 def checkpoint(self):
1569 # pylint: disable=C0415
1570 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import checkpoint
1571 return checkpoint
1573 @staticmethod
1574 def recompute_handle_collector_ctx():
1575 # pylint: disable=C0415
1576 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import recompute_handle_collector_ctx
1577 return recompute_handle_collector_ctx()
1579 @staticmethod
1580 def recompute_handle(handle, session_id):
1581 # pylint: disable=C0415
1582 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import recompute_handle
1583 return recompute_handle(handle, session_id)
1585 @staticmethod
1586 def recompute_session_ctx(session_id, retain_on_unpack=False):
1587 # pylint: disable=C0415
1588 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import recompute_session_ctx
1589 return recompute_session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack)
1591 @staticmethod
1592 def clear_recompute_session(session_id):
1593 # pylint: disable=C0415
1594 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint import clear_recompute_session
1595 return clear_recompute_session(session_id)
1597 @staticmethod
1598 def is_compiling() -> bool:
1599 """Return whether execution is currently captured by ``torch.compile``."""
1600 return torch.compiler.is_compiling()
1602 @staticmethod
1603 def checkpoint_wrapper(module, **checkpoint_kwargs):
1604 # pylint: disable=C0415
1605 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1606 return ckpt_wrapper(module, **checkpoint_kwargs)
1608 @staticmethod
1609 def checkpoint_exclude_wrapper(module: Any, *, save_output: bool = True) -> Any:
1610 """Wrap a module or callable whose activations should not be recomputed.
1612 Args:
1613 module: PyTorch Module or callable to exclude from checkpoint replay.
1614 save_output: Whether to retain the excluded region output for replay.
1616 Returns:
1617 The platform-specific checkpoint exclusion wrapper.
1618 """
1619 # pylint: disable=C0415
1620 from hyper_parallel.platform.torch.activation_checkpoint.checkpoint_exclude_wrapper import (
1621 checkpoint_exclude_wrapper,
1622 )
1623 return checkpoint_exclude_wrapper(module, save_output=save_output)
1625 @staticmethod
1626 def swap_wrapper(module, policy_fn=None, group_swap=False):
1627 # pylint: disable=C0415
1628 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_wrapper
1629 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1631 @staticmethod
1632 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1633 # pylint: disable=C0415
1634 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_tensor_wrapper
1635 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1637 @staticmethod
1638 def get_class_activation_wrapper():
1639 # pylint: disable=C0415
1640 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import ActivationWrapper
1641 return ActivationWrapper
1643 @property
1644 def noop_context_fn(self):
1645 return noop_context_fn
1647 @staticmethod
1648 def ignore_sac_ops(ignore_ops: list[object | None]) -> None:
1649 # pylint: disable=C0415
1650 from hyper_parallel.platform.torch.activation_checkpoint.sac import ignore_sac_ops
1651 ignore_sac_ops(ignore_ops)
1653 @staticmethod
1654 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1655 # pylint: disable=C0415
1656 from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts
1657 return create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation, group_swap)
1659 @staticmethod
1660 def create_native_selective_checkpoint_contexts(policy_fn: Callable) -> Any:
1661 """Create Torch-native selective-checkpoint contexts for compile."""
1662 # pylint: disable=C0415
1663 from hyper_parallel.platform.torch.activation_checkpoint.native_compile import (
1664 create_native_selective_checkpoint_contexts,
1665 )
1666 return create_native_selective_checkpoint_contexts(policy_fn)
1668 @staticmethod
1669 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1670 # pylint: disable=C0415
1671 from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import AsyncSaveOnCpu
1672 return AsyncSaveOnCpu(policy_fn, group_swap=group_swap)
1674 @staticmethod
1675 def get_element_size(tensor):
1676 """Get Tensor Element Size"""
1677 return tensor.element_size()
1679 @staticmethod
1680 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1681 """Allocate an uninitialized 1-D tensor buffer."""
1682 if pin_memory:
1683 return torch.empty(numel, dtype=dtype, device='cpu', pin_memory=True)
1684 return torch.empty(numel, dtype=dtype, device=device)
1686 @staticmethod
1687 def tensor_to_numpy(tensor) -> np.ndarray:
1688 """Convert PyTorch tensor to numpy array."""
1689 return tensor.cpu().numpy()
1691 @staticmethod
1692 def from_numpy(np_array):
1693 """Create a host (CPU) PyTorch tensor from a numpy array."""
1694 return torch.from_numpy(np_array)
1696 @staticmethod
1697 def clip_grad_norm_(
1698 parameters, max_norm, norm_type=2.0,
1699 error_if_nonfinite=False, foreach=None,
1700 ):
1701 # pylint: disable=C0415
1702 from hyper_parallel.platform.torch.clip_grad import (
1703 clip_grad_norm_ as _clip_grad_norm,
1704 )
1705 return _clip_grad_norm(
1706 parameters, max_norm, norm_type,
1707 error_if_nonfinite=error_if_nonfinite, foreach=foreach,
1708 )
1710 @staticmethod
1711 def profiler_record(name):
1712 """Profiler context manager for recording operations using torch.profiler."""
1713 return torch.profiler.record_function(name)
1715 def cast_fp_tensor(self, dtype, x):
1716 """
1717 Cast floating-point tensor to target dtype if applicable.
1718 """
1719 if (
1720 not isinstance(x, torch.Tensor)
1721 or not torch.is_floating_point(x)
1722 or x.dtype == dtype
1723 ):
1724 return x
1725 return x.to(dtype)
1727 def apply_to_tensors(self, fn, container):
1728 """Recursively apply to all tensor in different kinds of container types."""
1730 def apply(x):
1732 if isinstance(x, torch.Tensor):
1733 return fn(x)
1734 if hasattr(x, "__dataclass_fields__"):
1735 dc = dataclasses.replace(x)
1736 changes = {
1737 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
1738 }
1739 return dataclasses.replace(dc, **changes)
1740 if isinstance(x, OrderedDict):
1741 od = x.__class__()
1742 for key, value in x.items():
1743 od[key] = apply(value)
1744 return od
1745 if isinstance(x, PackedSequence):
1746 apply(x.data)
1747 return x
1748 if isinstance(x, dict):
1749 return {key: apply(value) for key, value in x.items()}
1750 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
1751 res = (apply(el) for el in x)
1752 return type(x)(*res)
1753 if isinstance(x, (list, tuple, set)):
1754 return type(x)(apply(el) for el in x)
1755 return x
1757 return apply(container)
1760 @property
1761 def meta_device(self):
1762 return torch.device("meta")
1764 def init_on_device(self, device, include_buffers=False):
1765 return _init_on_device(device, include_buffers=include_buffers)
1767 def str_to_dtype(self, dtype_str: str) -> torch.dtype:
1768 """Map ``torch.<type>`` strings from checkpoint metadata to ``torch.dtype``."""
1769 parts = dtype_str.split(".", 1)
1770 if len(parts) != 2:
1771 raise ValueError(
1772 f"Expected dtype string like 'torch.float32', got {dtype_str!r}."
1773 )
1774 prefix, name = parts
1775 if prefix != "torch":
1776 raise ValueError(
1777 f"Expected PyTorch dtype string with prefix 'torch', got {dtype_str!r}."
1778 )
1779 dtype = getattr(torch, name)
1780 if isinstance(dtype, torch.dtype):
1781 return dtype
1782 raise ValueError(f"{dtype_str!r} does not resolve to a torch.dtype.")
1784 def list_to_size(self, size_list: list[int]) -> torch.Size:
1785 return torch.Size(size_list)