Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / platform.py: 42%
868 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""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 get_rank():
937 """
938 Get the rank of the current process in the distributed group.
940 Returns:
941 int: The rank of the current process.
942 """
943 return get_rank_id()
945 @staticmethod
946 def get_global_rank(group, group_rank):
947 """
948 Get the global rank from a group rank.
950 Args:
951 group (str): The process group name.
952 group_rank (int): The rank within the group.
954 Returns:
955 int: The global rank.
956 """
957 return dist.get_global_rank(group, group_rank)
959 @staticmethod
960 def get_world_size():
961 """
962 Get the total number of processes in the distributed group.
964 Returns:
965 int: The world size.
966 """
967 return get_group_size()
969 @staticmethod
970 def get_op_name(func):
971 """
972 Extract the operation name from a function.
974 Args:
975 func: The function to extract the name from.
977 Returns:
978 str: The operation name.
979 """
980 return func.name
982 @staticmethod
983 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
984 data = _ensure_contiguous(data)
985 # rank_list is accepted for torch parity; MindSpore keeps the existing group order.
986 output, _ = comm_func.all_gather_into_tensor(None, data, group=group)
987 if concat_dim == 0:
988 return output
989 output_tensors = ms.ops.Split(output_num=concat_size)(output)
990 return ms.mint.concat(output_tensors, concat_dim)
992 @staticmethod
993 def chunk(data, split_dim, split_size, index):
994 return ms.ops.Split(axis=split_dim, output_num=split_size)(data)[index]
996 @staticmethod
997 def differentiable_all_to_all(input_data, output_shape, group):
998 input_data = _ensure_contiguous(input_data)
999 output_tensor, _ = comm_func.all_to_all_single(
1000 output_shape,
1001 input_data,
1002 group=group,
1003 async_op=False
1004 )
1005 return output_tensor
1007 @staticmethod
1008 def tensor_type_cast(input_data, cast_type):
1009 """Cast tensor to specified data type."""
1010 type_mapping = {
1011 'float32': ms.float32,
1012 'float16': ms.float16,
1013 'int64': ms.int64,
1014 'int32': ms.int32
1015 }
1016 if cast_type not in type_mapping:
1017 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")
1018 return input_data.to(type_mapping[cast_type])
1020 @staticmethod
1021 def differentiable_all_reduce(data, op, group):
1022 data = _ensure_contiguous(data)
1023 output, _ = comm_func.all_reduce(data, op, group)
1024 return output
1026 @staticmethod
1027 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
1028 data = _ensure_contiguous(data)
1029 if axis > 0:
1030 data = ms.mint.concat(ms.ops.Split(axis=axis, output_num=dev_num)(data), dim=0)
1031 output_tensor, _ = comm_func.reduce_scatter_tensor(None, data, 'sum', group)
1032 if op == 'avg':
1033 output_tensor = output_tensor / dev_num
1034 return output_tensor
1036 @staticmethod
1037 def init_parameters(module, stage_index):
1038 return _init_parameters(module, stage_index)
1040 # pylint: disable=W0212
1041 @staticmethod
1042 def update_param_data(param, data):
1043 """update param data"""
1044 if isinstance(param, DTensorBase):
1045 param.set_data(data)
1046 else:
1047 param._update_data(data)
1049 @staticmethod
1050 def load_into_param(param, data):
1051 copy_tensor = MindSporePlatform.empty_like(data)
1052 copy_tensor.copy_(data)
1053 if isinstance(param, DTensorBase):
1054 param.set_data(copy_tensor)
1055 else:
1056 param._update(copy_tensor)
1058 @staticmethod
1059 def get_cell_construct(cell):
1060 return cell.construct
1062 @staticmethod
1063 def get_cells_and_names(cell):
1064 return cell.cells_and_names()
1066 @staticmethod
1067 def get_modules(module):
1068 return module.cells()
1070 @staticmethod
1071 def search_parameter_by_name(cell, param_name: str):
1072 """
1073 Find the parent Module of the parameter, the parameter's name in the parent Module, and the parameter.
1074 Return value: (parent Module instance, parameter's name in parent Module, parameter object).
1075 Returns None if not found.
1076 """
1077 # Remove the "self." prefix from param_name (to maintain compatibility with original logic)
1078 param_name = param_name.replace("self.", "")
1079 # Case 1: The parameter is a direct parameter of the current Module (not in any sub-Module)
1080 if param_name in cell._params:
1081 return (cell, param_name, cell._params[param_name])
1083 # Case 2: The parameter is in a sub-Module (supports multi-level nesting, e.g., "net_b.dense1.weight")
1084 if "." in param_name:
1085 # Split into: sub-Module path + parameter name (e.g., "net_b.dense1" + "weight")
1086 cell_path, param_key = param_name.rsplit(".", 1)
1087 try:
1088 # Locate the sub-Module where the parameter resides (supports multi-level paths)
1089 target_cell = cell.get_sub_cell(cell_path)
1090 # Check if the sub-Module directly contains this parameter
1091 if param_key in target_cell._params:
1092 return target_cell, param_key, target_cell._params[param_key]
1093 except AttributeError:
1094 # Sub-Module path does not exist or the parameter is not in that sub-Module
1095 pass
1097 # Traverse all sub-Modules (recursively) to search for the parameter
1098 for _, child_cell in cell._cells.items():
1099 if isinstance(child_cell, Cell):
1100 # Recursively search within the sub-Module
1101 result = MindSporePlatform.search_parameter_by_name(child_cell, param_name)
1102 if result is not None:
1103 return result
1105 return None
1107 @staticmethod
1108 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
1109 """
1110 Modify the original parameter in a Module or sub-Module using the search result
1111 Args:
1112 cell: The cell which parameter is to update
1113 result: A tuple contains parent Module, parameter key and old parameter.
1114 new_param: New Parameter object (used to replace the original parameter)
1115 """
1116 parent_cell, param_key, _ = result
1117 # Key operation: directly modify the _params dictionary of the parent Module (original storage location)
1118 parent_cell._params[param_key] = new_param
1120 if param_key in parent_cell.__dict__:
1121 parent_cell.__dict__[param_key] = new_param
1122 parent_cell._params_list[param_key] = new_param
1123 return True
1125 @staticmethod
1126 def set_layout_into_parameter(param, layout):
1127 """Set layout in to parameter"""
1128 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel
1129 from hyper_parallel.core.dtensor.layout import _infer_slice_shape_by_layout, \
1130 _get_slice_tensor_by_layout # pylint: disable=import-outside-toplevel
1131 if isinstance(param, DTensor):
1132 raise ValueError(f"Parameter {param.name} has been configured layout, cannot be set repeatedly.")
1133 param_info = param.param_info
1134 requires_grad = param.requires_grad
1135 name = param.name
1136 slice_shape = _infer_slice_shape_by_layout(param.shape, layout)
1138 if not param.has_init:
1139 # has been init, get slice data
1140 param_dtensor = DTensor.from_local(
1141 _get_slice_tensor_by_layout(param, layout).value(), layout.mesh, layout.alias_placements
1142 )
1143 param = Parameter(param_dtensor, name=name, requires_grad=requires_grad)
1144 param.param_info = param_info
1145 else:
1146 # has not been init, need to modify init shape
1147 param.init_mode.shape = slice_shape
1148 param_dtensor = DTensor.from_local(param.init_mode, layout.mesh, layout.alias_placements)
1149 param = Parameter(param_dtensor, name=name, requires_grad=requires_grad)
1150 param.param_info = param_info
1151 return param
1153 @staticmethod
1154 def get_param_local_shape(param):
1155 """get param local shape"""
1156 if isinstance(param, DTensorBase):
1157 return param.local_shape
1158 return param.shape
1160 @staticmethod
1161 def get_param_local_data(param):
1162 """get param local shape"""
1163 if isinstance(param, DTensorBase):
1164 return param.to_local()
1165 return param
1167 @staticmethod
1168 def get_param_type_size(param):
1169 return type_size_in_bytes(param.dtype)
1171 @staticmethod
1172 def is_tensor(obj: Any) -> bool:
1173 """Return True if ``obj`` is a ``mindspore.Tensor``."""
1174 return isinstance(obj, Tensor)
1176 @staticmethod
1177 def get_tensor_storage_size(tensor: Any) -> int:
1178 """Return serialized byte size (numel * itemsize) for a MindSpore tensor."""
1179 if not MindSporePlatform.is_tensor(tensor):
1180 raise TypeError(
1181 f"MindSporePlatform.get_tensor_storage_size expects mindspore.Tensor, got {type(tensor)!r}"
1182 )
1183 return int(tensor.numel()) * int(tensor.itemsize)
1185 @staticmethod
1186 def new_zero_parameter(param_shape, param_type, requires_grad, device):
1187 param = Parameter(initializer("zeros", param_shape, param_type), requires_grad=requires_grad)
1188 if device in ("GPU", "Ascend"):
1189 return param.to(device)
1190 return param
1192 @staticmethod
1193 def new_tensor(tensor_shape, tensor_type, device):
1194 tensor = Tensor(shape=tensor_shape, dtype=tensor_type)
1195 if device in ("GPU", "Ascend"):
1196 return tensor.to(device)
1197 return tensor
1199 @staticmethod
1200 def full_like(tensor, fill_value, dtype=None):
1201 return mint.full_like(tensor, fill_value, dtype=dtype)
1203 @staticmethod
1204 def isend(tensor, dst=None, group=None, tag=0):
1205 return dist.isend(tensor, dst, group, tag)
1207 @staticmethod
1208 def irecv(tensor, src=None, group=None, tag=0):
1209 return dist.irecv(tensor, src, group, tag)
1211 @staticmethod
1212 def p2p_op(op_type, tensor, peer, group=None):
1213 # pylint: disable=C0415
1214 from mindspore.mint.distributed import P2POp
1215 return P2POp(op_type, tensor, peer, group)
1217 @staticmethod
1218 def batch_isend_irecv(p2p_ops):
1219 """Launch a peer-batched P2P group.
1221 MindSpore's ``batch_isend_irecv`` lowers the whole list to a single
1222 ``HcclBatchISendIRecv`` kernel on one comm stream and returns a list
1223 with one packaging ``CommHandle``; we hand that single handle back so
1224 callers can defer the whole batch's wait to one consumption point.
1225 A send and a recv to the same peer therefore overlap on the duplex
1226 link inside this one kernel.
1227 """
1228 # pylint: disable=C0415
1229 from mindspore.mint.distributed import batch_isend_irecv
1230 if not p2p_ops:
1231 return None
1232 handles = batch_isend_irecv(p2p_ops)
1233 return handles[0] if handles else None
1235 @staticmethod
1236 def p2p_exchange(tensor, peer_rank: int, group=None): # pylint: disable=unused-argument
1237 raise NotImplementedError(
1238 "p2p_exchange is not yet supported on the MindSpore platform."
1239 )
1241 @staticmethod
1242 def send_object_list(obj_list, dst=None, group=None):
1243 # pylint: disable=C0415
1244 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import send_object_list
1245 send_object_list(obj_list, dst, group)
1247 @staticmethod
1248 def recv_object_list(obj_list, src=None, group=None):
1249 # pylint: disable=C0415
1250 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import recv_object_list
1251 recv_object_list(obj_list, src, group)
1253 @staticmethod
1254 def set_tensor_requires_grad(input_tensor):
1255 """
1256 set requires grad flag for input tensor
1257 """
1258 input_tensor.requires_grad_()
1260 @staticmethod
1261 def _normalize_group_options(pg_options: Any) -> Any:
1262 if not isinstance(pg_options, dict) or "hccl_config" not in pg_options:
1263 return pg_options
1264 from mindspore._c_expression import GroupOptions # pylint: disable=C0415
1266 options = GroupOptions()
1267 options.hccl_config = pg_options["hccl_config"]
1268 return options
1270 @staticmethod
1271 def _create_group_with_options(group_name: str, rank_list: list[int], pg_options: Any = None) -> None:
1272 """Create a MindSpore communication group with optional backend-specific options."""
1273 if pg_options is None:
1274 new_group(rank_ids=rank_list, group=group_name)
1275 return
1276 try:
1277 new_group(
1278 rank_ids=rank_list,
1279 group=group_name,
1280 options=MindSporePlatform._normalize_group_options(pg_options),
1281 )
1282 except (ImportError, RuntimeError, TypeError, ValueError):
1283 new_group(rank_ids=rank_list, group=group_name)
1285 def _create_group(self, rank_list, pg_options: Any = None):
1286 world_group = self._maybe_reuse_world_group(rank_list)
1287 if world_group is not None:
1288 return world_group
1290 group_name = str(tuple(sorted(rank_list)))
1291 self._create_group_with_options(group_name, rank_list, pg_options=pg_options)
1292 EXISTING_COMM_GROUPS[group_name] = group_name
1293 return group_name
1295 @staticmethod
1296 def all_gather_into_tensor(data, group_info, async_op=False):
1297 group_name = group_info if isinstance(group_info, str) else group_info.group_name
1298 rank_size = get_group_size(group_name) if isinstance(group_info, str) else group_info.rank_size
1299 output_shape = list(data.shape)
1300 output_shape[0] *= rank_size
1301 return _mindspore_all_gather_single(data, output_shape, group_name, async_op=async_op)
1303 @staticmethod
1304 def all_gather_single(input_tensor, output_shape, group, async_op=False):
1305 return _mindspore_all_gather_single(input_tensor, output_shape, group, async_op=async_op)
1307 @staticmethod
1308 def all_reduce(data, group_info, async_op=False):
1309 if isinstance(group_info, str):
1310 handle = dist.all_reduce(data, group=group_info, async_op=async_op)
1311 else:
1312 handle = dist.all_reduce(data, group=group_info.group_name, async_op=async_op)
1313 return data, handle
1315 @staticmethod
1316 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
1317 if group_src is not None:
1318 ranks = MindSporePlatform.get_process_group_ranks(group)
1319 src = ranks[group_src]
1320 handle = dist.broadcast(data, src, group, async_op)
1321 if async_op:
1322 handle.wait()
1323 return data
1325 @staticmethod
1326 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
1327 group_name = group if isinstance(group, str) else getattr(group, "group_name", group)
1328 if group_src is not None:
1329 ranks = MindSporePlatform.get_process_group_ranks(group)
1330 src = ranks[group_src]
1331 if scatter_list is None:
1332 # MindSpore mint.scatter validates scatter_list on every rank; PyTorch passes None on receivers.
1333 rank_size = get_group_size(group_name)
1334 scatter_list = [output] * rank_size
1335 else:
1336 scatter_list = [c.contiguous() if hasattr(c, "is_contiguous") and not c.is_contiguous() else c
1337 for c in scatter_list]
1338 handle = dist.scatter(output, scatter_list, src, group_name, async_op=async_op)
1339 if async_op and handle is not None:
1340 handle.wait()
1341 return output
1343 @staticmethod
1344 def reduce_scatter_tensor(data, group_info, async_op=False):
1345 group_name = group_info if isinstance(group_info, str) else group_info.group_name
1346 rank_size = get_group_size(group_name) if isinstance(group_info, str) else group_info.rank_size
1347 output_shape = list(data.shape)
1348 output_shape[0] //= rank_size
1349 return _mindspore_reduce_scatter_single(data, output_shape, group_name, async_op=async_op)
1351 @staticmethod
1352 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1353 return _mindspore_reduce_scatter_single(input_tensor, output_shape, group, async_op=async_op)
1355 @staticmethod
1356 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1357 return _mindspore_all_to_all_single(input_tensor, output_shape, group, async_op=async_op)
1359 @staticmethod
1360 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1361 handle_box=None):
1362 return _MSAsyncAllGatherFunction.apply(
1363 x, work, out_perm, group, world_size, gather_dim, handle_box
1364 )
1366 @staticmethod
1367 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim, # pylint: disable=unused-argument
1368 handle_box=None):
1369 return _MSAsyncA2AFunction.apply(
1370 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1371 )
1373 @staticmethod
1374 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1375 """Launch an asynchronous, differentiable all-to-all-single.
1377 Token a2a entry point used by ``CommComputeOverlap``-driven MoE
1378 wrappers. The kernel is queued on the HCCL group's stream and
1379 the host returns immediately, so the calling thread can proceed
1380 to the next sync hook (notify + rendezvous) before the
1381 collective finishes — this is what enables the comm/compute
1382 overlap window on the paired thread.
1384 Args:
1385 input_tensor: **1-D** tensor — the caller is responsible for
1386 flattening multi-dim inputs beforehand.
1387 input_splits: ``list[int]`` — **element** counts sent to each
1388 rank (not row counts). For an originally
1389 ``(N, D)`` tensor, each entry is ``rows_i * D``.
1390 output_splits: ``list[int]`` — element counts received from each rank.
1391 group: Process group.
1393 Returns:
1394 ``AsyncCollectiveTensor`` of shape ``(sum(output_splits),)`` that
1395 defers ``CommHandle.wait()`` to the first consumer op via
1396 ``__ms_dispatch__``.
1398 Raises:
1399 ValueError: if ``input_tensor`` is not 1-D.
1401 Note:
1402 The 1-D + element-count contract diverges from the Torch
1403 implementation (which accepts N-D input + row-count splits).
1404 The divergence is intentional for now: it lets the MS path
1405 call the inner primitive directly and avoid the cross-stream
1406 race that ``comm_func.all_to_all_single``'s trailing reshape
1407 triggers under ``MS_ALLOC_CONF=memory_tracker:True`` —
1408 see :meth:`_MSAsyncA2ALazyBwd._issue_async_a2a`.
1409 """
1410 if input_tensor.ndim != 1:
1411 raise ValueError(
1412 "MindSporePlatform.differentiable_all_to_all_single_async requires a 1-D "
1413 f"input_tensor (got ndim={input_tensor.ndim}, shape={tuple(input_tensor.shape)}). "
1414 "Flatten the tensor and convert row-count splits to element counts before calling."
1415 )
1416 return _MSAsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1418 @staticmethod
1419 def differentiable_sync_hook(x, hook_name: str, coordinator):
1420 """Fire a HookCoordinator rendezvous on forward and backward.
1422 Args:
1423 x: Input tensor — returned unchanged.
1424 hook_name: One of:
1425 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1426 full rendezvous on both directions.
1427 * ``"CHUNK_START"`` — chunk-entry hook on
1428 forward; pairs with ``D_LAST.bwd`` so the
1429 BWD thread's combine.bwd of the last layer
1430 is bracketed by a barrier-synced sync point.
1431 Skipped on backward.
1432 * ``"D_LAST"`` — closing D of the last MoE
1433 layer in a chunk. Forward: ``notify_dispatched``
1434 only (no Attention follows so rendezvous is
1435 skipped). Backward: full rendezvous via D's
1436 BWD role; paired with ``CHUNK_START`` on FWD.
1437 coordinator: The :class:`HookCoordinator` driving the
1438 rendezvous protocol.
1440 Returns:
1441 ``x`` unchanged.
1443 Note:
1444 Two-thread compatibility on MindSpore PyNative is not yet
1445 fully verified. The HookCoordinator + ``_Function``
1446 primitives are individually thread-safe, but the
1447 interaction with MindSpore's autograd execution model
1448 under ``threading.Thread`` should be PoC-tested before
1449 production use.
1450 """
1451 return _MSSyncHookFunction.apply(x, hook_name, coordinator)
1453 @staticmethod
1454 def parameters_dict(cell: Cell):
1455 return cell.parameters_and_names()
1457 @staticmethod
1458 def get_tensor_transform():
1459 return _tensor_transform
1461 @staticmethod
1462 def construct_strided_slice(x, begin, end, stride):
1463 return ms.ops.strided_slice(x, begin, end, stride)
1465 @staticmethod
1466 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1467 # pylint: disable=C0415
1468 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import _MicroBatch
1469 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1471 @staticmethod
1472 def get_model_state_dict(model, *, options=None):
1473 raise NotImplementedError(
1474 "get_model_state_dict is not yet supported on MindSpore"
1475 )
1477 @staticmethod
1478 def save_checkpoint(cell: Union[Cell, dict], file_path: str, ckpt_format: str = "safetensors") -> None:
1479 if isinstance(cell, dict):
1480 save_dict = {}
1481 for k, v in cell.items():
1482 if isinstance(v, Parameter):
1483 save_dict[k] = v
1484 elif isinstance(v, Tensor):
1485 save_dict[k] = Parameter(v, name=k)
1486 else:
1487 save_dict[k] = v
1488 else:
1489 save_dict = cell._params
1490 ms.save_checkpoint(save_obj=save_dict, ckpt_file_name=file_path, format=ckpt_format)
1492 @staticmethod
1493 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
1494 return ms.load_checkpoint(ckpt_file_name=file_path, format=ckpt_format)
1496 @staticmethod
1497 def get_symmetric_memory_handler():
1498 # pylint: disable=C0415
1499 from hyper_parallel.platform.mindspore.symmetric_memory import MSSymmetricMemoryHandler
1500 symmetric_memory = MSSymmetricMemoryHandler()
1501 return symmetric_memory
1503 @staticmethod
1504 def get_multicore_handler():
1505 """Create and return a MindSpore multicore handler instance."""
1506 # pylint: disable=C0415
1507 from hyper_parallel.platform.mindspore.multicore import MSMulticoreHandler
1508 return MSMulticoreHandler()
1510 def new_stream(self):
1511 return ms.runtime.Stream()
1513 def get_stream_context(self):
1514 return ms.runtime.StreamCtx
1516 @staticmethod
1517 def all_gather_object(object_list, obj, group=None) -> None:
1518 """
1519 Gathers objects from the given group into object list.
1521 Args:
1522 object_list (list[Any]): Define the output list, which size equal to the size of group.
1523 obj (Any): The object on current rank and in given process group.
1524 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1525 global group.
1527 Returns:
1528 None. Objs are gathered into ``object_list``.
1529 """
1530 dist.all_gather_object(object_list, obj, group)
1532 @staticmethod
1533 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1534 """
1535 Synchronize all processes in the given communication group.
1537 Args:
1538 group (str, optional): The communication group to work on. Default is ``None``,
1539 meaning the default world group.
1540 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1541 device_ids (list[int], optional): Reserved parameter on Ascend. Default: ``None``.
1543 Returns:
1544 CommHandle if ``async_op`` is True; otherwise ``None``.
1545 """
1546 return dist.barrier(group, async_op, device_ids)
1548 @staticmethod
1549 def init_process_group(
1550 backend: str = None,
1551 *,
1552 init_method: Optional[str] = None,
1553 timeout: Optional[timedelta] = None,
1554 world_size: int = -1,
1555 rank: int = -1,
1556 store: TCPStore = None,
1557 pg_options=None,
1558 device_id=None
1559 ) -> None:
1560 """
1561 Initialize global process group.
1563 Args:
1564 backend (str): The backend used to init process group. Default is ``"hccl"`` and now only support hccl.
1565 init_method (str, optional): URL specifying how to initialize the process group. Default is ``None``.
1566 timeout (timedelta, optional): Timeout for API executed. Default is ``None``.
1567 world_size (int): Number of processes. Default is ``-1``.
1568 rank (int, optional): Rank of the current process. Default is ``-1``.
1569 store (Store, optional): An object that stores key/value data, facilitating the exchange of inter-process
1570 communication addresses and connection information. Default is ``None``. Currently, only the
1571 ``TCPStore`` type is supported.
1572 pg_options (ProcessGroupOptions, optional): Reserved parameter. Current not take effect.
1573 device_id (int, optional): Reserved parameter. Current not take effect.
1574 """
1575 if backend is None:
1576 backend = "hccl"
1577 try:
1578 if dist.is_initialized():
1579 return
1580 except AttributeError:
1581 pass
1582 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1583 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1585 @staticmethod
1586 def destroy_process_group(group: Optional[str] = None) -> None:
1587 """
1588 Destroy given process group.
1590 Args:
1591 group (str, optional): Specify the group to destroy. Default: ``None`` means ``hccl_world_group``. If group
1592 is None or "hccl_world_group", destroy global process group and all process groups relative to global
1593 process group.
1594 """
1595 if group in EXISTING_COMM_GROUPS.values():
1596 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1597 for k in keys_to_destroy:
1598 del EXISTING_COMM_GROUPS[k]
1599 dist.destroy_process_group(group)
1601 @staticmethod
1602 def get_process_group_ranks(group: Optional[str] = None) -> list[int]:
1603 """
1604 Get all ranks in given process group.
1606 Args:
1607 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1609 Returns:
1610 List[int]: List of ranks in given process group.
1611 """
1612 return dist.get_process_group_ranks(group)
1614 @staticmethod
1615 def get_backend(group: Optional[str] = None) -> str:
1616 """
1617 Get the backend of given process group.
1619 Args:
1620 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1622 Returns:
1623 str: The backend of the group.
1624 """
1625 return dist.get_backend(group)
1627 @staticmethod
1628 def split_group(parent_pg: Optional[str] = None,
1629 split_ranks: Optional[list] = None,
1630 timeout: Optional[timedelta] = None,
1631 pg_options: Optional[Any] = None,
1632 group_desc: Optional[str] = None,
1633 ) -> str:
1634 """
1635 Create split group for a specific group rank in split_ranks, which group contains current rank id.
1637 Args:
1638 parent_pg (str, Optional): A process group which the goal group split from.
1639 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1640 timeout (Optional[timedelta]): Timeout for API executed. Default is ``None``.
1641 pg_options (Optional[Any]): Backend-specific group options. MindSpore can use
1642 ``{"hccl_config": {"hccl_op_expansion_mode": "AIV"}}`` to request AIV mode.
1643 group_desc (Optional[str]): Description of process group.
1645 Returns:
1646 str: The split group name.
1647 """
1648 if split_ranks is None or len(split_ranks) == 0:
1649 raise ValueError("split_ranks cannot be None or empty")
1651 rank_id = MindSporePlatform.get_rank()
1652 for split_rank in split_ranks:
1653 if rank_id in split_rank:
1654 world_group = MindSporePlatform._maybe_reuse_world_group(split_rank)
1655 if world_group is not None:
1656 return world_group
1657 split_group = MindSporePlatform.get_created_group(split_rank)
1658 if split_group:
1659 return split_group
1660 group_name = str(tuple(sorted(split_rank)))
1661 MindSporePlatform._create_group_with_options(group_name, split_rank, pg_options=pg_options)
1662 EXISTING_COMM_GROUPS[group_name] = group_name
1663 return group_name
1664 raise ValueError(f"Split group invalid rank, the Split_ranks {split_ranks} does not contain current rank"
1665 f" {rank_id}")
1667 @staticmethod
1668 def get_group_local_rank(group=None) -> int:
1669 """get group local rank id."""
1670 return dist.get_group_rank(group, MindSporePlatform.get_rank())
1672 @staticmethod
1673 def get_group_rank(group=None) -> int:
1674 return MindSporePlatform.get_group_local_rank(group)
1676 @staticmethod
1677 def no_grad():
1678 return _no_grad()
1680 @staticmethod
1681 def preserve_version_counter(tensor):
1682 from mindspore.common.api import _unsafe_preserve_version_counter # pylint: disable=C0415
1683 return _unsafe_preserve_version_counter(tensor)
1685 @staticmethod
1686 def relu(tensor):
1687 return mint.nn.functional.relu(tensor)
1689 @staticmethod
1690 def cat(tensors, dim=0):
1691 return mint.cat(tensors, dim=dim)
1693 @staticmethod
1694 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1695 return mint.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1697 def get_current_stream(self):
1698 return ms.runtime.current_stream()
1700 def new_event(self):
1701 return ms.runtime.Event()
1703 def tree_map(self, fn, tree):
1704 """
1705 Apply fn to each leaf in a nested structure (list / tuple / dict),
1706 preserving the original structure.
1707 """
1708 if isinstance(tree, dict):
1709 return type(tree)(
1710 (k, self.tree_map(fn, v)) for k, v in tree.items()
1711 )
1713 if isinstance(tree, tuple):
1714 return tuple(self.tree_map(fn, v) for v in tree)
1716 if isinstance(tree, list):
1717 return [self.tree_map(fn, v) for v in tree]
1719 # leaf
1720 return fn(tree)
1722 @staticmethod
1723 def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
1724 return module.register_forward_pre_hook(hook, with_kwargs=with_kwargs)
1726 @staticmethod
1727 def register_full_backward_hook(module, hook, prepend=False):
1728 return module.register_backward_hook(hook)
1730 @staticmethod
1731 def register_full_backward_pre_hook(module, hook, prepend=False):
1732 return module.register_backward_pre_hook(hook)
1734 @property
1735 def checkpoint(self):
1736 return ms.recompute
1738 @staticmethod
1739 def checkpoint_wrapper(module, **checkpoint_kwargs):
1740 # pylint: disable=C0415
1741 from hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1742 return ckpt_wrapper(module, **checkpoint_kwargs)
1744 @staticmethod
1745 def swap_wrapper(module, policy_fn=None, group_swap=False):
1746 # pylint: disable=C0415
1747 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_wrapper
1748 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1750 @staticmethod
1751 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1752 # pylint: disable=C0415
1753 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_tensor_wrapper
1754 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1756 @staticmethod
1757 def get_class_activation_wrapper():
1758 # pylint: disable=C0415
1759 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper
1760 return ActivationWrapper
1762 @property
1763 def noop_context_fn(self):
1764 return null_context_fn
1766 @staticmethod
1767 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1768 # pylint: disable=C0415
1769 from hyper_parallel.platform.mindspore.activation_checkpoint.sac import create_selective_checkpoint_contexts
1770 return create_selective_checkpoint_contexts(policy_fn_or_list,
1771 allow_cache_entry_mutation=allow_cache_entry_mutation,
1772 group_swap=group_swap)
1774 @staticmethod
1775 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1776 # pylint: disable=C0415
1777 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import AsyncSaveOnCpu
1778 return AsyncSaveOnCpu(policy_fn=policy_fn, group_swap=group_swap)
1780 @staticmethod
1781 def recompute_handle_collector_ctx():
1782 # pylint: disable=C0415
1783 from mindspore.common.recompute import _recompute_handle_collector_ctx
1784 return _recompute_handle_collector_ctx()
1786 @staticmethod
1787 def recompute_handle(handle, session_id):
1788 return handle.recompute(session_id)
1790 @staticmethod
1791 def recompute_session_ctx(session_id, retain_on_unpack=False):
1792 # pylint: disable=C0415
1793 from mindspore.common.recompute import _recompute_session_ctx
1794 return _recompute_session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack)
1796 @staticmethod
1797 def clear_recompute_session(session_id):
1798 # pylint: disable=C0415
1799 from mindspore.common.recompute import _clear_recompute_session
1800 return _clear_recompute_session(session_id)
1802 _MS_DEVICE_MAP = {
1803 "npu": "Ascend",
1804 "ascend": "Ascend",
1805 "gpu": "GPU",
1806 "cpu": "cpu",
1807 "": "cpu",
1808 }
1810 @staticmethod
1811 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1812 """Allocate an uninitialized 1-D tensor buffer."""
1813 if pin_memory:
1814 return mint.empty((numel,), dtype=dtype, device="cpu", pin_memory=True)
1815 if device is None:
1816 return mint.empty((numel,), dtype=dtype)
1817 device_type = str(device).split(":", maxsplit=1)[0].lower()
1818 ms_device = MindSporePlatform._MS_DEVICE_MAP.get(device_type)
1819 if ms_device is None:
1820 raise ValueError(
1821 f"Unsupported device type '{device_type}' for MindSpore; "
1822 f"supported: {sorted(MindSporePlatform._MS_DEVICE_MAP)}"
1823 )
1824 if ms_device == "cpu":
1825 return mint.empty((numel,), dtype=dtype, device="cpu")
1826 return mint.empty((numel,), dtype=dtype, device=ms_device)
1828 @staticmethod
1829 def get_element_size(tensor):
1830 """Get Tensor Element Size"""
1831 return tensor.itemsize
1833 @staticmethod
1834 def tensor_to_numpy(tensor) -> np.ndarray:
1835 """Convert MindSpore tensor to numpy array."""
1836 return tensor.asnumpy()
1838 @staticmethod
1839 def from_numpy(np_array):
1840 """Create a host (CPU) MindSpore tensor from a numpy array."""
1841 return ms.from_numpy(np_array)
1843 @staticmethod
1845 def clip_grad_norm_(
1846 parameters, max_norm, norm_type=2.0,
1847 error_if_nonfinite=False, foreach=None,
1848 ):
1849 raise NotImplementedError(
1850 "clip_grad_norm_ is not yet supported on MindSpore"
1851 )
1853 @property
1854 def meta_device(self):
1855 return "meta"
1857 def init_on_device(self, device, include_buffers=False):
1858 return _init_on_device(device, include_buffers=include_buffers)
1860 def cast_fp_tensor(self, dtype, x):
1861 """
1862 Cast floating-point tensor to target dtype if applicable.
1863 """
1864 if (
1865 not isinstance(x, ms.Tensor)
1866 or not ms.ops.is_floating_point(x)
1867 or x.dtype == dtype
1868 ):
1869 return x
1870 return x.to(dtype)
1872 def apply_to_tensors(self, fn, container):
1873 """Recursively apply to all tensor in different kinds of container types."""
1875 def apply(x):
1876 if isinstance(x, ms.Tensor):
1877 return fn(x)
1878 if hasattr(x, "__dataclass_fields__"):
1879 dc = dataclasses.replace(x)
1880 changes = {
1881 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
1882 }
1883 return dataclasses.replace(dc, **changes)
1884 if isinstance(x, OrderedDict):
1885 od = x.__class__()
1886 for key, value in x.items():
1887 od[key] = apply(value)
1888 return od
1889 if isinstance(x, dict):
1890 return {key: apply(value) for key, value in x.items()}
1891 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
1892 res = (apply(el) for el in x)
1893 return type(x)(*res)
1894 if isinstance(x, (list, tuple, set)):
1895 return type(x)(apply(el) for el in x)
1896 return x
1898 return apply(container)
1900 @staticmethod
1901 def profiler_record(name):
1902 """Profiler context manager for recording operations using mindspore.profiler."""
1903 return ms.profiler.common.record_function.RecordFunction(name)
1905 def str_to_dtype(self, dtype_str: str) -> Any:
1906 """Resolve checkpoint dtype strings (``mindspore.*`` or short ``str(Tensor.dtype)`` e.g. ``Float32``)."""
1907 if "." in dtype_str:
1908 prefix, name = dtype_str.split(".", 1)
1909 if prefix == "mindspore":
1910 return getattr(ms, name)
1911 dtype = getattr(ms, dtype_str.lower(), None)
1912 if dtype is not None:
1913 return dtype
1914 raise ValueError(
1915 f"Expected dtype string like 'mindspore.float32' or 'Float32', got {dtype_str!r}."
1916 )
1918 def list_to_size(self, size_list: list[int]) -> tuple[int, ...]:
1919 return tuple(size_list)
1921 @staticmethod
1922 def _maybe_reuse_world_group(rank_list):
1923 """Reuse the default world group for full-world rank lists."""
1924 normalized = tuple(sorted(rank_list))
1925 world_ranks = tuple(range(MindSporePlatform.get_world_size()))
1926 if normalized != world_ranks:
1927 return None
1929 EXISTING_COMM_GROUPS[str(normalized)] = GlobalComm.WORLD_COMM_GROUP
1930 return GlobalComm.WORLD_COMM_GROUP