Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / platform.py: 44%
892 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""MindSpore platform api"""
16from datetime import timedelta
17from typing import Any, Optional, Union
18import dataclasses
19from collections import OrderedDict
21import numpy as np
22import mindspore as ms
23import mindspore.common.dtype as mstype
24from mindspore.mint.distributed import TCPStore
26from mindspore.nn import Cell
27from mindspore import mint
28from mindspore.common.api import _no_grad
29from mindspore.common._grad_function import _Function
30from mindspore.common.dtype import type_size_in_bytes
31from mindspore.common.recompute import null_context_fn
32from mindspore.common.parameter import Parameter
33from mindspore.common.tensor import Tensor
34from mindspore.common.initializer import initializer
35from mindspore.communication import GlobalComm
36from mindspore.communication import get_group_size
37from mindspore.communication import create_group as new_group
38from mindspore.communication import get_rank as get_rank_id
39from mindspore.ops import communication as ops_comm
40from mindspore.ops.function import comm_func
41# Private MindSpore symbols used by ``_MSAsyncA2ALazyBwd._issue_async_a2a`` to
42# bypass the trailing reshape that ``comm_func.all_to_all_single`` performs on
43# the default compute stream before the async ``CommHandle.wait()`` fires —
44# see that helper's docstring for the full rationale. If a future MindSpore
45# release moves or renames either symbol, this module will fail to import
46# loudly (intended — silently falling back to ``comm_func.all_to_all_single``
47# would re-introduce the race).
48from mindspore.ops.function.comm_func import _deal_comm_outputs
49from mindspore.ops.auto_generate.gen_ops_prim import inner_comm_all_to_all_v_op
50from mindspore._c_expression import TensorTransform
51import mindspore.mint.distributed as dist
53from hyper_parallel.platform.platform import Platform, PlatformType, EXISTING_COMM_GROUPS
54from hyper_parallel.platform.mindspore.dtensor import DTensorBase
55from hyper_parallel.platform.mindspore.pipeline_parallel.stage import PipelineStageBase
56from hyper_parallel.platform.mindspore.parameter_init import init_parameters as _init_parameters
57from hyper_parallel.platform.mindspore.init_weights import (
58 init_on_device as _init_on_device,
59 _install_cell_to_empty_patch,
60)
62comm_func.set_comm_ops_inplace(False)
63_tensor_transform = TensorTransform.get_instance()
66# pylint: disable=C0103
69def _a2a_reconstruct_ms(out_perm: Tensor, concat_dim: int) -> Tensor:
70 """Reconstruct A2A result from raw out_perm buffer."""
71 new_ndim = out_perm.dim()
72 chunk_in_perm = concat_dim + 1
73 recon_perm = list(range(1, chunk_in_perm)) + [0] + list(range(chunk_in_perm, new_ndim))
74 x_recon = out_perm.permute(recon_perm).contiguous()
75 shape = list(x_recon.shape)
76 merged = shape[concat_dim] * shape[concat_dim + 1]
77 return x_recon.reshape(shape[:concat_dim] + [merged] + shape[concat_dim + 2:])
80def _normalize_dim(dim: int, ndim: int) -> int:
81 """Normalize a possibly-negative dimension index."""
82 return dim + ndim if dim < 0 else dim
85def _move_dim_to_front(tensor: Tensor, dim: int) -> Tensor:
86 """Move ``dim`` to the front while preserving the other dimensions' order."""
87 dim = _normalize_dim(dim, tensor.dim())
88 if dim == 0:
89 return tensor.contiguous()
90 perm = [dim] + [i for i in range(tensor.dim()) if i != dim]
91 return tensor.permute(perm).contiguous()
94def _move_dim_from_front(tensor: Tensor, dim: int) -> Tensor:
95 """Inverse of :func:`_move_dim_to_front`."""
96 dim = _normalize_dim(dim, tensor.dim())
97 if dim == 0:
98 return tensor.contiguous()
99 perm = [dim] + [i for i in range(tensor.dim()) if i != dim]
100 inverse = [0] * len(perm)
101 for idx, value in enumerate(perm):
102 inverse[value] = idx
103 return tensor.permute(inverse).contiguous()
106def _normalize_all_to_all_single_result(result, output: Tensor) -> tuple[Tensor, object]:
107 """Normalize MindSpore all_to_all_single return values to ``(output, handle)``."""
108 if isinstance(result, tuple):
109 if len(result) != 2:
110 raise ValueError(
111 "mindspore all_to_all_single returned an unexpected tuple "
112 f"with length {len(result)}"
113 )
114 return result
115 return output, result
118def _normalize_all_gather_single_result(result, output: Tensor) -> tuple[Tensor, object]:
119 """Normalize MindSpore all_gather_into_tensor return values to ``(output, handle)``."""
120 if isinstance(result, tuple):
121 if len(result) != 2:
122 raise ValueError(
123 "mindspore all_gather_into_tensor returned an unexpected tuple "
124 f"with length {len(result)}"
125 )
126 return result
127 return output, result
130def _normalize_reduce_scatter_single_result(result, output: Tensor) -> tuple[Tensor, object]:
131 """Normalize MindSpore reduce_scatter_tensor return values to ``(output, handle)``."""
132 if isinstance(result, tuple):
133 if len(result) != 2:
134 raise ValueError(
135 "mindspore reduce_scatter_tensor returned an unexpected tuple "
136 f"with length {len(result)}"
137 )
138 return result
139 return output, result
142def _mindspore_all_to_all_single(input_tensor: Tensor, output_shape, group, async_op=False) -> tuple[Tensor, object]:
143 """Launch MindSpore all_to_all_single and normalize return values."""
144 output = mint.empty(tuple(output_shape), dtype=input_tensor.dtype)
145 result = ops_comm.all_to_all_single(output, input_tensor, group=group, async_op=async_op)
146 normalized_output, handle = _normalize_all_to_all_single_result(result, output)
147 if not async_op:
148 return normalized_output, None
149 return normalized_output, handle
152def _mindspore_all_gather_single(input_tensor: Tensor, output_shape, group, async_op=False) -> tuple[Tensor, object]:
153 """Launch MindSpore all_gather_into_tensor and normalize return values."""
154 output = mint.empty(tuple(output_shape), dtype=input_tensor.dtype)
155 result = ops_comm.all_gather_into_tensor(output, input_tensor, group=group, async_op=async_op)
156 normalized_output, handle = _normalize_all_gather_single_result(result, output)
157 if not async_op:
158 return normalized_output, None
159 return normalized_output, handle
162def _mindspore_reduce_scatter_single(
163 input_tensor: Tensor, output_shape, group, async_op=False
164) -> tuple[Tensor, object]:
165 """Launch MindSpore reduce_scatter_tensor and normalize return values."""
166 output = mint.empty(tuple(output_shape), dtype=input_tensor.dtype)
167 result = ops_comm.reduce_scatter_tensor(output, input_tensor, group=group, async_op=async_op)
168 normalized_output, handle = _normalize_reduce_scatter_single_result(result, output)
169 if not async_op:
170 return normalized_output, None
171 return normalized_output, handle
174class AsyncCollectiveTensor(Tensor):
175 """MindSpore Tensor subclass that defers ``CommHandle.wait()`` to
176 the first op that reads it.
178 Mimics PyTorch's ``AsyncCollectiveTensor`` using MindSpore's
179 per-tensor ``__ms_dispatch__`` mechanism. Constructed by calling
180 ``AsyncCollectiveTensor(inner_tensor, work)`` — :meth:`__new__`
181 invokes ``Tensor._make_subclass`` which (per MindSpore C++ side)
182 sets ``has_ms_dispatch=true`` on the new tensor because this class
183 defines ``__ms_dispatch__``. All subsequent ops involving this
184 tensor are routed through that callback.
186 Stream-side ``CommHandle.wait()`` (host non-blocking) means the
187 overlap window between the async a2a issue and the first consumer
188 op is preserved: the wait is only inserted on the consumer stream
189 at the consumer dispatch site, not at the a2a issue site.
191 Note:
192 Currently every op (including view ops like reshape /
193 transpose / permute) triggers ``work.wait()`` + unwrap.
194 Once MindSpore exposes schema alias annotations on
195 :class:`OpFunc` (planned per discussion with the MS team),
196 this class can mirror PyTorch's ``_is_view_op`` to keep
197 view chains lazy and stretch the overlap window further.
199 Attributes:
200 elem: The underlying regular Tensor (PyTorch's
201 ``AsyncCollectiveTensor.elem``). Returned by
202 :meth:`_wait_and_unwrap` after the wait fires
203 so downstream ops see a plain Tensor type.
204 completed: Whether ``work.wait()`` has already been
205 triggered (idempotency guard).
206 _pending_work: The async ``CommHandle`` returned by MindSpore.
207 PyTorch's equivalent class doesn't carry this
208 because PyTorch tracks tensor→work via the
209 global ``wait_tensor()`` aten op + c10d
210 registry. MindSpore has no such infra, so we
211 have to stash the handle on the wrapper itself.
212 """
214 __slots__ = ("elem", "completed", "_pending_work")
216 @staticmethod
217 def __new__(cls, inner: Tensor, work): # pylint: disable=W0613
218 """Construct a wrapper tensor sharing storage with ``inner``.
220 ``Tensor._make_subclass`` returns a tensor of class ``cls``
221 that shares storage with ``inner``. MindSpore C++ side then
222 sets ``has_ms_dispatch=true`` because ``cls`` defines
223 ``__ms_dispatch__``. Per-instance state is set in
224 :meth:`__init__`.
225 """
226 return Tensor._make_subclass(cls, inner) # pylint: disable=W0212
228 def __init__(self, inner: Tensor, work): # pylint: disable=W0231
229 """Initialize wrapper state (does NOT call ``super().__init__``).
231 Skipping ``Tensor.__init__`` is intentional: the parent
232 constructor would re-interpret ``inner`` as raw input data
233 and ``work`` as a dtype, corrupting the tensor that
234 :meth:`__new__` already built via ``Tensor._make_subclass``.
235 """
236 self.elem = inner
237 self.completed = work is None
238 self._pending_work = work
240 def _wait_and_unwrap(self) -> Tensor:
241 """Trigger ``work.wait()`` (idempotent) and return ``elem``.
243 Mirrors PyTorch's ``trigger_wait``: returns the underlying
244 regular Tensor so downstream ops see a plain ``Tensor``
245 instance, not an ``AsyncCollectiveTensor`` (avoids re-entering
246 ``__ms_dispatch__`` on every subsequent op).
247 """
248 if not self.completed:
249 work = self._pending_work
250 if work is not None:
251 work.wait() # stream-side: inserts streamWaitEvent on current stream
252 self.completed = True
253 return self.elem
255 @classmethod
256 def __ms_dispatch__(cls, func, args, kwargs=None):
257 """Per-tensor dispatch callback invoked for every op touching a
258 :class:`AsyncCollectiveTensor` instance.
260 Must be a ``@classmethod`` so MindSpore's C++-side invocation
261 (``tensor_py_reg.cc`` retrieves the attribute from the class
262 and calls it as ``handler(op_func, packed_args, kwargs)`` —
263 three positional args, no ``self`` binding) lines up with the
264 signature ``(cls, func, args, kwargs)``. Mirrors PyTorch's
265 ``__torch_dispatch__`` decoration on ``AsyncCollectiveTensor``.
267 Currently every op triggers wait + unwrap on any
268 ``AsyncCollectiveTensor`` arg, then runs the op on the
269 underlying inner tensors. This is the conservative
270 correctness-first behavior: it always defers the wait at
271 least until the first op consumes the tensor (which is later
272 than calling ``work.wait()`` immediately at a2a issue site,
273 so the overlap window is preserved across the
274 ``sync_hook("B")`` window).
276 TODO: when MindSpore exposes schema alias annotations on
277 ``func`` (the ``OpFunc`` parameter), add a fast path that
278 keeps view ops (reshape / transpose / permute / etc.) lazy
279 and only triggers wait on real data-touching ops, mirroring
280 PyTorch's ``_is_view_op`` in
281 ``torch/distributed/_functional_collectives.py``. Until that
282 annotation is available, treating views as real ops just
283 shortens the overlap window for view-heavy paths — it does
284 not affect correctness.
285 """
286 args = args if args is not None else ()
287 kwargs = kwargs if kwargs is not None else {}
288 unwrapped_args = tuple(
289 a._wait_and_unwrap() if isinstance(a, cls) else a # pylint: disable=W0212
290 for a in args
291 )
292 unwrapped_kwargs = {
293 k: (v._wait_and_unwrap() if isinstance(v, cls) else v) # pylint: disable=W0212
294 for k, v in kwargs.items()
295 }
296 return func(*unwrapped_args, **unwrapped_kwargs)
298 # ------------------------------------------------------------------
299 # Data-export overrides
300 # ------------------------------------------------------------------
301 # The methods below all read raw tensor data (or print it) and
302 # bypass ``__ms_dispatch__`` because they are Python-level methods
303 # on ``Tensor``, not MindSpore ops. Without these overrides they
304 # would access ``self``'s data buffer before the pending async a2a
305 # has finished, returning stale / uninitialized values. Each
306 # override forces a stream-side wait via ``_wait_and_unwrap`` and
307 # delegates to the same method on the underlying inner tensor.
308 #
309 # Methods deliberately NOT overridden:
310 # ``__len__`` — metadata only (returns shape[0]); no data read.
311 # ``__hash__`` — id-based on MindSpore Tensor; no data read.
312 # ``__contains__`` — uses ``(elem == self).any().item()`` which
313 # dispatches through ``==`` so wait fires
314 # transitively before the chain reaches data.
315 # ``__getitem__`` — slicing dispatches through ``__ms_dispatch__``.
316 # ``__format__`` — calls ``__repr__`` which we override.
318 def asnumpy(self):
319 """Convert to numpy ndarray; waits the pending a2a first."""
320 return self._wait_and_unwrap().asnumpy()
322 def numpy(self):
323 """Alias of :meth:`asnumpy` — same wait + unwrap path."""
324 return self._wait_and_unwrap().numpy()
326 def __array__(self, dtype=None):
327 """``np.array(t)`` protocol; waits + delegates to inner tensor."""
328 return self._wait_and_unwrap().__array__(dtype)
330 def get_bytes(self):
331 """Raw byte serialization; must wait before reading the buffer."""
332 return self._wait_and_unwrap().get_bytes()
334 def tolist(self):
335 """Convert to nested Python list; waits first."""
336 return self._wait_and_unwrap().tolist()
338 def item(self):
339 """Extract scalar value (0-d tensor); waits first."""
340 return self._wait_and_unwrap().item()
342 def __bool__(self):
343 """``bool(t)`` / ``if t:``; reads scalar value, must wait."""
344 return bool(self._wait_and_unwrap())
346 def __int__(self):
347 """``int(t)``; reads scalar value, must wait."""
348 return int(self._wait_and_unwrap())
350 def __float__(self):
351 """``float(t)``; reads scalar value, must wait."""
352 return float(self._wait_and_unwrap())
354 def __index__(self):
355 """Python index protocol; uses scalar value, must wait."""
356 return self._wait_and_unwrap().__index__()
358 def __repr__(self):
359 """Eager debug print; force wait so the printout reflects real data.
361 Mirrors PyTorch's ``AsyncCollectiveTensor.__repr__`` style by
362 labelling the wrapper so a stray ``print(t)`` doesn't silently
363 hide the lazy nature of the value.
364 """
365 return f"AsyncCollectiveTensor({self._wait_and_unwrap()})"
367 def __str__(self):
368 """``str(t)`` / format printing; falls through to :meth:`__repr__`."""
369 return self.__repr__()
371 def __iter__(self):
372 """Iterate over dim-0 slices; one wait, then iterate inner."""
373 return iter(self._wait_and_unwrap())
376class _MSAsyncA2ALazyBwd(_Function):
377 """Async all-to-all whose forward and backward both return
378 :class:`AsyncCollectiveTensor`, deferring ``CommHandle.wait()``
379 to the first consumer op via ``__ms_dispatch__``.
381 Mirrors the Torch ``_AsyncA2ALazyBwd`` semantics: the kernel is
382 queued on the HCCL group's stream, host returns immediately, and
383 the wait fires lazily on the consumer's stream — giving the
384 paired thread a window to dispatch its compute concurrently.
385 """
387 @staticmethod
388 def _issue_async_a2a(flat_input, send_splits, recv_splits, group):
389 """Issue an async all-to-all-v on a 1-D flat tensor.
391 Bypasses ``comm_func.all_to_all_single``: that wrapper appends an
392 unconditional ``result.reshape((-1,) + recv_shape_without_first_dim)``
393 on the default compute stream *before* the async ``CommHandle.wait()``
394 fires (the wait is deferred to the first consumer op via
395 :class:`AsyncCollectiveTensor`). MindSpore's mem_pool race_checker
396 (``MS_ALLOC_CONF=memory_tracker:True``) flags that trailing reshape
397 as a cross-stream race on the HCCL output, even though for 1-D
398 inputs it is a metadata-only no-op. Calling the inner primitive
399 directly skips the tracker-visible read on stream 0.
401 Args:
402 flat_input: 1-D tensor — must already be flattened by the caller.
403 send_splits: ``list[int]`` — element counts sent to each rank.
404 recv_splits: ``list[int]`` — element counts received from each rank.
405 group: Process group.
407 Returns:
408 ``(output_tensor, CommHandle)`` — the 1-D output and the async handle.
409 """
410 rank_size = get_group_size(group)
411 # Positional args follow the MS auto-generated primitive signature:
412 # ``(input, group, send_splits, recv_splits, rank_size, block)``.
413 # ``block=False`` selects the async path; the handle is returned in
414 # the raw tuple and unpacked by ``_deal_comm_outputs`` below.
415 raw = inner_comm_all_to_all_v_op(
416 flat_input, group, list(send_splits), list(recv_splits), rank_size,
417 False,
418 )
419 # ``_deal_comm_outputs(raw, is_async=True)`` mirrors the async branch
420 # inside ``comm_func.all_to_all_single`` — unpacks the primitive's raw
421 # output into ``(tensor, handle)`` without the trailing reshape.
422 return _deal_comm_outputs(raw, True)
424 @staticmethod
425 def forward(ctx, input_tensor, output_splits, input_splits, group): # pylint: disable=arguments-differ
426 """Launch async a2a; return :class:`AsyncCollectiveTensor`.
428 ``input_tensor`` must already be 1-D and the splits must be element
429 counts (not row counts). The caller is expected to flatten and
430 translate splits beforehand — see
431 :meth:`MindSporePlatform.differentiable_all_to_all_single_async`.
432 """
433 ctx.input_splits = input_splits
434 ctx.output_splits = output_splits
435 ctx.group = group
436 flat_input = input_tensor.reshape(-1)
437 actual_output, work = _MSAsyncA2ALazyBwd._issue_async_a2a(
438 flat_input, input_splits, output_splits, group,
439 )
440 return AsyncCollectiveTensor(actual_output, work)
442 @staticmethod
443 def backward(ctx, grad_output):
444 """Symmetric reverse a2a; returns :class:`AsyncCollectiveTensor`."""
445 # If grad_output is still lazy, force unwrap before issuing the
446 # reverse a2a (which is itself a "real" op on the data).
447 if isinstance(grad_output, AsyncCollectiveTensor):
448 grad_output = grad_output._wait_and_unwrap() # pylint: disable=W0212
449 flat_grad = grad_output.reshape(-1)
450 actual_grad, work = _MSAsyncA2ALazyBwd._issue_async_a2a(
451 flat_grad, ctx.output_splits, ctx.input_splits, ctx.group,
452 )
453 lazy_grad = AsyncCollectiveTensor(actual_grad, work)
454 return lazy_grad, None, None, None
457class _MSSyncHookFunction(_Function):
458 """Identity autograd op that fires HookCoordinator rendezvous on
459 forward and backward, mirroring the Torch ``_TorchSyncHookFunction``.
461 The role tables are intentionally identical to the Torch backend so
462 the dual-thread protocol (COMM-first dispatch ordering) is the same
463 on MindSpore.
465 Hook-name semantics:
467 - ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` — full rendezvous on both
468 forward and backward, using ``_FWD_ROLES`` / ``_BWD_ROLES``.
469 - ``"CHUNK_START"`` — pair-0 entry hook.
470 **Forward**: full rendezvous(COMPUTE) — pairs with
471 ``D_LAST.bwd`` so the BWD thread's combine.bwd of the last
472 layer is bracketed by a barrier-synced window.
473 **Backward**: paired with ``CHUNK_END.fwd`` as the BWD-side of
474 the exit barrier (roles ``(COMPUTE, COMPUTE)``).
475 - ``"D_LAST"`` — closing D hook of the last MoE layer in a chunk.
476 **Forward**: **pure skip** — neither notify nor rendezvous.
477 The C_last → combine COMM event is left un-notified so BWD's
478 COMPUTE waiter at ``A_0.bwd`` stays parked. This keeps FWD's
479 post-combine forward work serialised against BWD's Attn.bwd_0;
480 required because MS PyNative does not support concurrent
481 FWD-record + BWD-replay on its autograd executor. (The Torch
482 backend takes the looser ``notify(COMM) + skip`` path here for
483 more overlap — Torch autograd is thread-safe.)
484 **Backward**: full rendezvous using ``_BWD_ROLES["D"]``; this
485 is the very first BWD rendezvous and pairs with
486 ``CHUNK_START.fwd`` to bracket combine.bwd_last.
487 - ``"CHUNK_END"`` — pair-N exit hook (FWD side).
488 **Forward**: roles ``(COMM, COMPUTE)``. ``notify_dispatched``
489 sets the C_last event (waking BWD's A_0.bwd waiter), then
490 ``rendezvous(COMPUTE)`` parks FWD on the exit barrier so BWD's
491 Attn.bwd_0 runs with FWD already blocked — no concurrent
492 FWD-record + BWD-replay.
493 **Backward**: skipped (this would be the first node visited
494 in BWD replay; its partner ``D_LAST.bwd`` already pairs with
495 ``CHUNK_START.fwd`` on pair 0).
496 """
498 # Index encoding: 1 = COMM, 2 = COMPUTE.
499 _FWD_ROLES = {
500 # ``CHUNK_START``: chunk entry on FWD. No "previous" op on
501 # this thread within this overlap.run() — ``notify(COMPUTE)``
502 # is a no-op anyway. Next role is COMPUTE so FWD parks on
503 # ``_comm_dispatched.wait`` for BWD's ``D_LAST.bwd`` COMM.
504 "CHUNK_START": (2, 2),
505 "A": (2, 1), # prev=Attention COMPUTE | next=dispatch COMM
506 "B": (1, 2), # prev=dispatch COMM | next=module COMPUTE
507 "C": (2, 1), # prev=module COMPUTE | next=combine COMM
508 "D": (1, 2), # prev=combine COMM | next=Attention COMPUTE
509 # ``CHUNK_END``: chunk-exit hook on FWD. Does two things in
510 # one place — both critical for MS PyNative correctness:
511 # 1. ``notify_dispatched(COMM)`` sets the C_last event from
512 # C_last's rendezvous(COMM). ``D_LAST.fwd`` deliberately
513 # does NOT notify (it is a pure skip) so BWD's COMPUTE
514 # waiter at ``A_0.bwd`` stays parked until FWD has
515 # finished all chunk-local forward work (post-combine
516 # sort/index_select/multiply).
517 # 2. ``rendezvous(COMPUTE)`` parks FWD on the exit barrier.
518 # By the time BWD wakes from step 1 and starts
519 # Attn.bwd_0, FWD is already blocked at this barrier —
520 # no concurrent FWD-record + BWD-replay window.
521 "CHUNK_END": (1, 2),
522 }
523 _BWD_ROLES = {
524 # ``CHUNK_START.bwd`` is intentionally NOT engaged here.
525 # MS PyNative's autograd may skip the backward node if the
526 # chunk input lacks ``requires_grad`` (the value of
527 # ``x.grad`` is unused downstream), which would leave the
528 # pair-8 BWD partner unmatched and deadlock FWD's
529 # ``CHUNK_END`` barrier. pair-8 BWD is instead taken out of
530 # band: the OVERLAP_B_F callback's ``bwd_fn`` makes one
531 # explicit ``coordinator.rendezvous(COMPUTE)`` after
532 # ``backward_one_chunk`` returns, paired with FWD's
533 # ``CHUNK_END.fwd`` rendezvous.
534 # ``D_LAST`` on backward routes through D's BWD role (COMM
535 # next: the upcoming combine.bwd) — see the docstring above
536 # for why we no longer skip.
537 "D": (2, 1), # prev=Attn.bwd COMPUTE | next=combine.bwd COMM
538 "C": (1, 2), # prev=combine.bwd COMM | next=module.bwd COMPUTE
539 "B": (2, 1), # prev=module.bwd COMPUTE | next=dispatch.bwd COMM
540 "A": (1, 2), # prev=dispatch.bwd COMM | next=Attn.bwd COMPUTE
541 }
542 _ROLE_CACHE = None
544 @staticmethod
545 def _role_enum(idx: int):
546 """Lazy import of HookRole to avoid a circular import at module load."""
547 if _MSSyncHookFunction._ROLE_CACHE is None:
548 # pylint: disable=C0415
549 from hyper_parallel.core.pipeline_parallel.hook_coordinator import HookRole
550 _MSSyncHookFunction._ROLE_CACHE = (None, HookRole.COMM, HookRole.COMPUTE)
551 return _MSSyncHookFunction._ROLE_CACHE[idx]
553 @staticmethod
554 def _passthrough(x):
555 """Identity passthrough that defeats MS autograd's identity-output handling.
557 When :meth:`forward` returns its input unchanged, MS PyNative's
558 ``FunctionBase.apply`` sees ``is_same_as_input=True`` on the output
559 and inserts a ``ViewAsSelfWithNoGrad`` (a ``view(self, self.shape)``
560 kernel) on the current compute stream. If the input is an
561 :class:`AsyncCollectiveTensor` whose lazy ``CommHandle.wait()`` has
562 not yet fired, that view runs on the default stream while the HCCL
563 kernel is still writing the same memory on the comm stream — flagged
564 by MS's mem_pool ``race_checker`` (``MS_ALLOC_CONF=memory_tracker:True``).
566 Returning a freshly wrapped :class:`AsyncCollectiveTensor` keeps the
567 same underlying buffer and pending work, but yields a new
568 ``shared_ptr<Tensor>`` so ``is_same_as_input`` is ``False`` and no
569 autograd view is emitted. For regular tensors the original
570 passthrough is safe (the view sits on the same stream as the data).
572 Note:
573 The clone shares ``_pending_work`` with the original but keeps
574 an independent ``completed`` flag. Two assumptions:
576 * ``CommHandle.wait()`` is idempotent — relied on whenever both
577 wrappers end up being consumed (matches the existing
578 :meth:`AsyncCollectiveTensor._wait_and_unwrap` pattern, which
579 also does not null out ``_pending_work`` after waiting).
580 * Per-wrapper ``completed`` is intentional: a ``wait()`` on
581 stream A does not synchronize stream B, so each consumer
582 stream must be free to re-issue its own wait.
583 """
584 if isinstance(x, AsyncCollectiveTensor):
585 new_wrapper = AsyncCollectiveTensor(x.elem, x._pending_work) # pylint: disable=W0212
586 new_wrapper.completed = x.completed
587 return new_wrapper
588 return x
590 @staticmethod
591 def forward(ctx, x, hook_name, coordinator): # pylint: disable=arguments-differ
592 """Fire forward-direction rendezvous and return ``x`` unchanged."""
593 ctx.hook_name = hook_name
594 ctx.coordinator = coordinator
595 if not coordinator.is_enabled():
596 return _MSSyncHookFunction._passthrough(x)
597 if hook_name == "D_LAST":
598 # Pure skip — neither notify nor rendezvous. The
599 # C_last → combine COMM event is left un-notified on
600 # purpose so BWD's COMPUTE waiter at A_0.bwd stays parked
601 # until FWD reaches CHUNK_END.fwd. This keeps FWD's
602 # post-combine forward work (sort / index_select / probs
603 # mul / strided_slice) strictly serialised against BWD's
604 # Attn.bwd_0 — required because MS PyNative does not
605 # support concurrent FWD-record + BWD-replay on the
606 # autograd executor.
607 return _MSSyncHookFunction._passthrough(x)
608 prev_idx, next_idx = _MSSyncHookFunction._FWD_ROLES[hook_name]
609 role_of = _MSSyncHookFunction._role_enum
610 coordinator.notify_dispatched(role_of(prev_idx))
611 coordinator.rendezvous(role_of(next_idx))
612 return _MSSyncHookFunction._passthrough(x)
614 @staticmethod
615 def backward(ctx, grad_output):
616 """Mirror of :meth:`forward` using ``_BWD_ROLES``."""
617 hook_name = ctx.hook_name
618 coordinator = ctx.coordinator
619 if not coordinator.is_enabled():
620 return _MSSyncHookFunction._passthrough(grad_output), None, None
621 if hook_name in ("CHUNK_END", "CHUNK_START"):
622 # Both boundary hooks skip in backward:
623 # * ``CHUNK_END.bwd`` would fire FIRST in BWD replay (it
624 # wraps the chunk's last forward op). We do not want
625 # a rendezvous here — pair 0 is handled by
626 # ``D_LAST.bwd`` ↔ ``CHUNK_START.fwd``.
627 # * ``CHUNK_START.bwd`` would fire LAST. We do not
628 # rendezvous here either, because MS autograd may skip
629 # the node entirely when the chunk input lacks
630 # ``requires_grad`` (unused ``x.grad``). pair-8 BWD
631 # is taken out of band — see the role-table comment.
632 return _MSSyncHookFunction._passthrough(grad_output), None, None
633 # ``D_LAST.bwd`` reuses D's BWD role: it is the *first non-skip*
634 # BWD rendezvous and pairs with FWD's ``CHUNK_START`` to lock
635 # the combine.bwd_last launch inside a barrier-synced window.
636 role_name = "D" if hook_name == "D_LAST" else hook_name
637 prev_idx, next_idx = _MSSyncHookFunction._BWD_ROLES[role_name]
638 role_of = _MSSyncHookFunction._role_enum
639 coordinator.notify_dispatched(role_of(prev_idx))
640 coordinator.rendezvous(role_of(next_idx))
641 return _MSSyncHookFunction._passthrough(grad_output), None, None
644class _MSAsyncA2AFunction(_Function):
645 """Differentiable wrapper for pre-launched async all-to-all."""
647 @staticmethod
648 def forward(ctx, x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box): # pylint: disable=arguments-differ
649 """Wait for pre-launched async A2A and return reconstructed output."""
650 ctx.group = group
651 ctx.world_size = world_size
652 ctx.concat_dim = concat_dim
653 ctx.split_dim = split_dim
654 ctx.handle_box = handle_box
655 ctx.x_shape = tuple(x.shape)
656 work.wait()
657 return _a2a_reconstruct_ms(out_perm, concat_dim)
659 @staticmethod
660 def backward(ctx, grad_output):
661 """Launch async head->seq A2A for backward overlap, or return zero grad."""
662 if ctx.handle_box is not None:
663 g = grad_output.contiguous()
664 shape = list(g.shape)
665 seq_dim = ctx.concat_dim
666 s_full = shape[seq_dim]
667 ndim = len(shape) + 1
668 x_perm = g.reshape(
669 shape[:seq_dim] + [ctx.world_size, s_full // ctx.world_size] + shape[seq_dim + 1:]
670 ).permute(
671 [seq_dim] + list(range(seq_dim)) + list(range(seq_dim + 1, ndim))
672 ).contiguous()
673 out_perm, work = _mindspore_all_to_all_single(
674 x_perm,
675 list(x_perm.shape),
676 ctx.group,
677 async_op=True,
678 )
679 ctx.handle_box.append((work, out_perm))
680 return mint.zeros(ctx.x_shape, dtype=grad_output.dtype), None, None, None, None, None, None, None
683class _MSAsyncAllGatherFunction(_Function):
684 """Differentiable wrapper for pre-launched async all-gather."""
686 @staticmethod
687 def forward(ctx, x, work, out_perm, group, world_size, gather_dim, handle_box): # pylint: disable=arguments-differ
688 """Wait for pre-launched all-gather and reconstruct the gathered tensor."""
689 ctx.group = group
690 ctx.world_size = world_size
691 ctx.gather_dim = gather_dim
692 ctx.handle_box = handle_box
693 ctx.x_shape = tuple(x.shape)
694 work.wait()
695 return _move_dim_from_front(out_perm, gather_dim)
697 @staticmethod
698 def backward(ctx, grad_output):
699 """Launch reverse reduce-scatter for the all-gather."""
700 grad_perm = _move_dim_to_front(grad_output.contiguous(), ctx.gather_dim)
701 output_shape = list(grad_perm.shape)
702 if output_shape[0] % ctx.world_size != 0:
703 raise ValueError(
704 "all_gather backward expected gathered dimension to be divisible by world_size, "
705 f"got {output_shape[0]} and {ctx.world_size}."
706 )
707 output_shape[0] //= ctx.world_size
708 output, work = _mindspore_reduce_scatter_single(
709 grad_perm,
710 output_shape,
711 ctx.group,
712 async_op=True,
713 )
714 if ctx.handle_box is not None:
715 ctx.handle_box.append((work, output, ctx.gather_dim))
716 return mint.zeros(ctx.x_shape, dtype=grad_output.dtype), None, None, None, None, None, None
717 work.wait()
718 return _move_dim_from_front(output, ctx.gather_dim), None, None, None, None, None, None
721def _ensure_contiguous(x):
722 """Return a contiguous copy of *x* if not already contiguous."""
723 if not x.is_contiguous() or x.storage_offset() != 0:
724 x = x.contiguous()
725 return x
728class MindSporePlatform(Platform):
729 """MindSpore platform api"""
730 Tensor = Tensor
731 tensor = Tensor
732 Parameter = Parameter
733 Module = Cell
734 DTensorBase = DTensorBase
735 PipelineStageBase = PipelineStageBase
736 platform_type = PlatformType.MINDSPORE
737 tensor_dtype = mstype
738 dtype = ms.Type
739 Function = _Function
741 _custom_ops_cls = None
743 @property
744 def custom_ops(self):
745 """Return the MindSpore platform custom ops instance.
747 .. warning::
748 This is an experimental API that subject to change or deletion.
750 Returns:
751 MindSporeCustomOps: Custom ops class that delegates to DFunction
752 implementations wrapping Ascend NPU custom C++ kernels.
753 """
754 if self._custom_ops_cls is None:
755 from hyper_parallel.platform.mindspore.custom_ops.custom_ops import ( # pylint: disable=import-outside-toplevel
756 MindSporeCustomOps,
757 )
758 self._custom_ops_cls = MindSporeCustomOps
759 return self._custom_ops_cls
761 def __init__(self):
762 # Ensure MindSpore ``nn.Cell.to_empty`` is patched as soon as the
763 # MindSpore platform instance is created.
764 _install_cell_to_empty_patch()
766 @staticmethod
767 def is_linear_module(module) -> bool:
768 """Check whether *module* is a MindSpore ``Dense`` (linear) or ``mint.nn.Linear`` layer."""
769 return isinstance(module, (ms.nn.Dense, mint.nn.Linear))
771 @staticmethod
772 def is_embedding_module(module) -> bool:
773 """Check whether *module* is a MindSpore ``Embedding`` or ``mint.nn.Embedding`` layer."""
774 return isinstance(module, (ms.nn.Embedding, mint.nn.Embedding))
776 def device_count(self, device_handle):
777 """
778 Get the number of available devices.
780 Args:
781 device_handle: The device handle (e.g., ms.device_context).
783 Returns:
784 int: The number of available devices.
785 """
786 device_type = self.device_type()
787 if device_type == "cpu":
788 return device_handle.device_context.cpu.device_count()
789 if device_type == "gpu":
790 return device_handle.device_context.gpu.device_count()
791 return device_handle.device_context.ascend.device_count()
793 @staticmethod
794 def get_rng_state(device=None, device_handle=None):
795 """
796 Get the random number generator state.
798 Args:
799 device (Optional): The device to get RNG state from (not used in MindSpore).
800 device_handle (Optional): The device handle (not used in MindSpore).
802 Returns:
803 Tensor: The RNG state as a tensor.
804 """
805 _ = device, device_handle
806 return ms.get_rng_state()
808 @staticmethod
809 def set_rng_state(state, device=None, device_handle=None):
810 """
811 Set the random number generator state.
813 Args:
814 state (Tensor): The RNG state to set.
815 device (Optional): The device to set RNG state for (not used in MindSpore).
816 device_handle (Optional): The device handle (not used in MindSpore).
817 """
818 _ = device, device_handle
819 return ms.set_rng_state(state)
821 def device_type(self):
822 """
823 Get the current device type.
825 Returns:
826 str: The device type string ("npu" for Ascend, "gpu" for GPU, "cpu" for CPU).
827 """
828 device_type = ms.get_context("device_target")
829 if device_type == "Ascend":
830 return "npu"
831 return device_type.lower()
833 def device(self, device_idx=None):
834 """
835 Get the device type string.
837 Args:
838 device_idx (Optional[int]): The device index (not used in MindSpore).
840 Returns:
841 str: The device type string.
842 """
843 _ = device_idx
844 device_type = self.device_type()
845 return device_type
847 @staticmethod
848 def get_device_handle():
849 """
850 Get the MindSpore module as the device handle.
852 Returns:
853 module: The mindspore module.
854 """
855 return ms
857 @staticmethod
858 def manual_seed(seed):
859 """
860 Set the random seed for reproducibility.
862 Args:
863 seed (int): The random seed value.
865 Returns:
866 None
867 """
868 return ms.manual_seed(seed)
870 @staticmethod
871 def ones(size, dtype=None):
872 """
873 Create a tensor filled with ones.
875 Args:
876 size (tuple): The shape of the output tensor.
877 dtype (Optional[ms.Type]): The desired data type.
879 Returns:
880 Tensor: A tensor filled with ones.
881 """
882 return mint.ones(size, dtype=dtype)
884 @staticmethod
885 def zeros(size, dtype=None, device=None):
886 """
887 Create a tensor filled with zeros.
889 Args:
890 size (tuple): The shape of the output tensor.
891 dtype (Optional[ms.Type]): The desired data type.
892 device (Optional[ms.device]): The device to create the tensor on.
894 Returns:
895 Tensor: A tensor filled with zeros.
896 """
897 tensor = mint.zeros(size, dtype=dtype)
898 if device in ("GPU", "Ascend"):
899 return tensor.to(device)
900 return tensor
902 @staticmethod
903 def full(size, fill_value, dtype=None):
904 """
905 Create a tensor filled with a scalar value.
907 Args:
908 size (tuple): The shape of the output tensor.
909 fill_value (scalar): The value to fill the tensor with.
910 dtype (Optional[ms.Type]): The desired data type.
912 Returns:
913 Tensor: A tensor filled with the specified value.
914 """
915 return mint.full(size, fill_value, dtype=dtype)
917 @staticmethod
918 def empty(size, dtype=None, device=None): # pylint: disable=unused-argument
919 """
920 Create an uninitialized tensor.
922 Args:
923 size (tuple): The shape of the output tensor.
924 dtype (Optional[ms.Type]): The desired data type.
925 device: Accepted for cross-backend signature parity with the
926 Torch backend but ignored — under MindSpore the active
927 device is bound at process init via ``ms.set_device`` and
928 ``mint.empty`` allocates on it directly.
930 Returns:
931 Tensor: An uninitialized tensor.
932 """
933 return mint.empty(size, dtype=dtype)
935 @staticmethod
936 def rand(size, dtype=None, device=None): # pylint: disable=unused-argument
937 """Create a tensor filled with uniform random values in ``[0, 1)``."""
938 tensor = mint.rand(size, dtype=dtype)
939 if device in ("GPU", "Ascend"):
940 return tensor.to(device)
941 return tensor
943 @staticmethod
944 def randn(size, dtype=None, device=None): # pylint: disable=unused-argument
945 """Create a tensor filled with standard-normal random values."""
946 tensor = mint.randn(size, dtype=dtype)
947 if device in ("GPU", "Ascend"):
948 return tensor.to(device)
949 return tensor
951 @staticmethod
952 def get_rank():
953 """
954 Get the rank of the current process in the distributed group.
956 Returns:
957 int: The rank of the current process.
958 """
959 return get_rank_id()
961 @staticmethod
962 def get_global_rank(group, group_rank):
963 """
964 Get the global rank from a group rank.
966 Args:
967 group (str): The process group name.
968 group_rank (int): The rank within the group.
970 Returns:
971 int: The global rank.
972 """
973 return dist.get_global_rank(group, group_rank)
975 @staticmethod
976 def get_world_size():
977 """
978 Get the total number of processes in the distributed group.
980 Returns:
981 int: The world size.
982 """
983 return get_group_size()
985 @staticmethod
986 def get_op_name(func):
987 """
988 Extract the operation name from a function.
990 Args:
991 func: The function to extract the name from.
993 Returns:
994 str: The operation name.
995 """
996 return func.name
998 @staticmethod
999 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
1000 data = _ensure_contiguous(data)
1001 # rank_list is accepted for torch parity; MindSpore keeps the existing group order.
1002 output, _ = comm_func.all_gather_into_tensor(None, data, group=group)
1003 if concat_dim == 0:
1004 return output
1005 output_tensors = ms.ops.Split(output_num=concat_size)(output)
1006 return ms.mint.concat(output_tensors, concat_dim)
1008 @staticmethod
1009 def chunk(data, split_dim, split_size, index):
1010 return ms.ops.Split(axis=split_dim, output_num=split_size)(data)[index]
1012 @staticmethod
1013 def differentiable_all_to_all(input_data, output_shape, group):
1014 input_data = _ensure_contiguous(input_data)
1015 output_tensor, _ = comm_func.all_to_all_single(
1016 output_shape,
1017 input_data,
1018 group=group,
1019 async_op=False
1020 )
1021 return output_tensor
1023 @staticmethod
1024 def tensor_type_cast(input_data, cast_type):
1025 """Cast tensor to specified data type."""
1026 type_mapping = {
1027 'float32': ms.float32,
1028 'float16': ms.float16,
1029 'int64': ms.int64,
1030 'int32': ms.int32
1031 }
1032 if cast_type not in type_mapping:
1033 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")
1034 return input_data.to(type_mapping[cast_type])
1036 @staticmethod
1037 def differentiable_all_reduce(data, op, group):
1038 data = _ensure_contiguous(data)
1039 output, _ = comm_func.all_reduce(data, op, group)
1040 return output
1042 @staticmethod
1043 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
1044 data = _ensure_contiguous(data)
1045 if axis > 0:
1046 data = ms.mint.concat(ms.ops.Split(axis=axis, output_num=dev_num)(data), dim=0)
1047 output_tensor, _ = comm_func.reduce_scatter_tensor(None, data, 'sum', group)
1048 if op == 'avg':
1049 output_tensor = output_tensor / dev_num
1050 return output_tensor
1052 @staticmethod
1053 def init_parameters(module, stage_index):
1054 return _init_parameters(module, stage_index)
1056 # pylint: disable=W0212
1057 @staticmethod
1058 def update_param_data(param, data):
1059 """update param data"""
1060 if isinstance(param, DTensorBase):
1061 param.set_data(data)
1062 else:
1063 param._update_data(data)
1065 @staticmethod
1066 def load_into_param(param, data):
1067 copy_tensor = MindSporePlatform.empty_like(data)
1068 copy_tensor.copy_(data)
1069 if isinstance(param, DTensorBase):
1070 param.set_data(copy_tensor)
1071 else:
1072 param._update(copy_tensor)
1074 @staticmethod
1075 def get_cell_construct(cell):
1076 return cell.construct
1078 @staticmethod
1079 def get_cells_and_names(cell):
1080 return cell.cells_and_names()
1082 @staticmethod
1083 def get_modules(module):
1084 return module.cells()
1086 @staticmethod
1087 def search_parameter_by_name(cell, param_name: str):
1088 """
1089 Find the parent Module of the parameter, the parameter's name in the parent Module, and the parameter.
1090 Return value: (parent Module instance, parameter's name in parent Module, parameter object).
1091 Returns None if not found.
1092 """
1093 # Remove the "self." prefix from param_name (to maintain compatibility with original logic)
1094 param_name = param_name.replace("self.", "")
1095 # Case 1: The parameter is a direct parameter of the current Module (not in any sub-Module)
1096 if param_name in cell._params:
1097 return (cell, param_name, cell._params[param_name])
1099 # Case 2: The parameter is in a sub-Module (supports multi-level nesting, e.g., "net_b.dense1.weight")
1100 if "." in param_name:
1101 # Split into: sub-Module path + parameter name (e.g., "net_b.dense1" + "weight")
1102 cell_path, param_key = param_name.rsplit(".", 1)
1103 try:
1104 # Locate the sub-Module where the parameter resides (supports multi-level paths)
1105 target_cell = cell.get_sub_cell(cell_path)
1106 # Check if the sub-Module directly contains this parameter
1107 if param_key in target_cell._params:
1108 return target_cell, param_key, target_cell._params[param_key]
1109 except AttributeError:
1110 # Sub-Module path does not exist or the parameter is not in that sub-Module
1111 pass
1113 # Traverse all sub-Modules (recursively) to search for the parameter
1114 for _, child_cell in cell._cells.items():
1115 if isinstance(child_cell, Cell):
1116 # Recursively search within the sub-Module
1117 result = MindSporePlatform.search_parameter_by_name(child_cell, param_name)
1118 if result is not None:
1119 return result
1121 return None
1123 @staticmethod
1124 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
1125 """
1126 Modify the original parameter in a Module or sub-Module using the search result
1127 Args:
1128 cell: The cell which parameter is to update
1129 result: A tuple contains parent Module, parameter key and old parameter.
1130 new_param: New Parameter object (used to replace the original parameter)
1131 """
1132 parent_cell, param_key, _ = result
1133 # Key operation: directly modify the _params dictionary of the parent Module (original storage location)
1134 parent_cell._params[param_key] = new_param
1136 if param_key in parent_cell.__dict__:
1137 parent_cell.__dict__[param_key] = new_param
1138 parent_cell._params_list[param_key] = new_param
1139 return True
1141 @staticmethod
1142 def set_layout_into_parameter(param, layout):
1143 """Set layout in to parameter"""
1144 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel
1145 from hyper_parallel.core.dtensor.layout import _infer_slice_shape_by_layout, \
1146 _get_slice_tensor_by_layout # pylint: disable=import-outside-toplevel
1147 if isinstance(param, DTensor):
1148 raise ValueError(f"Parameter {param.name} has been configured layout, cannot be set repeatedly.")
1149 param_info = param.param_info
1150 requires_grad = param.requires_grad
1151 name = param.name
1152 slice_shape = _infer_slice_shape_by_layout(param.shape, layout)
1154 if not param.has_init:
1155 # has been init, get slice data
1156 param_dtensor = DTensor.from_local(
1157 _get_slice_tensor_by_layout(param, layout).value(), layout.mesh, layout.alias_placements
1158 )
1159 param = Parameter(param_dtensor, name=name, requires_grad=requires_grad)
1160 param.param_info = param_info
1161 else:
1162 # has not been init, need to modify init shape
1163 param.init_mode.shape = slice_shape
1164 param_dtensor = DTensor.from_local(param.init_mode, layout.mesh, layout.alias_placements)
1165 param = Parameter(param_dtensor, name=name, requires_grad=requires_grad)
1166 param.param_info = param_info
1167 return param
1169 @staticmethod
1170 def get_param_local_shape(param):
1171 """get param local shape"""
1172 if isinstance(param, DTensorBase):
1173 return param.local_shape
1174 return param.shape
1176 @staticmethod
1177 def get_param_local_data(param):
1178 """get param local shape"""
1179 if isinstance(param, DTensorBase):
1180 return param.to_local()
1181 return param
1183 @staticmethod
1184 def get_param_type_size(param):
1185 return type_size_in_bytes(param.dtype)
1187 @staticmethod
1188 def is_tensor(obj: Any) -> bool:
1189 """Return True if ``obj`` is a ``mindspore.Tensor``."""
1190 return isinstance(obj, Tensor)
1192 @staticmethod
1193 def get_tensor_storage_size(tensor: Any) -> int:
1194 """Return serialized byte size (numel * itemsize) for a MindSpore tensor."""
1195 if not MindSporePlatform.is_tensor(tensor):
1196 raise TypeError(
1197 f"MindSporePlatform.get_tensor_storage_size expects mindspore.Tensor, got {type(tensor)!r}"
1198 )
1199 return int(tensor.numel()) * int(tensor.itemsize)
1201 @staticmethod
1202 def new_zero_parameter(param_shape, param_type, requires_grad, device):
1203 param = Parameter(initializer("zeros", param_shape, param_type), requires_grad=requires_grad)
1204 if device in ("GPU", "Ascend"):
1205 return param.to(device)
1206 return param
1208 @staticmethod
1209 def new_tensor(tensor_shape, tensor_type, device):
1210 tensor = Tensor(shape=tensor_shape, dtype=tensor_type)
1211 if device in ("GPU", "Ascend"):
1212 return tensor.to(device)
1213 return tensor
1215 @staticmethod
1216 def full_like(tensor, fill_value, dtype=None):
1217 return mint.full_like(tensor, fill_value, dtype=dtype)
1219 @staticmethod
1220 def isend(tensor, dst=None, group=None, tag=0):
1221 return dist.isend(tensor, dst, group, tag)
1223 @staticmethod
1224 def irecv(tensor, src=None, group=None, tag=0):
1225 return dist.irecv(tensor, src, group, tag)
1227 @staticmethod
1228 def p2p_op(op_type, tensor, peer, group=None):
1229 # pylint: disable=C0415
1230 from mindspore.mint.distributed import P2POp
1231 return P2POp(op_type, tensor, peer, group)
1233 @staticmethod
1234 def batch_isend_irecv(p2p_ops):
1235 """Launch a peer-batched P2P group.
1237 MindSpore's ``batch_isend_irecv`` lowers the whole list to a single
1238 ``HcclBatchISendIRecv`` kernel on one comm stream and returns a list
1239 with one packaging ``CommHandle``; we hand that single handle back so
1240 callers can defer the whole batch's wait to one consumption point.
1241 A send and a recv to the same peer therefore overlap on the duplex
1242 link inside this one kernel.
1243 """
1244 # pylint: disable=C0415
1245 from mindspore.mint.distributed import batch_isend_irecv
1246 if not p2p_ops:
1247 return None
1248 handles = batch_isend_irecv(p2p_ops)
1249 return handles[0] if handles else None
1251 @staticmethod
1252 def prepare_batch_p2p_group(group: Any = None) -> None: # pylint: disable=unused-argument
1253 """Prepare a group for batched P2P operations.
1255 MindSpore does not require full-group participation before its first
1256 subset ``batch_isend_irecv`` call, so no synchronization is needed.
1258 Args:
1259 group: The communication group used by the batched P2P operations.
1260 ``None`` uses the default group.
1261 """
1263 @staticmethod
1264 def p2p_exchange(tensor, peer_rank: int, group=None): # pylint: disable=unused-argument
1265 raise NotImplementedError(
1266 "p2p_exchange is not yet supported on the MindSpore platform."
1267 )
1269 @staticmethod
1270 def send_object_list(obj_list, dst=None, group=None):
1271 # pylint: disable=C0415
1272 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import send_object_list
1273 send_object_list(obj_list, dst, group)
1275 @staticmethod
1276 def recv_object_list(obj_list, src=None, group=None):
1277 # pylint: disable=C0415
1278 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import recv_object_list
1279 recv_object_list(obj_list, src, group)
1281 @staticmethod
1282 def set_tensor_requires_grad(input_tensor):
1283 """
1284 set requires grad flag for input tensor
1285 """
1286 input_tensor.requires_grad_()
1288 @staticmethod
1289 def _normalize_group_options(pg_options: Any) -> Any:
1290 if not isinstance(pg_options, dict) or "hccl_config" not in pg_options:
1291 return pg_options
1292 from mindspore._c_expression import GroupOptions # pylint: disable=C0415
1294 options = GroupOptions()
1295 options.hccl_config = pg_options["hccl_config"]
1296 return options
1298 @staticmethod
1299 def _create_group_with_options(group_name: str, rank_list: list[int], pg_options: Any = None) -> None:
1300 """Create a MindSpore communication group with optional backend-specific options."""
1301 if pg_options is None:
1302 new_group(rank_ids=rank_list, group=group_name)
1303 return
1304 try:
1305 new_group(
1306 rank_ids=rank_list,
1307 group=group_name,
1308 options=MindSporePlatform._normalize_group_options(pg_options),
1309 )
1310 except (ImportError, RuntimeError, TypeError, ValueError):
1311 new_group(rank_ids=rank_list, group=group_name)
1313 def _create_group(self, rank_list, pg_options: Any = None):
1314 world_group = self._maybe_reuse_world_group(rank_list)
1315 if world_group is not None:
1316 return world_group
1318 group_name = str(tuple(sorted(rank_list)))
1319 self._create_group_with_options(group_name, rank_list, pg_options=pg_options)
1320 EXISTING_COMM_GROUPS[group_name] = group_name
1321 return group_name
1323 @staticmethod
1324 def all_gather_into_tensor(data, group_info, async_op=False):
1325 group_name = group_info if isinstance(group_info, str) else group_info.group_name
1326 rank_size = get_group_size(group_name) if isinstance(group_info, str) else group_info.rank_size
1327 output_shape = list(data.shape)
1328 output_shape[0] *= rank_size
1329 return _mindspore_all_gather_single(data, output_shape, group_name, async_op=async_op)
1331 @staticmethod
1332 def all_gather_single(input_tensor, output_shape, group, async_op=False):
1333 return _mindspore_all_gather_single(input_tensor, output_shape, group, async_op=async_op)
1335 @staticmethod
1336 def all_reduce(data, group_info, async_op=False):
1337 if isinstance(group_info, str):
1338 handle = dist.all_reduce(data, group=group_info, async_op=async_op)
1339 else:
1340 handle = dist.all_reduce(data, group=group_info.group_name, async_op=async_op)
1341 return data, handle
1343 @staticmethod
1344 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
1345 if group_src is not None:
1346 ranks = MindSporePlatform.get_process_group_ranks(group)
1347 src = ranks[group_src]
1348 handle = dist.broadcast(data, src, group, async_op)
1349 if async_op:
1350 handle.wait()
1351 return data
1353 @staticmethod
1354 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
1355 group_name = group if isinstance(group, str) else getattr(group, "group_name", group)
1356 if group_src is not None:
1357 ranks = MindSporePlatform.get_process_group_ranks(group)
1358 src = ranks[group_src]
1359 if scatter_list is None:
1360 # MindSpore mint.scatter validates scatter_list on every rank; PyTorch passes None on receivers.
1361 rank_size = get_group_size(group_name)
1362 scatter_list = [output] * rank_size
1363 else:
1364 scatter_list = [c.contiguous() if hasattr(c, "is_contiguous") and not c.is_contiguous() else c
1365 for c in scatter_list]
1366 handle = dist.scatter(output, scatter_list, src, group_name, async_op=async_op)
1367 if async_op and handle is not None:
1368 handle.wait()
1369 return output
1371 @staticmethod
1372 def reduce_scatter_tensor(data, group_info, async_op=False):
1373 group_name = group_info if isinstance(group_info, str) else group_info.group_name
1374 rank_size = get_group_size(group_name) if isinstance(group_info, str) else group_info.rank_size
1375 output_shape = list(data.shape)
1376 output_shape[0] //= rank_size
1377 return _mindspore_reduce_scatter_single(data, output_shape, group_name, async_op=async_op)
1379 @staticmethod
1380 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1381 return _mindspore_reduce_scatter_single(input_tensor, output_shape, group, async_op=async_op)
1383 @staticmethod
1384 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1385 return _mindspore_all_to_all_single(input_tensor, output_shape, group, async_op=async_op)
1387 @staticmethod
1388 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1389 handle_box=None):
1390 return _MSAsyncAllGatherFunction.apply(
1391 x, work, out_perm, group, world_size, gather_dim, handle_box
1392 )
1394 @staticmethod
1395 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim, # pylint: disable=unused-argument
1396 handle_box=None):
1397 return _MSAsyncA2AFunction.apply(
1398 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1399 )
1401 @staticmethod
1402 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1403 """Launch an asynchronous, differentiable all-to-all-single.
1405 Token a2a entry point used by ``CommComputeOverlap``-driven MoE
1406 wrappers. The kernel is queued on the HCCL group's stream and
1407 the host returns immediately, so the calling thread can proceed
1408 to the next sync hook (notify + rendezvous) before the
1409 collective finishes — this is what enables the comm/compute
1410 overlap window on the paired thread.
1412 Args:
1413 input_tensor: **1-D** tensor — the caller is responsible for
1414 flattening multi-dim inputs beforehand.
1415 input_splits: ``list[int]`` — **element** counts sent to each
1416 rank (not row counts). For an originally
1417 ``(N, D)`` tensor, each entry is ``rows_i * D``.
1418 output_splits: ``list[int]`` — element counts received from each rank.
1419 group: Process group.
1421 Returns:
1422 ``AsyncCollectiveTensor`` of shape ``(sum(output_splits),)`` that
1423 defers ``CommHandle.wait()`` to the first consumer op via
1424 ``__ms_dispatch__``.
1426 Raises:
1427 ValueError: if ``input_tensor`` is not 1-D.
1429 Note:
1430 The 1-D + element-count contract diverges from the Torch
1431 implementation (which accepts N-D input + row-count splits).
1432 The divergence is intentional for now: it lets the MS path
1433 call the inner primitive directly and avoid the cross-stream
1434 race that ``comm_func.all_to_all_single``'s trailing reshape
1435 triggers under ``MS_ALLOC_CONF=memory_tracker:True`` —
1436 see :meth:`_MSAsyncA2ALazyBwd._issue_async_a2a`.
1437 """
1438 if input_tensor.ndim != 1:
1439 raise ValueError(
1440 "MindSporePlatform.differentiable_all_to_all_single_async requires a 1-D "
1441 f"input_tensor (got ndim={input_tensor.ndim}, shape={tuple(input_tensor.shape)}). "
1442 "Flatten the tensor and convert row-count splits to element counts before calling."
1443 )
1444 return _MSAsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1446 @staticmethod
1447 def differentiable_sync_hook(x, hook_name: str, coordinator):
1448 """Fire a HookCoordinator rendezvous on forward and backward.
1450 Args:
1451 x: Input tensor — returned unchanged.
1452 hook_name: One of:
1453 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1454 full rendezvous on both directions.
1455 * ``"CHUNK_START"`` — chunk-entry hook on
1456 forward; pairs with ``D_LAST.bwd`` so the
1457 BWD thread's combine.bwd of the last layer
1458 is bracketed by a barrier-synced sync point.
1459 Skipped on backward.
1460 * ``"D_LAST"`` — closing D of the last MoE
1461 layer in a chunk. Forward: ``notify_dispatched``
1462 only (no Attention follows so rendezvous is
1463 skipped). Backward: full rendezvous via D's
1464 BWD role; paired with ``CHUNK_START`` on FWD.
1465 coordinator: The :class:`HookCoordinator` driving the
1466 rendezvous protocol.
1468 Returns:
1469 ``x`` unchanged.
1471 Note:
1472 Two-thread compatibility on MindSpore PyNative is not yet
1473 fully verified. The HookCoordinator + ``_Function``
1474 primitives are individually thread-safe, but the
1475 interaction with MindSpore's autograd execution model
1476 under ``threading.Thread`` should be PoC-tested before
1477 production use.
1478 """
1479 return _MSSyncHookFunction.apply(x, hook_name, coordinator)
1481 @staticmethod
1482 def parameters_dict(cell: Cell):
1483 return cell.parameters_and_names()
1485 @staticmethod
1486 def buffers_dict(cell: Cell) -> Any:
1487 """Return all named buffers registered by the cell tree."""
1488 return cell.named_buffers()
1490 @staticmethod
1491 def get_tensor_transform():
1492 return _tensor_transform
1494 @staticmethod
1495 def construct_strided_slice(x, begin, end, stride):
1496 return ms.ops.strided_slice(x, begin, end, stride)
1498 @staticmethod
1499 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1500 # pylint: disable=C0415
1501 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import _MicroBatch
1502 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1504 @staticmethod
1505 def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
1506 """Get the state dictionary of a model (not yet supported on MindSpore).
1508 Args:
1509 model: The model to extract state from.
1510 options: Optional configuration for state dict extraction.
1512 Returns:
1513 dict: The state dictionary containing model parameters and buffers.
1515 Raises:
1516 NotImplementedError: MindSpore support is not yet implemented.
1517 """
1518 raise NotImplementedError(
1519 "get_model_state_dict is not yet supported on MindSpore"
1520 )
1522 @staticmethod
1523 def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
1524 """Set the state dictionary of a model (not yet supported on MindSpore).
1526 Args:
1527 model: The model to load state into.
1528 model_state_dict: The state dict to load into the model.
1529 options: Optional configuration for state dict loading.
1531 Returns:
1532 None.
1534 Raises:
1535 NotImplementedError: MindSpore support is not yet implemented.
1536 """
1537 raise NotImplementedError(
1538 "set_model_state_dict is not yet supported on MindSpore"
1539 )
1541 @staticmethod
1542 def save_checkpoint(cell: Union[Cell, dict], file_path: str, ckpt_format: str = "safetensors") -> None:
1543 if isinstance(cell, dict):
1544 save_dict = {}
1545 for k, v in cell.items():
1546 if isinstance(v, Parameter):
1547 save_dict[k] = v
1548 elif isinstance(v, Tensor):
1549 save_dict[k] = Parameter(v, name=k)
1550 else:
1551 save_dict[k] = v
1552 else:
1553 save_dict = cell._params
1554 ms.save_checkpoint(save_obj=save_dict, ckpt_file_name=file_path, format=ckpt_format)
1556 @staticmethod
1557 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
1558 return ms.load_checkpoint(ckpt_file_name=file_path, format=ckpt_format)
1560 @staticmethod
1561 def get_symmetric_memory_handler():
1562 # pylint: disable=C0415
1563 from hyper_parallel.platform.mindspore.symmetric_memory import MSSymmetricMemoryHandler
1564 symmetric_memory = MSSymmetricMemoryHandler()
1565 return symmetric_memory
1567 @staticmethod
1568 def get_multicore_handler():
1569 """Create and return a MindSpore multicore handler instance."""
1570 # pylint: disable=C0415
1571 from hyper_parallel.platform.mindspore.multicore import MSMulticoreHandler
1572 return MSMulticoreHandler()
1574 def new_stream(self):
1575 return ms.runtime.Stream()
1577 def get_stream_context(self):
1578 return ms.runtime.StreamCtx
1580 @staticmethod
1581 def all_gather_object(object_list, obj, group=None) -> None:
1582 """
1583 Gathers objects from the given group into object list.
1585 Args:
1586 object_list (list[Any]): Define the output list, which size equal to the size of group.
1587 obj (Any): The object on current rank and in given process group.
1588 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1589 global group.
1591 Returns:
1592 None. Objs are gathered into ``object_list``.
1593 """
1594 dist.all_gather_object(object_list, obj, group)
1596 @staticmethod
1597 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1598 """
1599 Synchronize all processes in the given communication group.
1601 Args:
1602 group (str, optional): The communication group to work on. Default is ``None``,
1603 meaning the default world group.
1604 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1605 device_ids (list[int], optional): Reserved parameter on Ascend. Default: ``None``.
1607 Returns:
1608 CommHandle if ``async_op`` is True; otherwise ``None``.
1609 """
1610 return dist.barrier(group, async_op, device_ids)
1612 @staticmethod
1613 def init_process_group(
1614 backend: str = None,
1615 *,
1616 init_method: Optional[str] = None,
1617 timeout: Optional[timedelta] = None,
1618 world_size: int = -1,
1619 rank: int = -1,
1620 store: TCPStore = None,
1621 pg_options=None,
1622 device_id=None
1623 ) -> None:
1624 """
1625 Initialize global process group.
1627 Args:
1628 backend (str): The backend used to init process group. Default is ``"hccl"`` and now only support hccl.
1629 init_method (str, optional): URL specifying how to initialize the process group. Default is ``None``.
1630 timeout (timedelta, optional): Timeout for API executed. Default is ``None``.
1631 world_size (int): Number of processes. Default is ``-1``.
1632 rank (int, optional): Rank of the current process. Default is ``-1``.
1633 store (Store, optional): An object that stores key/value data, facilitating the exchange of inter-process
1634 communication addresses and connection information. Default is ``None``. Currently, only the
1635 ``TCPStore`` type is supported.
1636 pg_options (ProcessGroupOptions, optional): Reserved parameter. Current not take effect.
1637 device_id (int, optional): Reserved parameter. Current not take effect.
1638 """
1639 if backend is None:
1640 backend = "hccl"
1641 try:
1642 if dist.is_initialized():
1643 return
1644 except AttributeError:
1645 pass
1646 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1647 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1649 @staticmethod
1650 def destroy_process_group(group: Optional[str] = None) -> None:
1651 """
1652 Destroy given process group.
1654 Args:
1655 group (str, optional): Specify the group to destroy. Default: ``None`` means ``hccl_world_group``. If group
1656 is None or "hccl_world_group", destroy global process group and all process groups relative to global
1657 process group.
1658 """
1659 if group in EXISTING_COMM_GROUPS.values():
1660 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1661 for k in keys_to_destroy:
1662 del EXISTING_COMM_GROUPS[k]
1663 dist.destroy_process_group(group)
1665 @staticmethod
1666 def get_process_group_ranks(group: Optional[str] = None) -> list[int]:
1667 """
1668 Get all ranks in given process group.
1670 Args:
1671 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1673 Returns:
1674 List[int]: List of ranks in given process group.
1675 """
1676 return dist.get_process_group_ranks(group)
1678 @staticmethod
1679 def get_backend(group: Optional[str] = None) -> str:
1680 """
1681 Get the backend of given process group.
1683 Args:
1684 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1686 Returns:
1687 str: The backend of the group.
1688 """
1689 return dist.get_backend(group)
1691 @staticmethod
1692 def split_group(parent_pg: Optional[str] = None,
1693 split_ranks: Optional[list] = None,
1694 timeout: Optional[timedelta] = None,
1695 pg_options: Optional[Any] = None,
1696 group_desc: Optional[str] = None,
1697 ) -> str:
1698 """
1699 Create split group for a specific group rank in split_ranks, which group contains current rank id.
1701 Args:
1702 parent_pg (str, Optional): A process group which the goal group split from.
1703 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1704 timeout (Optional[timedelta]): Timeout for API executed. Default is ``None``.
1705 pg_options (Optional[Any]): Backend-specific group options. MindSpore can use
1706 ``{"hccl_config": {"hccl_op_expansion_mode": "AIV"}}`` to request AIV mode.
1707 group_desc (Optional[str]): Description of process group.
1709 Returns:
1710 str: The split group name.
1711 """
1712 if split_ranks is None or len(split_ranks) == 0:
1713 raise ValueError("split_ranks cannot be None or empty")
1715 rank_id = MindSporePlatform.get_rank()
1716 for split_rank in split_ranks:
1717 if rank_id in split_rank:
1718 world_group = MindSporePlatform._maybe_reuse_world_group(split_rank)
1719 if world_group is not None:
1720 return world_group
1721 split_group = MindSporePlatform.get_created_group(split_rank)
1722 if split_group:
1723 return split_group
1724 group_name = str(tuple(sorted(split_rank)))
1725 MindSporePlatform._create_group_with_options(group_name, split_rank, pg_options=pg_options)
1726 EXISTING_COMM_GROUPS[group_name] = group_name
1727 return group_name
1728 raise ValueError(f"Split group invalid rank, the Split_ranks {split_ranks} does not contain current rank"
1729 f" {rank_id}")
1731 @staticmethod
1732 def get_group_local_rank(group=None) -> int:
1733 """get group local rank id."""
1734 return dist.get_group_rank(group, MindSporePlatform.get_rank())
1736 @staticmethod
1737 def get_group_rank(group=None) -> int:
1738 return MindSporePlatform.get_group_local_rank(group)
1740 @staticmethod
1741 def no_grad():
1742 return _no_grad()
1744 @staticmethod
1745 def preserve_version_counter(tensor):
1746 from mindspore.common.api import _unsafe_preserve_version_counter # pylint: disable=C0415
1747 return _unsafe_preserve_version_counter(tensor)
1749 @staticmethod
1750 def relu(tensor):
1751 return mint.nn.functional.relu(tensor)
1753 @staticmethod
1754 def cat(tensors, dim=0):
1755 return mint.cat(tensors, dim=dim)
1757 @staticmethod
1758 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1759 return mint.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1761 def get_current_stream(self):
1762 return ms.runtime.current_stream()
1764 def new_event(self):
1765 return ms.runtime.Event()
1767 def tree_map(self, fn, tree):
1768 """
1769 Apply fn to each leaf in a nested structure (list / tuple / dict),
1770 preserving the original structure.
1771 """
1772 if isinstance(tree, dict):
1773 return type(tree)(
1774 (k, self.tree_map(fn, v)) for k, v in tree.items()
1775 )
1777 if isinstance(tree, tuple):
1778 return tuple(self.tree_map(fn, v) for v in tree)
1780 if isinstance(tree, list):
1781 return [self.tree_map(fn, v) for v in tree]
1783 # leaf
1784 return fn(tree)
1786 @staticmethod
1787 def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
1788 return module.register_forward_pre_hook(hook, with_kwargs=with_kwargs)
1790 @staticmethod
1791 def register_full_backward_hook(module, hook, prepend=False):
1792 return module.register_backward_hook(hook)
1794 @staticmethod
1795 def register_full_backward_pre_hook(module, hook, prepend=False):
1796 return module.register_backward_pre_hook(hook)
1798 @property
1799 def checkpoint(self):
1800 return ms.recompute
1802 @staticmethod
1803 def checkpoint_wrapper(module, **checkpoint_kwargs):
1804 # pylint: disable=C0415
1805 from hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1806 return ckpt_wrapper(module, **checkpoint_kwargs)
1808 @staticmethod
1809 def checkpoint_exclude_wrapper(module: Any) -> Any:
1810 """Wrap a Cell or callable whose activations should not be recomputed.
1812 Args:
1813 module: MindSpore Cell or callable to exclude from checkpoint replay.
1815 Returns:
1816 The platform-specific checkpoint exclusion wrapper.
1817 """
1818 # pylint: disable=C0415
1819 from hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_exclude_wrapper import checkpoint_exclude_wrapper
1820 return checkpoint_exclude_wrapper(module)
1822 @staticmethod
1823 def swap_wrapper(module, policy_fn=None, group_swap=False):
1824 # pylint: disable=C0415
1825 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_wrapper
1826 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1828 @staticmethod
1829 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1830 # pylint: disable=C0415
1831 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_tensor_wrapper
1832 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1834 @staticmethod
1835 def get_class_activation_wrapper():
1836 # pylint: disable=C0415
1837 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper
1838 return ActivationWrapper
1840 @property
1841 def noop_context_fn(self):
1842 return null_context_fn
1844 @staticmethod
1845 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1846 # pylint: disable=C0415
1847 from hyper_parallel.platform.mindspore.activation_checkpoint.sac import create_selective_checkpoint_contexts
1848 return create_selective_checkpoint_contexts(policy_fn_or_list,
1849 allow_cache_entry_mutation=allow_cache_entry_mutation,
1850 group_swap=group_swap)
1852 @staticmethod
1853 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1854 # pylint: disable=C0415
1855 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import AsyncSaveOnCpu
1856 return AsyncSaveOnCpu(policy_fn=policy_fn, group_swap=group_swap)
1858 @staticmethod
1859 def recompute_handle_collector_ctx():
1860 # pylint: disable=C0415
1861 from mindspore.common.recompute import _recompute_handle_collector_ctx
1862 return _recompute_handle_collector_ctx()
1864 @staticmethod
1865 def recompute_handle(handle, session_id):
1866 return handle.recompute(session_id)
1868 @staticmethod
1869 def recompute_session_ctx(session_id, retain_on_unpack=False):
1870 # pylint: disable=C0415
1871 from mindspore.common.recompute import _recompute_session_ctx
1872 return _recompute_session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack)
1874 @staticmethod
1875 def clear_recompute_session(session_id):
1876 # pylint: disable=C0415
1877 from mindspore.common.recompute import _clear_recompute_session
1878 return _clear_recompute_session(session_id)
1880 _MS_DEVICE_MAP = {
1881 "npu": "Ascend",
1882 "ascend": "Ascend",
1883 "gpu": "GPU",
1884 "cpu": "cpu",
1885 "": "cpu",
1886 }
1888 @staticmethod
1889 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1890 """Allocate an uninitialized 1-D tensor buffer."""
1891 if pin_memory:
1892 return mint.empty((numel,), dtype=dtype, device="cpu", pin_memory=True)
1893 if device is None:
1894 return mint.empty((numel,), dtype=dtype)
1895 device_type = str(device).split(":", maxsplit=1)[0].lower()
1896 ms_device = MindSporePlatform._MS_DEVICE_MAP.get(device_type)
1897 if ms_device is None:
1898 raise ValueError(
1899 f"Unsupported device type '{device_type}' for MindSpore; "
1900 f"supported: {sorted(MindSporePlatform._MS_DEVICE_MAP)}"
1901 )
1902 if ms_device == "cpu":
1903 return mint.empty((numel,), dtype=dtype, device="cpu")
1904 return mint.empty((numel,), dtype=dtype, device=ms_device)
1906 @staticmethod
1907 def get_element_size(tensor):
1908 """Get Tensor Element Size"""
1909 return tensor.itemsize
1911 @staticmethod
1912 def tensor_to_numpy(tensor) -> np.ndarray:
1913 """Convert MindSpore tensor to numpy array."""
1914 return tensor.asnumpy()
1916 @staticmethod
1917 def from_numpy(np_array):
1918 """Create a host (CPU) MindSpore tensor from a numpy array."""
1919 return ms.from_numpy(np_array)
1921 @staticmethod
1923 def clip_grad_norm_(
1924 parameters, max_norm, norm_type=2.0,
1925 error_if_nonfinite=False, foreach=None,
1926 ):
1927 raise NotImplementedError(
1928 "clip_grad_norm_ is not yet supported on MindSpore"
1929 )
1931 @property
1932 def meta_device(self):
1933 return "meta"
1935 def init_on_device(self, device, include_buffers=False):
1936 return _init_on_device(device, include_buffers=include_buffers)
1938 def cast_fp_tensor(self, dtype, x):
1939 """
1940 Cast floating-point tensor to target dtype if applicable.
1941 """
1942 if (
1943 not isinstance(x, ms.Tensor)
1944 or not ms.ops.is_floating_point(x)
1945 or x.dtype == dtype
1946 ):
1947 return x
1948 return x.to(dtype)
1950 def apply_to_tensors(self, fn, container):
1951 """Recursively apply to all tensor in different kinds of container types."""
1953 def apply(x):
1954 if isinstance(x, ms.Tensor):
1955 return fn(x)
1956 if hasattr(x, "__dataclass_fields__"):
1957 dc = dataclasses.replace(x)
1958 changes = {
1959 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
1960 }
1961 return dataclasses.replace(dc, **changes)
1962 if isinstance(x, OrderedDict):
1963 od = x.__class__()
1964 for key, value in x.items():
1965 od[key] = apply(value)
1966 return od
1967 if isinstance(x, dict):
1968 return {key: apply(value) for key, value in x.items()}
1969 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
1970 res = (apply(el) for el in x)
1971 return type(x)(*res)
1972 if isinstance(x, (list, tuple, set)):
1973 return type(x)(apply(el) for el in x)
1974 return x
1976 return apply(container)
1978 @staticmethod
1979 def profiler_record(name):
1980 """Profiler context manager for recording operations using mindspore.profiler."""
1981 return ms.profiler.common.record_function.RecordFunction(name)
1983 def str_to_dtype(self, dtype_str: str) -> Any:
1984 """Resolve checkpoint dtype strings (``mindspore.*`` or short ``str(Tensor.dtype)`` e.g. ``Float32``)."""
1985 if "." in dtype_str:
1986 prefix, name = dtype_str.split(".", 1)
1987 if prefix == "mindspore":
1988 return getattr(ms, name)
1989 dtype = getattr(ms, dtype_str.lower(), None)
1990 if dtype is not None:
1991 return dtype
1992 raise ValueError(
1993 f"Expected dtype string like 'mindspore.float32' or 'Float32', got {dtype_str!r}."
1994 )
1996 def list_to_size(self, size_list: list[int]) -> tuple[int, ...]:
1997 return tuple(size_list)
1999 @staticmethod
2000 def _maybe_reuse_world_group(rank_list):
2001 """Reuse the default world group for full-world rank lists."""
2002 normalized = tuple(sorted(rank_list))
2003 world_ranks = tuple(range(MindSporePlatform.get_world_size()))
2004 if normalized != world_ranks:
2005 return None
2007 EXISTING_COMM_GROUPS[str(normalized)] = GlobalComm.WORLD_COMM_GROUP
2008 return GlobalComm.WORLD_COMM_GROUP