Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / platform.py: 41%
848 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""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, group=None, async_op=False):
1317 handle = dist.broadcast(data, src, group, async_op)
1318 if async_op:
1319 handle.wait()
1320 return data
1322 @staticmethod
1323 def reduce_scatter_tensor(data, group_info, async_op=False):
1324 group_name = group_info if isinstance(group_info, str) else group_info.group_name
1325 rank_size = get_group_size(group_name) if isinstance(group_info, str) else group_info.rank_size
1326 output_shape = list(data.shape)
1327 output_shape[0] //= rank_size
1328 return _mindspore_reduce_scatter_single(data, output_shape, group_name, async_op=async_op)
1330 @staticmethod
1331 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1332 return _mindspore_reduce_scatter_single(input_tensor, output_shape, group, async_op=async_op)
1334 @staticmethod
1335 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1336 return _mindspore_all_to_all_single(input_tensor, output_shape, group, async_op=async_op)
1338 @staticmethod
1339 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1340 handle_box=None):
1341 return _MSAsyncAllGatherFunction.apply(
1342 x, work, out_perm, group, world_size, gather_dim, handle_box
1343 )
1345 @staticmethod
1346 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim, # pylint: disable=unused-argument
1347 handle_box=None):
1348 return _MSAsyncA2AFunction.apply(
1349 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1350 )
1352 @staticmethod
1353 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1354 """Launch an asynchronous, differentiable all-to-all-single.
1356 Token a2a entry point used by ``CommComputeOverlap``-driven MoE
1357 wrappers. The kernel is queued on the HCCL group's stream and
1358 the host returns immediately, so the calling thread can proceed
1359 to the next sync hook (notify + rendezvous) before the
1360 collective finishes — this is what enables the comm/compute
1361 overlap window on the paired thread.
1363 Args:
1364 input_tensor: **1-D** tensor — the caller is responsible for
1365 flattening multi-dim inputs beforehand.
1366 input_splits: ``list[int]`` — **element** counts sent to each
1367 rank (not row counts). For an originally
1368 ``(N, D)`` tensor, each entry is ``rows_i * D``.
1369 output_splits: ``list[int]`` — element counts received from each rank.
1370 group: Process group.
1372 Returns:
1373 ``AsyncCollectiveTensor`` of shape ``(sum(output_splits),)`` that
1374 defers ``CommHandle.wait()`` to the first consumer op via
1375 ``__ms_dispatch__``.
1377 Raises:
1378 ValueError: if ``input_tensor`` is not 1-D.
1380 Note:
1381 The 1-D + element-count contract diverges from the Torch
1382 implementation (which accepts N-D input + row-count splits).
1383 The divergence is intentional for now: it lets the MS path
1384 call the inner primitive directly and avoid the cross-stream
1385 race that ``comm_func.all_to_all_single``'s trailing reshape
1386 triggers under ``MS_ALLOC_CONF=memory_tracker:True`` —
1387 see :meth:`_MSAsyncA2ALazyBwd._issue_async_a2a`.
1388 """
1389 if input_tensor.ndim != 1:
1390 raise ValueError(
1391 "MindSporePlatform.differentiable_all_to_all_single_async requires a 1-D "
1392 f"input_tensor (got ndim={input_tensor.ndim}, shape={tuple(input_tensor.shape)}). "
1393 "Flatten the tensor and convert row-count splits to element counts before calling."
1394 )
1395 return _MSAsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1397 @staticmethod
1398 def differentiable_sync_hook(x, hook_name: str, coordinator):
1399 """Fire a HookCoordinator rendezvous on forward and backward.
1401 Args:
1402 x: Input tensor — returned unchanged.
1403 hook_name: One of:
1404 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1405 full rendezvous on both directions.
1406 * ``"CHUNK_START"`` — chunk-entry hook on
1407 forward; pairs with ``D_LAST.bwd`` so the
1408 BWD thread's combine.bwd of the last layer
1409 is bracketed by a barrier-synced sync point.
1410 Skipped on backward.
1411 * ``"D_LAST"`` — closing D of the last MoE
1412 layer in a chunk. Forward: ``notify_dispatched``
1413 only (no Attention follows so rendezvous is
1414 skipped). Backward: full rendezvous via D's
1415 BWD role; paired with ``CHUNK_START`` on FWD.
1416 coordinator: The :class:`HookCoordinator` driving the
1417 rendezvous protocol.
1419 Returns:
1420 ``x`` unchanged.
1422 Note:
1423 Two-thread compatibility on MindSpore PyNative is not yet
1424 fully verified. The HookCoordinator + ``_Function``
1425 primitives are individually thread-safe, but the
1426 interaction with MindSpore's autograd execution model
1427 under ``threading.Thread`` should be PoC-tested before
1428 production use.
1429 """
1430 return _MSSyncHookFunction.apply(x, hook_name, coordinator)
1432 @staticmethod
1433 def parameters_dict(cell: Cell):
1434 return cell.parameters_and_names()
1436 @staticmethod
1437 def get_tensor_transform():
1438 return _tensor_transform
1440 @staticmethod
1441 def construct_strided_slice(x, begin, end, stride):
1442 return ms.ops.strided_slice(x, begin, end, stride)
1444 @staticmethod
1445 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1446 # pylint: disable=C0415
1447 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import _MicroBatch
1448 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1450 @staticmethod
1451 def get_model_state_dict(model, *, options=None):
1452 raise NotImplementedError(
1453 "get_model_state_dict is not yet supported on MindSpore"
1454 )
1456 @staticmethod
1457 def save_checkpoint(cell: Union[Cell, dict], file_path: str, ckpt_format: str = "safetensors") -> None:
1458 if isinstance(cell, dict):
1459 save_dict = {}
1460 for k, v in cell.items():
1461 if isinstance(v, Parameter):
1462 save_dict[k] = v
1463 elif isinstance(v, Tensor):
1464 save_dict[k] = Parameter(v, name=k)
1465 else:
1466 save_dict[k] = v
1467 else:
1468 save_dict = cell._params
1469 ms.save_checkpoint(save_obj=save_dict, ckpt_file_name=file_path, format=ckpt_format)
1471 @staticmethod
1472 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
1473 return ms.load_checkpoint(ckpt_file_name=file_path, format=ckpt_format)
1475 @staticmethod
1476 def get_symmetric_memory_handler():
1477 # pylint: disable=C0415
1478 from hyper_parallel.platform.mindspore.symmetric_memory import MSSymmetricMemoryHandler
1479 symmetric_memory = MSSymmetricMemoryHandler()
1480 return symmetric_memory
1482 @staticmethod
1483 def get_multicore_handler():
1484 """Create and return a MindSpore multicore handler instance."""
1485 # pylint: disable=C0415
1486 from hyper_parallel.platform.mindspore.multicore import MSMulticoreHandler
1487 return MSMulticoreHandler()
1489 def new_stream(self):
1490 return ms.runtime.Stream()
1492 def get_stream_context(self):
1493 return ms.runtime.StreamCtx
1495 @staticmethod
1496 def all_gather_object(object_list, obj, group=None) -> None:
1497 """
1498 Gathers objects from the given group into object list.
1500 Args:
1501 object_list (list[Any]): Define the output list, which size equal to the size of group.
1502 obj (Any): The object on current rank and in given process group.
1503 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1504 global group.
1506 Returns:
1507 None. Objs are gathered into ``object_list``.
1508 """
1509 dist.all_gather_object(object_list, obj, group)
1511 @staticmethod
1512 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1513 """
1514 Synchronize all processes in the given communication group.
1516 Args:
1517 group (str, optional): The communication group to work on. Default is ``None``,
1518 meaning the default world group.
1519 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1520 device_ids (list[int], optional): Reserved parameter on Ascend. Default: ``None``.
1522 Returns:
1523 CommHandle if ``async_op`` is True; otherwise ``None``.
1524 """
1525 return dist.barrier(group, async_op, device_ids)
1527 @staticmethod
1528 def init_process_group(
1529 backend: str = None,
1530 *,
1531 init_method: Optional[str] = None,
1532 timeout: Optional[timedelta] = None,
1533 world_size: int = -1,
1534 rank: int = -1,
1535 store: TCPStore = None,
1536 pg_options=None,
1537 device_id=None
1538 ) -> None:
1539 """
1540 Initialize global process group.
1542 Args:
1543 backend (str): The backend used to init process group. Default is ``"hccl"`` and now only support hccl.
1544 init_method (str, optional): URL specifying how to initialize the process group. Default is ``None``.
1545 timeout (timedelta, optional): Timeout for API executed. Default is ``None``.
1546 world_size (int): Number of processes. Default is ``-1``.
1547 rank (int, optional): Rank of the current process. Default is ``-1``.
1548 store (Store, optional): An object that stores key/value data, facilitating the exchange of inter-process
1549 communication addresses and connection information. Default is ``None``. Currently, only the
1550 ``TCPStore`` type is supported.
1551 pg_options (ProcessGroupOptions, optional): Reserved parameter. Current not take effect.
1552 device_id (int, optional): Reserved parameter. Current not take effect.
1553 """
1554 if backend is None:
1555 backend = "hccl"
1556 try:
1557 if dist.is_initialized():
1558 return
1559 except AttributeError:
1560 pass
1561 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1562 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1564 @staticmethod
1565 def destroy_process_group(group: Optional[str] = None) -> None:
1566 """
1567 Destroy given process group.
1569 Args:
1570 group (str, optional): Specify the group to destroy. Default: ``None`` means ``hccl_world_group``. If group
1571 is None or "hccl_world_group", destroy global process group and all process groups relative to global
1572 process group.
1573 """
1574 if group in EXISTING_COMM_GROUPS.values():
1575 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1576 for k in keys_to_destroy:
1577 del EXISTING_COMM_GROUPS[k]
1578 dist.destroy_process_group(group)
1580 @staticmethod
1581 def get_process_group_ranks(group: Optional[str] = None) -> list[int]:
1582 """
1583 Get all ranks in given process group.
1585 Args:
1586 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1588 Returns:
1589 List[int]: List of ranks in given process group.
1590 """
1591 return dist.get_process_group_ranks(group)
1593 @staticmethod
1594 def get_backend(group: Optional[str] = None) -> str:
1595 """
1596 Get the backend of given process group.
1598 Args:
1599 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1601 Returns:
1602 str: The backend of the group.
1603 """
1604 return dist.get_backend(group)
1606 @staticmethod
1607 def split_group(parent_pg: Optional[str] = None,
1608 split_ranks: Optional[list] = None,
1609 timeout: Optional[timedelta] = None,
1610 pg_options: Optional[Any] = None,
1611 group_desc: Optional[str] = None,
1612 ) -> str:
1613 """
1614 Create split group for a specific group rank in split_ranks, which group contains current rank id.
1616 Args:
1617 parent_pg (str, Optional): A process group which the goal group split from.
1618 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1619 timeout (Optional[timedelta]): Timeout for API executed. Default is ``None``.
1620 pg_options (Optional[Any]): Backend-specific group options. MindSpore can use
1621 ``{"hccl_config": {"hccl_op_expansion_mode": "AIV"}}`` to request AIV mode.
1622 group_desc (Optional[str]): Description of process group.
1624 Returns:
1625 str: The split group name.
1626 """
1627 if split_ranks is None or len(split_ranks) == 0:
1628 raise ValueError("split_ranks cannot be None or empty")
1630 rank_id = MindSporePlatform.get_rank()
1631 for split_rank in split_ranks:
1632 if rank_id in split_rank:
1633 world_group = MindSporePlatform._maybe_reuse_world_group(split_rank)
1634 if world_group is not None:
1635 return world_group
1636 split_group = MindSporePlatform.get_created_group(split_rank)
1637 if split_group:
1638 return split_group
1639 group_name = str(tuple(sorted(split_rank)))
1640 MindSporePlatform._create_group_with_options(group_name, split_rank, pg_options=pg_options)
1641 EXISTING_COMM_GROUPS[group_name] = group_name
1642 return group_name
1643 raise ValueError(f"Split group invalid rank, the Split_ranks {split_ranks} does not contain current rank"
1644 f" {rank_id}")
1646 @staticmethod
1647 def get_group_local_rank(group=None) -> int:
1648 """get group local rank id."""
1649 return dist.get_group_rank(group, MindSporePlatform.get_rank())
1651 @staticmethod
1652 def no_grad():
1653 return _no_grad()
1655 @staticmethod
1656 def preserve_version_counter(tensor):
1657 from mindspore.common.api import _unsafe_preserve_version_counter # pylint: disable=C0415
1658 return _unsafe_preserve_version_counter(tensor)
1660 @staticmethod
1661 def relu(tensor):
1662 return mint.nn.functional.relu(tensor)
1664 @staticmethod
1665 def cat(tensors, dim=0):
1666 return mint.cat(tensors, dim=dim)
1668 @staticmethod
1669 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1670 return mint.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1672 def get_current_stream(self):
1673 return ms.runtime.current_stream()
1675 def new_event(self):
1676 return ms.runtime.Event()
1678 def tree_map(self, fn, tree):
1679 """
1680 Apply fn to each leaf in a nested structure (list / tuple / dict),
1681 preserving the original structure.
1682 """
1683 if isinstance(tree, dict):
1684 return type(tree)(
1685 (k, self.tree_map(fn, v)) for k, v in tree.items()
1686 )
1688 if isinstance(tree, tuple):
1689 return tuple(self.tree_map(fn, v) for v in tree)
1691 if isinstance(tree, list):
1692 return [self.tree_map(fn, v) for v in tree]
1694 # leaf
1695 return fn(tree)
1697 @staticmethod
1698 def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
1699 return module.register_forward_pre_hook(hook, with_kwargs=with_kwargs)
1701 @staticmethod
1702 def register_full_backward_hook(module, hook, prepend=False):
1703 return module.register_backward_hook(hook)
1705 @staticmethod
1706 def register_full_backward_pre_hook(module, hook, prepend=False):
1707 return module.register_backward_pre_hook(hook)
1709 @property
1710 def checkpoint(self):
1711 return ms.recompute
1713 @staticmethod
1714 def checkpoint_wrapper(module, **checkpoint_kwargs):
1715 # pylint: disable=C0415
1716 from hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1717 return ckpt_wrapper(module, **checkpoint_kwargs)
1719 @staticmethod
1720 def swap_wrapper(module, policy_fn=None, group_swap=False):
1721 # pylint: disable=C0415
1722 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_wrapper
1723 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1725 @staticmethod
1726 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1727 # pylint: disable=C0415
1728 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_tensor_wrapper
1729 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1731 @staticmethod
1732 def get_class_activation_wrapper():
1733 # pylint: disable=C0415
1734 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper
1735 return ActivationWrapper
1737 @property
1738 def noop_context_fn(self):
1739 return null_context_fn
1741 @staticmethod
1742 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1743 # pylint: disable=C0415
1744 from hyper_parallel.platform.mindspore.activation_checkpoint.sac import create_selective_checkpoint_contexts
1745 return create_selective_checkpoint_contexts(policy_fn_or_list,
1746 allow_cache_entry_mutation=allow_cache_entry_mutation,
1747 group_swap=group_swap)
1749 @staticmethod
1750 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1751 # pylint: disable=C0415
1752 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import AsyncSaveOnCpu
1753 return AsyncSaveOnCpu(policy_fn=policy_fn, group_swap=group_swap)
1755 @staticmethod
1756 def recompute_handle_collector_ctx():
1757 # pylint: disable=C0415
1758 from mindspore.common.recompute import _recompute_handle_collector_ctx
1759 return _recompute_handle_collector_ctx()
1761 @staticmethod
1762 def recompute_handle(handle, session_id):
1763 return handle.recompute(session_id)
1765 @staticmethod
1766 def recompute_session_ctx(session_id, retain_on_unpack=False):
1767 # pylint: disable=C0415
1768 from mindspore.common.recompute import _recompute_session_ctx
1769 return _recompute_session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack)
1771 @staticmethod
1772 def clear_recompute_session(session_id):
1773 # pylint: disable=C0415
1774 from mindspore.common.recompute import _clear_recompute_session
1775 return _clear_recompute_session(session_id)
1777 _MS_DEVICE_MAP = {
1778 "npu": "Ascend",
1779 "ascend": "Ascend",
1780 "gpu": "GPU",
1781 "cpu": "cpu",
1782 "": "cpu",
1783 }
1785 @staticmethod
1786 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1787 """Allocate an uninitialized 1-D tensor buffer."""
1788 if pin_memory:
1789 return mint.empty((numel,), dtype=dtype, device="cpu", pin_memory=True)
1790 if device is None:
1791 return mint.empty((numel,), dtype=dtype)
1792 device_type = str(device).split(":", maxsplit=1)[0].lower()
1793 ms_device = MindSporePlatform._MS_DEVICE_MAP.get(device_type)
1794 if ms_device is None:
1795 raise ValueError(
1796 f"Unsupported device type '{device_type}' for MindSpore; "
1797 f"supported: {sorted(MindSporePlatform._MS_DEVICE_MAP)}"
1798 )
1799 if ms_device == "cpu":
1800 return mint.empty((numel,), dtype=dtype, device="cpu")
1801 return mint.empty((numel,), dtype=dtype, device=ms_device)
1803 @staticmethod
1804 def get_element_size(tensor):
1805 """Get Tensor Element Size"""
1806 return tensor.itemsize
1808 @staticmethod
1809 def tensor_to_numpy(tensor) -> np.ndarray:
1810 """Convert MindSpore tensor to numpy array."""
1811 return tensor.asnumpy()
1813 @staticmethod
1814 def from_numpy(np_array):
1815 """Create a host (CPU) MindSpore tensor from a numpy array."""
1816 return ms.from_numpy(np_array)
1818 @staticmethod
1820 def clip_grad_norm_(
1821 parameters, max_norm, norm_type=2.0,
1822 error_if_nonfinite=False, foreach=None,
1823 ):
1824 raise NotImplementedError(
1825 "clip_grad_norm_ is not yet supported on MindSpore"
1826 )
1828 @property
1829 def meta_device(self):
1830 return "meta"
1832 def init_on_device(self, device, include_buffers=False):
1833 return _init_on_device(device, include_buffers=include_buffers)
1835 def cast_fp_tensor(self, dtype, x):
1836 """
1837 Cast floating-point tensor to target dtype if applicable.
1838 """
1839 if (
1840 not isinstance(x, ms.Tensor)
1841 or not ms.ops.is_floating_point(x)
1842 or x.dtype == dtype
1843 ):
1844 return x
1845 return x.to(dtype)
1847 def apply_to_tensors(self, fn, container):
1848 """Recursively apply to all tensor in different kinds of container types."""
1850 def apply(x):
1851 if isinstance(x, ms.Tensor):
1852 return fn(x)
1853 if hasattr(x, "__dataclass_fields__"):
1854 dc = dataclasses.replace(x)
1855 changes = {
1856 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
1857 }
1858 return dataclasses.replace(dc, **changes)
1859 if isinstance(x, OrderedDict):
1860 od = x.__class__()
1861 for key, value in x.items():
1862 od[key] = apply(value)
1863 return od
1864 if isinstance(x, dict):
1865 return {key: apply(value) for key, value in x.items()}
1866 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
1867 res = (apply(el) for el in x)
1868 return type(x)(*res)
1869 if isinstance(x, (list, tuple, set)):
1870 return type(x)(apply(el) for el in x)
1871 return x
1873 return apply(container)
1875 @staticmethod
1876 def profiler_record(name):
1877 """Profiler context manager for recording operations using mindspore.profiler."""
1878 return ms.profiler.common.record_function.RecordFunction(name)
1880 def str_to_dtype(self, dtype_str: str) -> Any:
1881 """Resolve checkpoint dtype strings (``mindspore.*`` or short ``str(Tensor.dtype)`` e.g. ``Float32``)."""
1882 if "." in dtype_str:
1883 prefix, name = dtype_str.split(".", 1)
1884 if prefix == "mindspore":
1885 return getattr(ms, name)
1886 dtype = getattr(ms, dtype_str.lower(), None)
1887 if dtype is not None:
1888 return dtype
1889 raise ValueError(
1890 f"Expected dtype string like 'mindspore.float32' or 'Float32', got {dtype_str!r}."
1891 )
1893 def list_to_size(self, size_list: list[int]) -> tuple[int, ...]:
1894 return tuple(size_list)
1896 @staticmethod
1897 def _maybe_reuse_world_group(rank_list):
1898 """Reuse the default world group for full-world rank lists."""
1899 normalized = tuple(sorted(rank_list))
1900 world_ranks = tuple(range(MindSporePlatform.get_world_size()))
1901 if normalized != world_ranks:
1902 return None
1904 EXISTING_COMM_GROUPS[str(normalized)] = GlobalComm.WORLD_COMM_GROUP
1905 return GlobalComm.WORLD_COMM_GROUP