Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / platform.py: 46%
951 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""MindSpore platform api"""
16from datetime import timedelta
17from typing import Any, Optional, Sequence, 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
174def _mindspore_variable_all_gather(
175 input_tensor: Tensor,
176 output_splits: Sequence[int],
177 group: str,
178) -> Tensor:
179 """Gather variable dim-zero rows with MindSpore ``AllGatherV``."""
180 if input_tensor.ndim == 0:
181 raise ValueError("variable all-gather input must have at least one dimension")
182 splits = tuple(output_splits)
183 if not splits:
184 raise ValueError("output_splits must contain at least one group rank")
185 if any(not isinstance(rows, int) or isinstance(rows, bool) or rows < 0 for rows in splits):
186 raise ValueError(f"output_splits must contain non-negative integers, got {splits!r}")
188 feature_shape = tuple(input_tensor.shape[1:])
189 feature_numel = int(np.prod(feature_shape, dtype=np.int64)) if feature_shape else 1
190 element_splits = [rows * feature_numel for rows in splits]
191 if sum(element_splits) == 0:
192 return input_tensor.reshape((0, *feature_shape))
193 element_splits_tensor = ms.Tensor(element_splits, dtype=ms.int64)
194 gathered = ms.ops.AllGatherV(group=group)(
195 input_tensor.reshape((-1,)),
196 element_splits_tensor,
197 )
198 return gathered.reshape((sum(splits), *feature_shape))
201def _validate_variable_row_splits(
202 input_tensor: Tensor,
203 input_splits: Sequence[int],
204 output_splits: Sequence[int],
205 group: str,
206) -> tuple[list[int], list[int]]:
207 """Validate dim-zero row splits for a variable all-to-all."""
208 normalized_input = list(input_splits)
209 normalized_output = list(output_splits)
210 if not normalized_input or len(normalized_input) != len(normalized_output):
211 raise ValueError(
212 "input_splits and output_splits must be non-empty and have the same length, "
213 f"got input_splits={normalized_input!r}, output_splits={normalized_output!r}"
214 )
215 if any(
216 not isinstance(value, int) or isinstance(value, bool) or value < 0
217 for value in (*normalized_input, *normalized_output)
218 ):
219 raise ValueError(
220 "input_splits and output_splits must contain non-negative integers, "
221 f"got input_splits={normalized_input!r}, output_splits={normalized_output!r}"
222 )
223 if input_tensor.ndim == 0:
224 raise ValueError("variable all-to-all input_tensor must have at least one dimension")
225 if sum(normalized_input) != input_tensor.shape[0]:
226 raise ValueError(
227 "sum(input_splits) must equal input_tensor.shape[0], "
228 f"got sum={sum(normalized_input)}, shape={tuple(input_tensor.shape)!r}"
229 )
230 group_size = get_group_size(group)
231 if len(normalized_input) != group_size:
232 raise ValueError(
233 "split metadata length must equal the process-group size, "
234 f"got splits={len(normalized_input)}, group_size={group_size}"
235 )
236 return normalized_input, normalized_output
239def _mindspore_variable_all_to_all(
240 input_tensor: Tensor,
241 input_splits: Sequence[int],
242 output_splits: Sequence[int],
243 group: str,
244) -> Tensor:
245 """Run one synchronous N-D variable all-to-all using dim-zero row splits."""
246 output_shape = (sum(output_splits), *tuple(input_tensor.shape[1:]))
247 output, _ = comm_func.all_to_all_single(
248 output_shape,
249 input_tensor,
250 output_split_sizes=list(output_splits),
251 input_split_sizes=list(input_splits),
252 group=group,
253 async_op=False,
254 )
255 return output
258class _MSDifferentiableAllToAllSingle(_Function):
259 """Variable all-to-all with a symmetric reverse-A2A backward."""
261 @staticmethod
262 def forward( # pylint: disable=arguments-differ
263 ctx: Any,
264 input_tensor: Tensor,
265 output_splits: Sequence[int],
266 input_splits: Sequence[int],
267 group: str,
268 ) -> Tensor:
269 """Exchange dim-zero rows and record the reverse split metadata."""
270 ctx.input_splits = input_splits
271 ctx.output_splits = output_splits
272 ctx.group = group
273 return _mindspore_variable_all_to_all(
274 input_tensor,
275 input_splits,
276 output_splits,
277 group,
278 )
280 @staticmethod
281 def backward(ctx: Any, grad_output: Tensor) -> tuple:
282 """Route gradients through the reverse variable all-to-all."""
283 grad_input = _mindspore_variable_all_to_all(
284 grad_output,
285 ctx.output_splits,
286 ctx.input_splits,
287 ctx.group,
288 )
289 return grad_input, None, None, None
292class AsyncCollectiveTensor(Tensor):
293 """MindSpore Tensor subclass that defers ``CommHandle.wait()`` to
294 the first op that reads it.
296 Mimics PyTorch's ``AsyncCollectiveTensor`` using MindSpore's
297 per-tensor ``__ms_dispatch__`` mechanism. Constructed by calling
298 ``AsyncCollectiveTensor(inner_tensor, work)`` — :meth:`__new__`
299 invokes ``Tensor._make_subclass`` which (per MindSpore C++ side)
300 sets ``has_ms_dispatch=true`` on the new tensor because this class
301 defines ``__ms_dispatch__``. All subsequent ops involving this
302 tensor are routed through that callback.
304 Stream-side ``CommHandle.wait()`` (host non-blocking) means the
305 overlap window between the async a2a issue and the first consumer
306 op is preserved: the wait is only inserted on the consumer stream
307 at the consumer dispatch site, not at the a2a issue site.
309 Note:
310 Currently every op (including view ops like reshape /
311 transpose / permute) triggers ``work.wait()`` + unwrap.
312 Once MindSpore exposes schema alias annotations on
313 :class:`OpFunc` (planned per discussion with the MS team),
314 this class can mirror PyTorch's ``_is_view_op`` to keep
315 view chains lazy and stretch the overlap window further.
317 Attributes:
318 elem: The underlying regular Tensor (PyTorch's
319 ``AsyncCollectiveTensor.elem``). Returned by
320 :meth:`_wait_and_unwrap` after the wait fires
321 so downstream ops see a plain Tensor type.
322 completed: Whether ``work.wait()`` has already been
323 triggered (idempotency guard).
324 _pending_work: The async ``CommHandle`` returned by MindSpore.
325 PyTorch's equivalent class doesn't carry this
326 because PyTorch tracks tensor→work via the
327 global ``wait_tensor()`` aten op + c10d
328 registry. MindSpore has no such infra, so we
329 have to stash the handle on the wrapper itself.
330 """
332 __slots__ = ("elem", "completed", "_pending_work")
334 @staticmethod
335 def __new__(cls, inner: Tensor, work): # pylint: disable=W0613
336 """Construct a wrapper tensor sharing storage with ``inner``.
338 ``Tensor._make_subclass`` returns a tensor of class ``cls``
339 that shares storage with ``inner``. MindSpore C++ side then
340 sets ``has_ms_dispatch=true`` because ``cls`` defines
341 ``__ms_dispatch__``. Per-instance state is set in
342 :meth:`__init__`.
343 """
344 return Tensor._make_subclass(cls, inner) # pylint: disable=W0212
346 def __init__(self, inner: Tensor, work): # pylint: disable=W0231
347 """Initialize wrapper state (does NOT call ``super().__init__``).
349 Skipping ``Tensor.__init__`` is intentional: the parent
350 constructor would re-interpret ``inner`` as raw input data
351 and ``work`` as a dtype, corrupting the tensor that
352 :meth:`__new__` already built via ``Tensor._make_subclass``.
353 """
354 self.elem = inner
355 self.completed = work is None
356 self._pending_work = work
358 def _wait_and_unwrap(self) -> Tensor:
359 """Trigger ``work.wait()`` (idempotent) and return ``elem``.
361 Mirrors PyTorch's ``trigger_wait``: returns the underlying
362 regular Tensor so downstream ops see a plain ``Tensor``
363 instance, not an ``AsyncCollectiveTensor`` (avoids re-entering
364 ``__ms_dispatch__`` on every subsequent op).
365 """
366 if not self.completed:
367 work = self._pending_work
368 if work is not None:
369 work.wait() # stream-side: inserts streamWaitEvent on current stream
370 self.completed = True
371 return self.elem
373 @classmethod
374 def __ms_dispatch__(cls, func, args, kwargs=None):
375 """Per-tensor dispatch callback invoked for every op touching a
376 :class:`AsyncCollectiveTensor` instance.
378 Must be a ``@classmethod`` so MindSpore's C++-side invocation
379 (``tensor_py_reg.cc`` retrieves the attribute from the class
380 and calls it as ``handler(op_func, packed_args, kwargs)`` —
381 three positional args, no ``self`` binding) lines up with the
382 signature ``(cls, func, args, kwargs)``. Mirrors PyTorch's
383 ``__torch_dispatch__`` decoration on ``AsyncCollectiveTensor``.
385 Currently every op triggers wait + unwrap on any
386 ``AsyncCollectiveTensor`` arg, then runs the op on the
387 underlying inner tensors. This is the conservative
388 correctness-first behavior: it always defers the wait at
389 least until the first op consumes the tensor (which is later
390 than calling ``work.wait()`` immediately at a2a issue site,
391 so the overlap window is preserved across the
392 ``sync_hook("B")`` window).
394 TODO: when MindSpore exposes schema alias annotations on
395 ``func`` (the ``OpFunc`` parameter), add a fast path that
396 keeps view ops (reshape / transpose / permute / etc.) lazy
397 and only triggers wait on real data-touching ops, mirroring
398 PyTorch's ``_is_view_op`` in
399 ``torch/distributed/_functional_collectives.py``. Until that
400 annotation is available, treating views as real ops just
401 shortens the overlap window for view-heavy paths — it does
402 not affect correctness.
403 """
404 args = args if args is not None else ()
405 kwargs = kwargs if kwargs is not None else {}
406 unwrapped_args = tuple(
407 a._wait_and_unwrap() if isinstance(a, cls) else a # pylint: disable=W0212
408 for a in args
409 )
410 unwrapped_kwargs = {
411 k: (v._wait_and_unwrap() if isinstance(v, cls) else v) # pylint: disable=W0212
412 for k, v in kwargs.items()
413 }
414 return func(*unwrapped_args, **unwrapped_kwargs)
416 # ------------------------------------------------------------------
417 # Data-export overrides
418 # ------------------------------------------------------------------
419 # The methods below all read raw tensor data (or print it) and
420 # bypass ``__ms_dispatch__`` because they are Python-level methods
421 # on ``Tensor``, not MindSpore ops. Without these overrides they
422 # would access ``self``'s data buffer before the pending async a2a
423 # has finished, returning stale / uninitialized values. Each
424 # override forces a stream-side wait via ``_wait_and_unwrap`` and
425 # delegates to the same method on the underlying inner tensor.
426 #
427 # Methods deliberately NOT overridden:
428 # ``__len__`` — metadata only (returns shape[0]); no data read.
429 # ``__hash__`` — id-based on MindSpore Tensor; no data read.
430 # ``__contains__`` — uses ``(elem == self).any().item()`` which
431 # dispatches through ``==`` so wait fires
432 # transitively before the chain reaches data.
433 # ``__getitem__`` — slicing dispatches through ``__ms_dispatch__``.
434 # ``__format__`` — calls ``__repr__`` which we override.
436 def asnumpy(self):
437 """Convert to numpy ndarray; waits the pending a2a first."""
438 return self._wait_and_unwrap().asnumpy()
440 def numpy(self):
441 """Alias of :meth:`asnumpy` — same wait + unwrap path."""
442 return self._wait_and_unwrap().numpy()
444 def __array__(self, dtype=None):
445 """``np.array(t)`` protocol; waits + delegates to inner tensor."""
446 return self._wait_and_unwrap().__array__(dtype)
448 def get_bytes(self):
449 """Raw byte serialization; must wait before reading the buffer."""
450 return self._wait_and_unwrap().get_bytes()
452 def tolist(self):
453 """Convert to nested Python list; waits first."""
454 return self._wait_and_unwrap().tolist()
456 def item(self):
457 """Extract scalar value (0-d tensor); waits first."""
458 return self._wait_and_unwrap().item()
460 def __bool__(self):
461 """``bool(t)`` / ``if t:``; reads scalar value, must wait."""
462 return bool(self._wait_and_unwrap())
464 def __int__(self):
465 """``int(t)``; reads scalar value, must wait."""
466 return int(self._wait_and_unwrap())
468 def __float__(self):
469 """``float(t)``; reads scalar value, must wait."""
470 return float(self._wait_and_unwrap())
472 def __index__(self):
473 """Python index protocol; uses scalar value, must wait."""
474 return self._wait_and_unwrap().__index__()
476 def __repr__(self):
477 """Eager debug print; force wait so the printout reflects real data.
479 Mirrors PyTorch's ``AsyncCollectiveTensor.__repr__`` style by
480 labelling the wrapper so a stray ``print(t)`` doesn't silently
481 hide the lazy nature of the value.
482 """
483 return f"AsyncCollectiveTensor({self._wait_and_unwrap()})"
485 def __str__(self):
486 """``str(t)`` / format printing; falls through to :meth:`__repr__`."""
487 return self.__repr__()
489 def __iter__(self):
490 """Iterate over dim-0 slices; one wait, then iterate inner."""
491 return iter(self._wait_and_unwrap())
494class _MSAsyncA2ALazyBwd(_Function):
495 """Async all-to-all whose forward and backward both return
496 :class:`AsyncCollectiveTensor`, deferring ``CommHandle.wait()``
497 to the first consumer op via ``__ms_dispatch__``.
499 Mirrors the Torch ``_AsyncA2ALazyBwd`` semantics: the kernel is
500 queued on the HCCL group's stream, host returns immediately, and
501 the wait fires lazily on the consumer's stream — giving the
502 paired thread a window to dispatch its compute concurrently.
503 """
505 @staticmethod
506 def _issue_async_a2a(flat_input, send_splits, recv_splits, group):
507 """Issue an async all-to-all-v on a 1-D flat tensor.
509 Bypasses ``comm_func.all_to_all_single``: that wrapper appends an
510 unconditional ``result.reshape((-1,) + recv_shape_without_first_dim)``
511 on the default compute stream *before* the async ``CommHandle.wait()``
512 fires (the wait is deferred to the first consumer op via
513 :class:`AsyncCollectiveTensor`). MindSpore's mem_pool race_checker
514 (``MS_ALLOC_CONF=memory_tracker:True``) flags that trailing reshape
515 as a cross-stream race on the HCCL output, even though for 1-D
516 inputs it is a metadata-only no-op. Calling the inner primitive
517 directly skips the tracker-visible read on stream 0.
519 Args:
520 flat_input: 1-D tensor — must already be flattened by the caller.
521 send_splits: ``list[int]`` — element counts sent to each rank.
522 recv_splits: ``list[int]`` — element counts received from each rank.
523 group: Process group.
525 Returns:
526 ``(output_tensor, CommHandle)`` — the 1-D output and the async handle.
527 """
528 rank_size = get_group_size(group)
529 # Positional args follow the MS auto-generated primitive signature:
530 # ``(input, group, send_splits, recv_splits, rank_size, block)``.
531 # ``block=False`` selects the async path; the handle is returned in
532 # the raw tuple and unpacked by ``_deal_comm_outputs`` below.
533 raw = inner_comm_all_to_all_v_op(
534 flat_input, group, list(send_splits), list(recv_splits), rank_size,
535 False,
536 )
537 # ``_deal_comm_outputs(raw, is_async=True)`` mirrors the async branch
538 # inside ``comm_func.all_to_all_single`` — unpacks the primitive's raw
539 # output into ``(tensor, handle)`` without the trailing reshape.
540 return _deal_comm_outputs(raw, True)
542 @staticmethod
543 def forward(ctx, input_tensor, output_splits, input_splits, group): # pylint: disable=arguments-differ
544 """Launch async a2a; return :class:`AsyncCollectiveTensor`.
546 ``input_tensor`` must already be 1-D and the splits must be element
547 counts (not row counts). The caller is expected to flatten and
548 translate splits beforehand — see
549 :meth:`MindSporePlatform.differentiable_all_to_all_single_async`.
550 """
551 ctx.input_splits = input_splits
552 ctx.output_splits = output_splits
553 ctx.group = group
554 flat_input = input_tensor.reshape(-1)
555 actual_output, work = _MSAsyncA2ALazyBwd._issue_async_a2a(
556 flat_input, input_splits, output_splits, group,
557 )
558 return AsyncCollectiveTensor(actual_output, work)
560 @staticmethod
561 def backward(ctx, grad_output): # pylint: disable=arguments-differ
562 """Symmetric reverse a2a; returns :class:`AsyncCollectiveTensor`."""
563 # If grad_output is still lazy, force unwrap before issuing the
564 # reverse a2a (which is itself a "real" op on the data).
565 if isinstance(grad_output, AsyncCollectiveTensor):
566 grad_output = grad_output._wait_and_unwrap() # pylint: disable=W0212
567 flat_grad = grad_output.reshape(-1)
568 actual_grad, work = _MSAsyncA2ALazyBwd._issue_async_a2a(
569 flat_grad, ctx.output_splits, ctx.input_splits, ctx.group,
570 )
571 lazy_grad = AsyncCollectiveTensor(actual_grad, work)
572 return lazy_grad, None, None, None
575class _MSSyncHookFunction(_Function):
576 """Identity autograd op that fires HookCoordinator rendezvous on
577 forward and backward, mirroring the Torch ``_TorchSyncHookFunction``.
579 The role tables are intentionally identical to the Torch backend so
580 the dual-thread protocol (COMM-first dispatch ordering) is the same
581 on MindSpore.
583 Hook-name semantics:
585 - ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` — full rendezvous on both
586 forward and backward, using ``_FWD_ROLES`` / ``_BWD_ROLES``.
587 - ``"CHUNK_START"`` — pair-0 entry hook.
588 **Forward**: full rendezvous(COMPUTE) — pairs with
589 ``D_LAST.bwd`` so the BWD thread's combine.bwd of the last
590 layer is bracketed by a barrier-synced window.
591 **Backward**: paired with ``CHUNK_END.fwd`` as the BWD-side of
592 the exit barrier (roles ``(COMPUTE, COMPUTE)``).
593 - ``"D_LAST"`` — closing D hook of the last MoE layer in a chunk.
594 **Forward**: **pure skip** — neither notify nor rendezvous.
595 The C_last → combine COMM event is left un-notified so BWD's
596 COMPUTE waiter at ``A_0.bwd`` stays parked. This keeps FWD's
597 post-combine forward work serialised against BWD's Attn.bwd_0;
598 required because MS PyNative does not support concurrent
599 FWD-record + BWD-replay on its autograd executor. (The Torch
600 backend takes the looser ``notify(COMM) + skip`` path here for
601 more overlap — Torch autograd is thread-safe.)
602 **Backward**: full rendezvous using ``_BWD_ROLES["D"]``; this
603 is the very first BWD rendezvous and pairs with
604 ``CHUNK_START.fwd`` to bracket combine.bwd_last.
605 - ``"CHUNK_END"`` — pair-N exit hook (FWD side).
606 **Forward**: roles ``(COMM, COMPUTE)``. ``notify_dispatched``
607 sets the C_last event (waking BWD's A_0.bwd waiter), then
608 ``rendezvous(COMPUTE)`` parks FWD on the exit barrier so BWD's
609 Attn.bwd_0 runs with FWD already blocked — no concurrent
610 FWD-record + BWD-replay.
611 **Backward**: skipped (this would be the first node visited
612 in BWD replay; its partner ``D_LAST.bwd`` already pairs with
613 ``CHUNK_START.fwd`` on pair 0).
614 """
616 # Index encoding: 1 = COMM, 2 = COMPUTE.
617 _FWD_ROLES = {
618 # ``CHUNK_START``: chunk entry on FWD. No "previous" op on
619 # this thread within this overlap.run() — ``notify(COMPUTE)``
620 # is a no-op anyway. Next role is COMPUTE so FWD parks on
621 # ``_comm_dispatched.wait`` for BWD's ``D_LAST.bwd`` COMM.
622 "CHUNK_START": (2, 2),
623 "A": (2, 1), # prev=Attention COMPUTE | next=dispatch COMM
624 "B": (1, 2), # prev=dispatch COMM | next=module COMPUTE
625 "C": (2, 1), # prev=module COMPUTE | next=combine COMM
626 "D": (1, 2), # prev=combine COMM | next=Attention COMPUTE
627 # ``CHUNK_END``: chunk-exit hook on FWD. Does two things in
628 # one place — both critical for MS PyNative correctness:
629 # 1. ``notify_dispatched(COMM)`` sets the C_last event from
630 # C_last's rendezvous(COMM). ``D_LAST.fwd`` deliberately
631 # does NOT notify (it is a pure skip) so BWD's COMPUTE
632 # waiter at ``A_0.bwd`` stays parked until FWD has
633 # finished all chunk-local forward work (post-combine
634 # sort/index_select/multiply).
635 # 2. ``rendezvous(COMPUTE)`` parks FWD on the exit barrier.
636 # By the time BWD wakes from step 1 and starts
637 # Attn.bwd_0, FWD is already blocked at this barrier —
638 # no concurrent FWD-record + BWD-replay window.
639 "CHUNK_END": (1, 2),
640 }
641 _BWD_ROLES = {
642 # ``CHUNK_START.bwd`` is intentionally NOT engaged here.
643 # MS PyNative's autograd may skip the backward node if the
644 # chunk input lacks ``requires_grad`` (the value of
645 # ``x.grad`` is unused downstream), which would leave the
646 # pair-8 BWD partner unmatched and deadlock FWD's
647 # ``CHUNK_END`` barrier. pair-8 BWD is instead taken out of
648 # band: the OVERLAP_B_F callback's ``bwd_fn`` makes one
649 # explicit ``coordinator.rendezvous(COMPUTE)`` after
650 # ``backward_one_chunk`` returns, paired with FWD's
651 # ``CHUNK_END.fwd`` rendezvous.
652 # ``D_LAST`` on backward routes through D's BWD role (COMM
653 # next: the upcoming combine.bwd) — see the docstring above
654 # for why we no longer skip.
655 "D": (2, 1), # prev=Attn.bwd COMPUTE | next=combine.bwd COMM
656 "C": (1, 2), # prev=combine.bwd COMM | next=module.bwd COMPUTE
657 "B": (2, 1), # prev=module.bwd COMPUTE | next=dispatch.bwd COMM
658 "A": (1, 2), # prev=dispatch.bwd COMM | next=Attn.bwd COMPUTE
659 }
660 _ROLE_CACHE = None
662 @staticmethod
663 def _role_enum(idx: int):
664 """Lazy import of HookRole to avoid a circular import at module load."""
665 if _MSSyncHookFunction._ROLE_CACHE is None:
666 # pylint: disable=C0415
667 from hyper_parallel.core.pipeline_parallel.hook_coordinator import HookRole
668 _MSSyncHookFunction._ROLE_CACHE = (None, HookRole.COMM, HookRole.COMPUTE)
669 return _MSSyncHookFunction._ROLE_CACHE[idx]
671 @staticmethod
672 def _passthrough(x):
673 """Identity passthrough that defeats MS autograd's identity-output handling.
675 When :meth:`forward` returns its input unchanged, MS PyNative's
676 ``FunctionBase.apply`` sees ``is_same_as_input=True`` on the output
677 and inserts a ``ViewAsSelfWithNoGrad`` (a ``view(self, self.shape)``
678 kernel) on the current compute stream. If the input is an
679 :class:`AsyncCollectiveTensor` whose lazy ``CommHandle.wait()`` has
680 not yet fired, that view runs on the default stream while the HCCL
681 kernel is still writing the same memory on the comm stream — flagged
682 by MS's mem_pool ``race_checker`` (``MS_ALLOC_CONF=memory_tracker:True``).
684 Returning a freshly wrapped :class:`AsyncCollectiveTensor` keeps the
685 same underlying buffer and pending work, but yields a new
686 ``shared_ptr<Tensor>`` so ``is_same_as_input`` is ``False`` and no
687 autograd view is emitted. For regular tensors the original
688 passthrough is safe (the view sits on the same stream as the data).
690 Note:
691 The clone shares ``_pending_work`` with the original but keeps
692 an independent ``completed`` flag. Two assumptions:
694 * ``CommHandle.wait()`` is idempotent — relied on whenever both
695 wrappers end up being consumed (matches the existing
696 :meth:`AsyncCollectiveTensor._wait_and_unwrap` pattern, which
697 also does not null out ``_pending_work`` after waiting).
698 * Per-wrapper ``completed`` is intentional: a ``wait()`` on
699 stream A does not synchronize stream B, so each consumer
700 stream must be free to re-issue its own wait.
701 """
702 if isinstance(x, AsyncCollectiveTensor):
703 new_wrapper = AsyncCollectiveTensor(x.elem, x._pending_work) # pylint: disable=W0212
704 new_wrapper.completed = x.completed
705 return new_wrapper
706 return x
708 @staticmethod
709 def forward(ctx, x, hook_name, coordinator): # pylint: disable=arguments-differ
710 """Fire forward-direction rendezvous and return ``x`` unchanged."""
711 ctx.hook_name = hook_name
712 ctx.coordinator = coordinator
713 if not coordinator.is_enabled():
714 return _MSSyncHookFunction._passthrough(x)
715 if hook_name == "D_LAST":
716 # Pure skip — neither notify nor rendezvous. The
717 # C_last → combine COMM event is left un-notified on
718 # purpose so BWD's COMPUTE waiter at A_0.bwd stays parked
719 # until FWD reaches CHUNK_END.fwd. This keeps FWD's
720 # post-combine forward work (sort / index_select / probs
721 # mul / strided_slice) strictly serialised against BWD's
722 # Attn.bwd_0 — required because MS PyNative does not
723 # support concurrent FWD-record + BWD-replay on the
724 # autograd executor.
725 return _MSSyncHookFunction._passthrough(x)
726 prev_idx, next_idx = _MSSyncHookFunction._FWD_ROLES[hook_name]
727 role_of = _MSSyncHookFunction._role_enum
728 coordinator.notify_dispatched(role_of(prev_idx))
729 coordinator.rendezvous(role_of(next_idx))
730 return _MSSyncHookFunction._passthrough(x)
732 @staticmethod
733 def backward(ctx, grad_output): # pylint: disable=arguments-differ
734 """Mirror of :meth:`forward` using ``_BWD_ROLES``."""
735 hook_name = ctx.hook_name
736 coordinator = ctx.coordinator
737 if not coordinator.is_enabled():
738 return _MSSyncHookFunction._passthrough(grad_output), None, None
739 if hook_name in ("CHUNK_END", "CHUNK_START"):
740 # Both boundary hooks skip in backward:
741 # * ``CHUNK_END.bwd`` would fire FIRST in BWD replay (it
742 # wraps the chunk's last forward op). We do not want
743 # a rendezvous here — pair 0 is handled by
744 # ``D_LAST.bwd`` ↔ ``CHUNK_START.fwd``.
745 # * ``CHUNK_START.bwd`` would fire LAST. We do not
746 # rendezvous here either, because MS autograd may skip
747 # the node entirely when the chunk input lacks
748 # ``requires_grad`` (unused ``x.grad``). pair-8 BWD
749 # is taken out of band — see the role-table comment.
750 return _MSSyncHookFunction._passthrough(grad_output), None, None
751 # ``D_LAST.bwd`` reuses D's BWD role: it is the *first non-skip*
752 # BWD rendezvous and pairs with FWD's ``CHUNK_START`` to lock
753 # the combine.bwd_last launch inside a barrier-synced window.
754 role_name = "D" if hook_name == "D_LAST" else hook_name
755 prev_idx, next_idx = _MSSyncHookFunction._BWD_ROLES[role_name]
756 role_of = _MSSyncHookFunction._role_enum
757 coordinator.notify_dispatched(role_of(prev_idx))
758 coordinator.rendezvous(role_of(next_idx))
759 return _MSSyncHookFunction._passthrough(grad_output), None, None
762class _MSAsyncA2AFunction(_Function):
763 """Differentiable wrapper for pre-launched async all-to-all."""
765 @staticmethod
766 def forward(ctx, x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box): # pylint: disable=arguments-differ
767 """Wait for pre-launched async A2A and return reconstructed output."""
768 ctx.group = group
769 ctx.world_size = world_size
770 ctx.concat_dim = concat_dim
771 ctx.split_dim = split_dim
772 ctx.handle_box = handle_box
773 ctx.x_shape = tuple(x.shape)
774 work.wait()
775 return _a2a_reconstruct_ms(out_perm, concat_dim)
777 @staticmethod
778 def backward(ctx, grad_output): # pylint: disable=arguments-differ
779 """Launch async head->seq A2A for backward overlap, or return zero grad."""
780 if ctx.handle_box is not None:
781 g = grad_output.contiguous()
782 shape = list(g.shape)
783 seq_dim = ctx.concat_dim
784 s_full = shape[seq_dim]
785 ndim = len(shape) + 1
786 x_perm = g.reshape(
787 shape[:seq_dim] + [ctx.world_size, s_full // ctx.world_size] + shape[seq_dim + 1:]
788 ).permute(
789 [seq_dim] + list(range(seq_dim)) + list(range(seq_dim + 1, ndim))
790 ).contiguous()
791 out_perm, work = _mindspore_all_to_all_single(
792 x_perm,
793 list(x_perm.shape),
794 ctx.group,
795 async_op=True,
796 )
797 ctx.handle_box.append((work, out_perm))
798 return mint.zeros(ctx.x_shape, dtype=grad_output.dtype), None, None, None, None, None, None, None
801class _MSAsyncAllGatherFunction(_Function):
802 """Differentiable wrapper for pre-launched async all-gather."""
804 @staticmethod
805 def forward(ctx, x, work, out_perm, group, world_size, gather_dim, handle_box): # pylint: disable=arguments-differ
806 """Wait for pre-launched all-gather and reconstruct the gathered tensor."""
807 ctx.group = group
808 ctx.world_size = world_size
809 ctx.gather_dim = gather_dim
810 ctx.handle_box = handle_box
811 ctx.x_shape = tuple(x.shape)
812 work.wait()
813 return _move_dim_from_front(out_perm, gather_dim)
815 @staticmethod
816 def backward(ctx, grad_output): # pylint: disable=arguments-differ
817 """Launch reverse reduce-scatter for the all-gather."""
818 grad_perm = _move_dim_to_front(grad_output.contiguous(), ctx.gather_dim)
819 output_shape = list(grad_perm.shape)
820 if output_shape[0] % ctx.world_size != 0:
821 raise ValueError(
822 "all_gather backward expected gathered dimension to be divisible by world_size, "
823 f"got {output_shape[0]} and {ctx.world_size}."
824 )
825 output_shape[0] //= ctx.world_size
826 output, work = _mindspore_reduce_scatter_single(
827 grad_perm,
828 output_shape,
829 ctx.group,
830 async_op=True,
831 )
832 if ctx.handle_box is not None:
833 ctx.handle_box.append((work, output, ctx.gather_dim))
834 return mint.zeros(ctx.x_shape, dtype=grad_output.dtype), None, None, None, None, None, None
835 work.wait()
836 return _move_dim_from_front(output, ctx.gather_dim), None, None, None, None, None, None
839def _ensure_contiguous(x):
840 """Return a contiguous copy of *x* if not already contiguous."""
841 if not x.is_contiguous() or x.storage_offset() != 0:
842 x = x.contiguous()
843 return x
846class MindSporePlatform(Platform):
847 """MindSpore platform api"""
848 Tensor = Tensor
849 tensor = Tensor
850 Parameter = Parameter
851 Module = Cell
852 DTensorBase = DTensorBase
853 PipelineStageBase = PipelineStageBase
854 platform_type = PlatformType.MINDSPORE
855 tensor_dtype = mstype
856 dtype = ms.Type
857 Function = _Function
859 _custom_ops_cls = None
861 @property
862 def custom_ops(self):
863 """Return the MindSpore platform custom ops instance.
865 .. warning::
866 This is an experimental API that subject to change or deletion.
868 Returns:
869 MindSporeCustomOps: Custom ops class that delegates to DFunction
870 implementations wrapping Ascend NPU custom C++ kernels.
871 """
872 if self._custom_ops_cls is None:
873 from hyper_parallel.platform.mindspore.custom_ops.custom_ops import ( # pylint: disable=import-outside-toplevel
874 MindSporeCustomOps,
875 )
876 self._custom_ops_cls = MindSporeCustomOps
877 return self._custom_ops_cls
879 @staticmethod
880 def get_swap_optimizer():
881 """Return the MindSpore optimizer-state swap wrapper class."""
882 from hyper_parallel.platform.mindspore.swap_optimizer.swap_optimizer import ( # pylint: disable=import-outside-toplevel
883 get_swap_optimizer,
884 )
885 return get_swap_optimizer()
887 def __init__(self):
888 # Ensure MindSpore ``nn.Cell.to_empty`` is patched as soon as the
889 # MindSpore platform instance is created.
890 _install_cell_to_empty_patch()
892 @staticmethod
893 def is_linear_module(module) -> bool:
894 """Check whether *module* is a MindSpore ``Dense`` (linear) or ``mint.nn.Linear`` layer."""
895 return isinstance(module, (ms.nn.Dense, mint.nn.Linear))
897 @staticmethod
898 def is_embedding_module(module) -> bool:
899 """Check whether *module* is a MindSpore ``Embedding`` or ``mint.nn.Embedding`` layer."""
900 return isinstance(module, (ms.nn.Embedding, mint.nn.Embedding))
902 def device_count(self, device_handle):
903 """
904 Get the number of available devices.
906 Args:
907 device_handle: The device handle (e.g., ms.device_context).
909 Returns:
910 int: The number of available devices.
911 """
912 device_type = self.device_type()
913 if device_type == "cpu":
914 return device_handle.device_context.cpu.device_count()
915 if device_type == "gpu":
916 return device_handle.device_context.gpu.device_count()
917 return device_handle.device_context.ascend.device_count()
919 @staticmethod
920 def get_rng_state(device=None, device_handle=None):
921 """
922 Get the random number generator state.
924 Args:
925 device (Optional): The device to get RNG state from (not used in MindSpore).
926 device_handle (Optional): The device handle (not used in MindSpore).
928 Returns:
929 Tensor: The RNG state as a tensor.
930 """
931 _ = device, device_handle
932 return ms.get_rng_state()
934 @staticmethod
935 def set_rng_state(state, device=None, device_handle=None):
936 """
937 Set the random number generator state.
939 Args:
940 state (Tensor): The RNG state to set.
941 device (Optional): The device to set RNG state for (not used in MindSpore).
942 device_handle (Optional): The device handle (not used in MindSpore).
943 """
944 _ = device, device_handle
945 return ms.set_rng_state(state)
947 def device_type(self):
948 """
949 Get the current device type.
951 Returns:
952 str: The device type string ("npu" for Ascend, "gpu" for GPU, "cpu" for CPU).
953 """
954 device_type = ms.get_context("device_target")
955 if device_type == "Ascend":
956 return "npu"
957 return device_type.lower()
959 def device(self, device_idx=None):
960 """
961 Get the device type string.
963 Args:
964 device_idx (Optional[int]): The device index (not used in MindSpore).
966 Returns:
967 str: The device type string.
968 """
969 _ = device_idx
970 device_type = self.device_type()
971 return device_type
973 @staticmethod
974 def get_device_handle():
975 """
976 Get the MindSpore module as the device handle.
978 Returns:
979 module: The mindspore module.
980 """
981 return ms
983 @staticmethod
984 def manual_seed(seed):
985 """
986 Set the random seed for reproducibility.
988 Args:
989 seed (int): The random seed value.
991 Returns:
992 None
993 """
994 return ms.manual_seed(seed)
996 @staticmethod
997 def ones(size, dtype=None):
998 """
999 Create a tensor filled with ones.
1001 Args:
1002 size (tuple): The shape of the output tensor.
1003 dtype (Optional[ms.Type]): The desired data type.
1005 Returns:
1006 Tensor: A tensor filled with ones.
1007 """
1008 return mint.ones(size, dtype=dtype)
1010 @staticmethod
1011 def zeros(size, dtype=None, device=None):
1012 """
1013 Create a tensor filled with zeros.
1015 Args:
1016 size (tuple): The shape of the output tensor.
1017 dtype (Optional[ms.Type]): The desired data type.
1018 device (Optional[ms.device]): The device to create the tensor on.
1020 Returns:
1021 Tensor: A tensor filled with zeros.
1022 """
1023 tensor = mint.zeros(size, dtype=dtype)
1024 if device in ("GPU", "Ascend"):
1025 return tensor.to(device)
1026 return tensor
1028 @staticmethod
1029 def full(size, fill_value, dtype=None):
1030 """
1031 Create a tensor filled with a scalar value.
1033 Args:
1034 size (tuple): The shape of the output tensor.
1035 fill_value (scalar): The value to fill the tensor with.
1036 dtype (Optional[ms.Type]): The desired data type.
1038 Returns:
1039 Tensor: A tensor filled with the specified value.
1040 """
1041 return mint.full(size, fill_value, dtype=dtype)
1043 @staticmethod
1044 def empty(size, dtype=None, device=None): # pylint: disable=unused-argument
1045 """
1046 Create an uninitialized tensor.
1048 Args:
1049 size (tuple): The shape of the output tensor.
1050 dtype (Optional[ms.Type]): The desired data type.
1051 device: Accepted for cross-backend signature parity with the
1052 Torch backend but ignored — under MindSpore the active
1053 device is bound at process init via ``ms.set_device`` and
1054 ``mint.empty`` allocates on it directly.
1056 Returns:
1057 Tensor: An uninitialized tensor.
1058 """
1059 return mint.empty(size, dtype=dtype)
1061 @staticmethod
1062 def rand(size, dtype=None, device=None): # pylint: disable=unused-argument
1063 """Create a tensor filled with uniform random values in ``[0, 1)``."""
1064 tensor = mint.rand(size, dtype=dtype)
1065 if device in ("GPU", "Ascend"):
1066 return tensor.to(device)
1067 return tensor
1069 @staticmethod
1070 def randn(size, dtype=None, device=None): # pylint: disable=unused-argument
1071 """Create a tensor filled with standard-normal random values."""
1072 tensor = mint.randn(size, dtype=dtype)
1073 if device in ("GPU", "Ascend"):
1074 return tensor.to(device)
1075 return tensor
1077 @staticmethod
1078 def get_rank():
1079 """
1080 Get the rank of the current process in the distributed group.
1082 Returns:
1083 int: The rank of the current process.
1084 """
1085 return get_rank_id()
1087 @staticmethod
1088 def get_global_rank(group, group_rank):
1089 """
1090 Get the global rank from a group rank.
1092 Args:
1093 group (str): The process group name.
1094 group_rank (int): The rank within the group.
1096 Returns:
1097 int: The global rank.
1098 """
1099 return dist.get_global_rank(group, group_rank)
1101 @staticmethod
1102 def get_world_size():
1103 """
1104 Get the total number of processes in the distributed group.
1106 Returns:
1107 int: The world size.
1108 """
1109 return get_group_size()
1111 @staticmethod
1112 def get_op_name(func):
1113 """
1114 Extract the operation name from a function.
1116 Args:
1117 func: The function to extract the name from.
1119 Returns:
1120 str: The operation name.
1121 """
1122 return func.name
1124 @staticmethod
1125 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
1126 data = _ensure_contiguous(data)
1127 # rank_list is accepted for torch parity; MindSpore keeps the existing group order.
1128 output, _ = comm_func.all_gather_into_tensor(None, data, group=group)
1129 if concat_dim == 0:
1130 return output
1131 output_tensors = ms.ops.Split(output_num=concat_size)(output)
1132 return ms.mint.concat(output_tensors, concat_dim)
1134 @staticmethod
1135 def chunk(data, split_dim, split_size, index):
1136 return ms.ops.Split(axis=split_dim, output_num=split_size)(data)[index]
1138 @staticmethod
1139 def differentiable_all_to_all(input_data, output_shape, group):
1140 input_data = _ensure_contiguous(input_data)
1141 output_tensor, _ = comm_func.all_to_all_single(
1142 output_shape,
1143 input_data,
1144 group=group,
1145 async_op=False
1146 )
1147 return output_tensor
1149 @staticmethod
1150 def tensor_type_cast(input_data, cast_type):
1151 """Cast tensor to specified data type."""
1152 type_mapping = {
1153 'float32': ms.float32,
1154 'float16': ms.float16,
1155 'int64': ms.int64,
1156 'int32': ms.int32
1157 }
1158 if cast_type not in type_mapping:
1159 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")
1160 return input_data.to(type_mapping[cast_type])
1162 @staticmethod
1163 def differentiable_all_reduce(data, op, group):
1164 data = _ensure_contiguous(data)
1165 output, _ = comm_func.all_reduce(data, op, group)
1166 return output
1168 @staticmethod
1169 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
1170 data = _ensure_contiguous(data)
1171 if axis > 0:
1172 data = ms.mint.concat(ms.ops.Split(axis=axis, output_num=dev_num)(data), dim=0)
1173 output_tensor, _ = comm_func.reduce_scatter_tensor(None, data, 'sum', group)
1174 if op == 'avg':
1175 output_tensor = output_tensor / dev_num
1176 return output_tensor
1178 @staticmethod
1179 def init_parameters(module, stage_index):
1180 return _init_parameters(module, stage_index)
1182 # pylint: disable=W0212
1183 @staticmethod
1184 def update_param_data(param, data):
1185 """update param data"""
1186 if isinstance(param, DTensorBase):
1187 param.set_data(data)
1188 else:
1189 param._update_data(data)
1191 @staticmethod
1192 def load_into_param(param, data):
1193 copy_tensor = MindSporePlatform.empty_like(data)
1194 copy_tensor.copy_(data)
1195 if isinstance(param, DTensorBase):
1196 param.set_data(copy_tensor)
1197 else:
1198 param._update(copy_tensor)
1200 @staticmethod
1201 def get_cell_construct(cell):
1202 return cell.construct
1204 @staticmethod
1205 def get_cells_and_names(cell):
1206 return cell.cells_and_names()
1208 @staticmethod
1209 def get_modules(module):
1210 return module.cells()
1212 @staticmethod
1213 def search_parameter_by_name(cell, param_name: str):
1214 """
1215 Find the parent Module of the parameter, the parameter's name in the parent Module, and the parameter.
1216 Return value: (parent Module instance, parameter's name in parent Module, parameter object).
1217 Returns None if not found.
1218 """
1219 # Remove the "self." prefix from param_name (to maintain compatibility with original logic)
1220 param_name = param_name.replace("self.", "")
1221 # Case 1: The parameter is a direct parameter of the current Module (not in any sub-Module)
1222 if param_name in cell._params:
1223 return (cell, param_name, cell._params[param_name])
1225 # Case 2: The parameter is in a sub-Module (supports multi-level nesting, e.g., "net_b.dense1.weight")
1226 if "." in param_name:
1227 # Split into: sub-Module path + parameter name (e.g., "net_b.dense1" + "weight")
1228 cell_path, param_key = param_name.rsplit(".", 1)
1229 try:
1230 # Locate the sub-Module where the parameter resides (supports multi-level paths)
1231 target_cell = cell.get_sub_cell(cell_path)
1232 # Check if the sub-Module directly contains this parameter
1233 if param_key in target_cell._params:
1234 return target_cell, param_key, target_cell._params[param_key]
1235 except AttributeError:
1236 # Sub-Module path does not exist or the parameter is not in that sub-Module
1237 pass
1239 # Traverse all sub-Modules (recursively) to search for the parameter
1240 for _, child_cell in cell._cells.items():
1241 if isinstance(child_cell, Cell):
1242 # Recursively search within the sub-Module
1243 result = MindSporePlatform.search_parameter_by_name(child_cell, param_name)
1244 if result is not None:
1245 return result
1247 return None
1249 @staticmethod
1250 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
1251 """
1252 Modify the original parameter in a Module or sub-Module using the search result
1253 Args:
1254 cell: The cell which parameter is to update
1255 result: A tuple contains parent Module, parameter key and old parameter.
1256 new_param: New Parameter object (used to replace the original parameter)
1257 """
1258 parent_cell, param_key, _ = result
1259 # Key operation: directly modify the _params dictionary of the parent Module (original storage location)
1260 parent_cell._params[param_key] = new_param
1262 if param_key in parent_cell.__dict__:
1263 parent_cell.__dict__[param_key] = new_param
1264 parent_cell._params_list[param_key] = new_param
1265 return True
1267 @staticmethod
1268 def set_layout_into_parameter(param, layout):
1269 """Set layout in to parameter"""
1270 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel
1271 from hyper_parallel.core.dtensor.layout import _infer_slice_shape_by_layout, \
1272 _get_slice_tensor_by_layout # pylint: disable=import-outside-toplevel
1273 if isinstance(param, DTensor):
1274 raise ValueError(f"Parameter {param.name} has been configured layout, cannot be set repeatedly.")
1275 param_info = param.param_info
1276 requires_grad = param.requires_grad
1277 name = param.name
1278 slice_shape = _infer_slice_shape_by_layout(param.shape, layout)
1280 if not param.has_init:
1281 # has been init, get slice data
1282 param_dtensor = DTensor.from_local(
1283 _get_slice_tensor_by_layout(param, layout).value(), layout.mesh, layout.alias_placements
1284 )
1285 param = Parameter(param_dtensor, name=name, requires_grad=requires_grad)
1286 param.param_info = param_info
1287 else:
1288 # has not been init, need to modify init shape
1289 param.init_mode.shape = slice_shape
1290 param_dtensor = DTensor.from_local(param.init_mode, layout.mesh, layout.alias_placements)
1291 param = Parameter(param_dtensor, name=name, requires_grad=requires_grad)
1292 param.param_info = param_info
1293 return param
1295 @staticmethod
1296 def get_param_local_shape(param):
1297 """get param local shape"""
1298 if isinstance(param, DTensorBase):
1299 return param.local_shape
1300 return param.shape
1302 @staticmethod
1303 def get_param_local_data(param):
1304 """get param local shape"""
1305 if isinstance(param, DTensorBase):
1306 return param.to_local()
1307 return param
1309 @staticmethod
1310 def get_param_type_size(param):
1311 return type_size_in_bytes(param.dtype)
1313 @staticmethod
1314 def is_tensor(obj: Any) -> bool:
1315 """Return True if ``obj`` is a ``mindspore.Tensor``."""
1316 return isinstance(obj, Tensor)
1318 @staticmethod
1319 def get_tensor_storage_size(tensor: Any) -> int:
1320 """Return serialized byte size (numel * itemsize) for a MindSpore tensor."""
1321 if not MindSporePlatform.is_tensor(tensor):
1322 raise TypeError(
1323 f"MindSporePlatform.get_tensor_storage_size expects mindspore.Tensor, got {type(tensor)!r}"
1324 )
1325 return int(tensor.numel()) * int(tensor.itemsize)
1327 @staticmethod
1328 def new_zero_parameter(param_shape, param_type, requires_grad, device):
1329 param = Parameter(initializer("zeros", param_shape, param_type), requires_grad=requires_grad)
1330 if device in ("GPU", "Ascend"):
1331 return param.to(device)
1332 return param
1334 @staticmethod
1335 def new_tensor(tensor_shape, tensor_type, device):
1336 tensor = Tensor(shape=tensor_shape, dtype=tensor_type)
1337 if device in ("GPU", "Ascend"):
1338 return tensor.to(device)
1339 return tensor
1341 @staticmethod
1342 def full_like(tensor, fill_value, dtype=None):
1343 return mint.full_like(tensor, fill_value, dtype=dtype)
1345 @staticmethod
1346 def isend(tensor, dst=None, group=None, tag=0):
1347 return dist.isend(tensor, dst, group, tag)
1349 @staticmethod
1350 def irecv(tensor, src=None, group=None, tag=0):
1351 return dist.irecv(tensor, src, group, tag)
1353 @staticmethod
1354 def p2p_op(op_type, tensor, peer, group=None):
1355 # pylint: disable=C0415
1356 from mindspore.mint.distributed import P2POp
1357 return P2POp(op_type, tensor, peer, group)
1359 @staticmethod
1360 def batch_isend_irecv(p2p_ops):
1361 """Launch a peer-batched P2P group.
1363 MindSpore's ``batch_isend_irecv`` lowers the whole list to a single
1364 ``HcclBatchISendIRecv`` kernel on one comm stream and returns a list
1365 with one packaging ``CommHandle``; we hand that single handle back so
1366 callers can defer the whole batch's wait to one consumption point.
1367 A send and a recv to the same peer therefore overlap on the duplex
1368 link inside this one kernel.
1369 """
1370 # pylint: disable=C0415
1371 from mindspore.mint.distributed import batch_isend_irecv
1372 if not p2p_ops:
1373 return None
1374 handles = batch_isend_irecv(p2p_ops)
1375 return handles[0] if handles else None
1377 @staticmethod
1378 def prepare_batch_p2p_group(group: Any = None) -> None: # pylint: disable=unused-argument
1379 """Prepare a group for batched P2P operations.
1381 MindSpore does not require full-group participation before its first
1382 subset ``batch_isend_irecv`` call, so no synchronization is needed.
1384 Args:
1385 group: The communication group used by the batched P2P operations.
1386 ``None`` uses the default group.
1387 """
1389 @staticmethod
1390 def p2p_exchange(tensor, peer_rank: int, group=None): # pylint: disable=unused-argument
1391 raise NotImplementedError(
1392 "p2p_exchange is not yet supported on the MindSpore platform."
1393 )
1395 @staticmethod
1396 def send_object_list(obj_list, dst=None, group=None):
1397 # pylint: disable=C0415
1398 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import send_object_list
1399 send_object_list(obj_list, dst, group)
1401 @staticmethod
1402 def recv_object_list(obj_list, src=None, group=None):
1403 # pylint: disable=C0415
1404 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import recv_object_list
1405 recv_object_list(obj_list, src, group)
1407 @staticmethod
1408 def set_tensor_requires_grad(input_tensor):
1409 """
1410 set requires grad flag for input tensor
1411 """
1412 input_tensor.requires_grad_()
1414 @staticmethod
1415 def _normalize_group_options(pg_options: Any) -> Any:
1416 if not isinstance(pg_options, dict) or "hccl_config" not in pg_options:
1417 return pg_options
1418 from mindspore._c_expression import GroupOptions # pylint: disable=C0415
1420 options = GroupOptions()
1421 options.hccl_config = pg_options["hccl_config"]
1422 return options
1424 @staticmethod
1425 def _create_group_with_options(group_name: str, rank_list: list[int], pg_options: Any = None) -> None:
1426 """Create a MindSpore communication group with optional backend-specific options."""
1427 if pg_options is None:
1428 new_group(rank_ids=rank_list, group=group_name)
1429 return
1430 try:
1431 new_group(
1432 rank_ids=rank_list,
1433 group=group_name,
1434 options=MindSporePlatform._normalize_group_options(pg_options),
1435 )
1436 except (ImportError, RuntimeError, TypeError, ValueError):
1437 new_group(rank_ids=rank_list, group=group_name)
1439 def _create_group(self, rank_list, pg_options: Any = None):
1440 world_group = self._maybe_reuse_world_group(rank_list)
1441 if world_group is not None:
1442 return world_group
1444 group_name = str(tuple(sorted(rank_list)))
1445 self._create_group_with_options(group_name, rank_list, pg_options=pg_options)
1446 EXISTING_COMM_GROUPS[group_name] = group_name
1447 return group_name
1449 @staticmethod
1450 def all_gather_into_tensor(data, group_info, async_op=False):
1451 group_name = group_info if isinstance(group_info, str) else group_info.group_name
1452 rank_size = get_group_size(group_name) if isinstance(group_info, str) else group_info.rank_size
1453 output_shape = list(data.shape)
1454 output_shape[0] *= rank_size
1455 return _mindspore_all_gather_single(data, output_shape, group_name, async_op=async_op)
1457 @staticmethod
1458 def all_gather_single(input_tensor, output_shape, group, async_op=False):
1459 return _mindspore_all_gather_single(input_tensor, output_shape, group, async_op=async_op)
1461 @staticmethod
1462 def all_reduce(data, group_info, async_op=False):
1463 if isinstance(group_info, str):
1464 handle = dist.all_reduce(data, group=group_info, async_op=async_op)
1465 else:
1466 handle = dist.all_reduce(data, group=group_info.group_name, async_op=async_op)
1467 return data, handle
1469 @staticmethod
1470 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
1471 if group_src is not None:
1472 ranks = MindSporePlatform.get_process_group_ranks(group)
1473 src = ranks[group_src]
1474 handle = dist.broadcast(data, src, group, async_op)
1475 if async_op:
1476 handle.wait()
1477 return data
1479 @staticmethod
1480 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
1481 group_name = group if isinstance(group, str) else getattr(group, "group_name", group)
1482 if group_src is not None:
1483 ranks = MindSporePlatform.get_process_group_ranks(group)
1484 src = ranks[group_src]
1485 if scatter_list is None:
1486 # MindSpore mint.scatter validates scatter_list on every rank; PyTorch passes None on receivers.
1487 rank_size = get_group_size(group_name)
1488 scatter_list = [output] * rank_size
1489 else:
1490 scatter_list = [c.contiguous() if hasattr(c, "is_contiguous") and not c.is_contiguous() else c
1491 for c in scatter_list]
1492 handle = dist.scatter(output, scatter_list, src, group_name, async_op=async_op)
1493 if async_op and handle is not None:
1494 handle.wait()
1495 return output
1497 @staticmethod
1498 def reduce_scatter_tensor(data, group_info, async_op=False):
1499 group_name = group_info if isinstance(group_info, str) else group_info.group_name
1500 rank_size = get_group_size(group_name) if isinstance(group_info, str) else group_info.rank_size
1501 output_shape = list(data.shape)
1502 output_shape[0] //= rank_size
1503 return _mindspore_reduce_scatter_single(data, output_shape, group_name, async_op=async_op)
1505 @staticmethod
1506 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
1507 return _mindspore_reduce_scatter_single(input_tensor, output_shape, group, async_op=async_op)
1509 @staticmethod
1510 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
1511 return _mindspore_all_to_all_single(input_tensor, output_shape, group, async_op=async_op)
1513 @staticmethod
1514 def differentiable_all_to_all_single(
1515 input_tensor: Tensor,
1516 input_splits: Sequence[int],
1517 output_splits: Sequence[int],
1518 group: str,
1519 ) -> Tensor:
1520 """Run a differentiable N-D variable all-to-all with dim-zero row splits."""
1521 input_splits, output_splits = _validate_variable_row_splits(
1522 input_tensor,
1523 input_splits,
1524 output_splits,
1525 group,
1526 )
1527 return _MSDifferentiableAllToAllSingle.apply(
1528 input_tensor,
1529 output_splits,
1530 input_splits,
1531 group,
1532 )
1534 @staticmethod
1535 def differentiable_variable_all_gather(
1536 input_tensor: Tensor, output_splits: Sequence[int], group: str) -> Tensor:
1537 """Gather variable dim-zero row shards with native ``AllGatherV``."""
1538 return _mindspore_variable_all_gather(input_tensor, output_splits, group)
1540 @staticmethod
1541 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
1542 handle_box=None):
1543 return _MSAsyncAllGatherFunction.apply(
1544 x, work, out_perm, group, world_size, gather_dim, handle_box
1545 )
1547 @staticmethod
1548 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim, # pylint: disable=unused-argument
1549 handle_box=None):
1550 return _MSAsyncA2AFunction.apply(
1551 x, work, out_perm, group, world_size, concat_dim, split_dim, handle_box
1552 )
1554 @staticmethod
1555 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
1556 """Launch an asynchronous, differentiable all-to-all-single.
1558 Token a2a entry point used by ``CommComputeOverlap``-driven MoE
1559 wrappers. The kernel is queued on the HCCL group's stream and
1560 the host returns immediately, so the calling thread can proceed
1561 to the next sync hook (notify + rendezvous) before the
1562 collective finishes — this is what enables the comm/compute
1563 overlap window on the paired thread.
1565 Args:
1566 input_tensor: **1-D** tensor — the caller is responsible for
1567 flattening multi-dim inputs beforehand.
1568 input_splits: ``list[int]`` — **element** counts sent to each
1569 rank (not row counts). For an originally
1570 ``(N, D)`` tensor, each entry is ``rows_i * D``.
1571 output_splits: ``list[int]`` — element counts received from each rank.
1572 group: Process group.
1574 Returns:
1575 ``AsyncCollectiveTensor`` of shape ``(sum(output_splits),)`` that
1576 defers ``CommHandle.wait()`` to the first consumer op via
1577 ``__ms_dispatch__``.
1579 Raises:
1580 ValueError: if ``input_tensor`` is not 1-D.
1582 Note:
1583 The 1-D + element-count contract diverges from the Torch
1584 implementation (which accepts N-D input + row-count splits).
1585 The divergence is intentional for now: it lets the MS path
1586 call the inner primitive directly and avoid the cross-stream
1587 race that ``comm_func.all_to_all_single``'s trailing reshape
1588 triggers under ``MS_ALLOC_CONF=memory_tracker:True`` —
1589 see :meth:`_MSAsyncA2ALazyBwd._issue_async_a2a`.
1590 """
1591 if input_tensor.ndim != 1:
1592 raise ValueError(
1593 "MindSporePlatform.differentiable_all_to_all_single_async requires a 1-D "
1594 f"input_tensor (got ndim={input_tensor.ndim}, shape={tuple(input_tensor.shape)}). "
1595 "Flatten the tensor and convert row-count splits to element counts before calling."
1596 )
1597 return _MSAsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)
1599 @staticmethod
1600 def differentiable_sync_hook(x, hook_name: str, coordinator):
1601 """Fire a HookCoordinator rendezvous on forward and backward.
1603 Args:
1604 x: Input tensor — returned unchanged.
1605 hook_name: One of:
1606 * ``"A"`` / ``"B"`` / ``"C"`` / ``"D"`` —
1607 full rendezvous on both directions.
1608 * ``"CHUNK_START"`` — chunk-entry hook on
1609 forward; pairs with ``D_LAST.bwd`` so the
1610 BWD thread's combine.bwd of the last layer
1611 is bracketed by a barrier-synced sync point.
1612 Skipped on backward.
1613 * ``"D_LAST"`` — closing D of the last MoE
1614 layer in a chunk. Forward: ``notify_dispatched``
1615 only (no Attention follows so rendezvous is
1616 skipped). Backward: full rendezvous via D's
1617 BWD role; paired with ``CHUNK_START`` on FWD.
1618 coordinator: The :class:`HookCoordinator` driving the
1619 rendezvous protocol.
1621 Returns:
1622 ``x`` unchanged.
1624 Note:
1625 Two-thread compatibility on MindSpore PyNative is not yet
1626 fully verified. The HookCoordinator + ``_Function``
1627 primitives are individually thread-safe, but the
1628 interaction with MindSpore's autograd execution model
1629 under ``threading.Thread`` should be PoC-tested before
1630 production use.
1631 """
1632 return _MSSyncHookFunction.apply(x, hook_name, coordinator)
1634 @staticmethod
1635 def parameters_dict(cell: Cell):
1636 return cell.parameters_and_names()
1638 @staticmethod
1639 def buffers_dict(cell: Cell) -> Any:
1640 """Return all named buffers registered by the cell tree."""
1641 return cell.named_buffers()
1643 @staticmethod
1644 def get_tensor_transform():
1645 return _tensor_transform
1647 @staticmethod
1648 def construct_strided_slice(x, begin, end, stride):
1649 return ms.ops.strided_slice(x, begin, end, stride)
1651 @staticmethod
1652 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1653 # pylint: disable=C0415
1654 from hyper_parallel.platform.mindspore.pipeline_parallel._utils import _MicroBatch
1655 return _MicroBatch(micro_batch_num, args_batch_dim, kwargs_batch_dim)
1657 @staticmethod
1658 def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
1659 """Get the state dictionary of a model (not yet supported on MindSpore).
1661 Args:
1662 model: The model to extract state from.
1663 options: Optional configuration for state dict extraction.
1665 Returns:
1666 dict: The state dictionary containing model parameters and buffers.
1668 Raises:
1669 NotImplementedError: MindSpore support is not yet implemented.
1670 """
1671 raise NotImplementedError(
1672 "get_model_state_dict is not yet supported on MindSpore"
1673 )
1675 @staticmethod
1676 def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
1677 """Set the state dictionary of a model (not yet supported on MindSpore).
1679 Args:
1680 model: The model to load state into.
1681 model_state_dict: The state dict to load into the model.
1682 options: Optional configuration for state dict loading.
1684 Returns:
1685 None.
1687 Raises:
1688 NotImplementedError: MindSpore support is not yet implemented.
1689 """
1690 raise NotImplementedError(
1691 "set_model_state_dict is not yet supported on MindSpore"
1692 )
1694 @staticmethod
1695 def save_checkpoint(cell: Union[Cell, dict], file_path: str, ckpt_format: str = "safetensors") -> None:
1696 if isinstance(cell, dict):
1697 save_dict = {}
1698 for k, v in cell.items():
1699 if isinstance(v, Parameter):
1700 save_dict[k] = v
1701 elif isinstance(v, Tensor):
1702 save_dict[k] = Parameter(v, name=k)
1703 else:
1704 save_dict[k] = v
1705 else:
1706 save_dict = cell._params
1707 ms.save_checkpoint(save_obj=save_dict, ckpt_file_name=file_path, format=ckpt_format)
1709 @staticmethod
1710 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
1711 return ms.load_checkpoint(ckpt_file_name=file_path, format=ckpt_format)
1713 @staticmethod
1714 def get_symmetric_memory_handler():
1715 # pylint: disable=C0415
1716 from hyper_parallel.platform.mindspore.symmetric_memory import MSSymmetricMemoryHandler
1717 symmetric_memory = MSSymmetricMemoryHandler()
1718 return symmetric_memory
1720 @staticmethod
1721 def get_multicore_handler():
1722 """Create and return a MindSpore multicore handler instance."""
1723 # pylint: disable=C0415
1724 from hyper_parallel.platform.mindspore.multicore import MSMulticoreHandler
1725 return MSMulticoreHandler()
1727 def new_stream(self):
1728 return ms.runtime.Stream()
1730 def get_stream_context(self):
1731 return ms.runtime.StreamCtx
1733 @staticmethod
1734 def all_gather_object(object_list, obj, group=None) -> None:
1735 """
1736 Gathers objects from the given group into object list.
1738 Args:
1739 object_list (list[Any]): Define the output list, which size equal to the size of group.
1740 obj (Any): The object on current rank and in given process group.
1741 group (ProcessGroup, optional): The process group to gather obj. Default is ``None``, and ``None`` means
1742 global group.
1744 Returns:
1745 None. Objs are gathered into ``object_list``.
1746 """
1747 dist.all_gather_object(object_list, obj, group)
1749 @staticmethod
1750 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1751 """
1752 Synchronize all processes in the given communication group.
1754 Args:
1755 group (str, optional): The communication group to work on. Default is ``None``,
1756 meaning the default world group.
1757 async_op (bool, optional): Whether this op should be asynchronous. Default: ``False``.
1758 device_ids (list[int], optional): Reserved parameter on Ascend. Default: ``None``.
1760 Returns:
1761 CommHandle if ``async_op`` is True; otherwise ``None``.
1762 """
1763 return dist.barrier(group, async_op, device_ids)
1765 @staticmethod
1766 def init_process_group(
1767 backend: str = None,
1768 *,
1769 init_method: Optional[str] = None,
1770 timeout: Optional[timedelta] = None,
1771 world_size: int = -1,
1772 rank: int = -1,
1773 store: TCPStore = None,
1774 pg_options=None,
1775 device_id=None
1776 ) -> None:
1777 """
1778 Initialize global process group.
1780 Args:
1781 backend (str): The backend used to init process group. Default is ``"hccl"`` and now only support hccl.
1782 init_method (str, optional): URL specifying how to initialize the process group. Default is ``None``.
1783 timeout (timedelta, optional): Timeout for API executed. Default is ``None``.
1784 world_size (int): Number of processes. Default is ``-1``.
1785 rank (int, optional): Rank of the current process. Default is ``-1``.
1786 store (Store, optional): An object that stores key/value data, facilitating the exchange of inter-process
1787 communication addresses and connection information. Default is ``None``. Currently, only the
1788 ``TCPStore`` type is supported.
1789 pg_options (ProcessGroupOptions, optional): Reserved parameter. Current not take effect.
1790 device_id (int, optional): Reserved parameter. Current not take effect.
1791 """
1792 if backend is None:
1793 backend = "hccl"
1794 try:
1795 if dist.is_initialized():
1796 return
1797 except AttributeError:
1798 pass
1799 dist.init_process_group(backend=backend, init_method=init_method, timeout=timeout, world_size=world_size,
1800 rank=rank, store=store, pg_options=pg_options, device_id=device_id)
1802 @staticmethod
1803 def destroy_process_group(group: Optional[str] = None) -> None:
1804 """
1805 Destroy given process group.
1807 Args:
1808 group (str, optional): Specify the group to destroy. Default: ``None`` means ``hccl_world_group``. If group
1809 is None or "hccl_world_group", destroy global process group and all process groups relative to global
1810 process group.
1811 """
1812 if group in EXISTING_COMM_GROUPS.values():
1813 keys_to_destroy = [k for k, v in EXISTING_COMM_GROUPS.items() if v == group]
1814 for k in keys_to_destroy:
1815 del EXISTING_COMM_GROUPS[k]
1816 dist.destroy_process_group(group)
1818 @staticmethod
1819 def get_process_group_ranks(group: Optional[str] = None) -> list[int]:
1820 """
1821 Get all ranks in given process group.
1823 Args:
1824 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1826 Returns:
1827 List[int]: List of ranks in given process group.
1828 """
1829 return dist.get_process_group_ranks(group)
1831 @staticmethod
1832 def get_backend(group: Optional[str] = None) -> str:
1833 """
1834 Get the backend of given process group.
1836 Args:
1837 group (str, optional): Specify the process group to work on. Default: ``None`` means ``hccl_world_group``.
1839 Returns:
1840 str: The backend of the group.
1841 """
1842 return dist.get_backend(group)
1844 @staticmethod
1845 def split_group(parent_pg: Optional[str] = None,
1846 split_ranks: Optional[list] = None,
1847 timeout: Optional[timedelta] = None,
1848 pg_options: Optional[Any] = None,
1849 group_desc: Optional[str] = None,
1850 ) -> str:
1851 """
1852 Create split group for a specific group rank in split_ranks, which group contains current rank id.
1854 Args:
1855 parent_pg (str, Optional): A process group which the goal group split from.
1856 split_ranks (Optional[list]): A list like ``list[list[int]]``.
1857 timeout (Optional[timedelta]): Timeout for API executed. Default is ``None``.
1858 pg_options (Optional[Any]): Backend-specific group options. MindSpore can use
1859 ``{"hccl_config": {"hccl_op_expansion_mode": "AIV"}}`` to request AIV mode.
1860 group_desc (Optional[str]): Description of process group.
1862 Returns:
1863 str: The split group name.
1864 """
1865 if split_ranks is None or len(split_ranks) == 0:
1866 raise ValueError("split_ranks cannot be None or empty")
1868 rank_id = MindSporePlatform.get_rank()
1869 for split_rank in split_ranks:
1870 if rank_id in split_rank:
1871 world_group = MindSporePlatform._maybe_reuse_world_group(split_rank)
1872 if world_group is not None:
1873 return world_group
1874 split_group = MindSporePlatform.get_created_group(split_rank)
1875 if split_group:
1876 return split_group
1877 group_name = str(tuple(sorted(split_rank)))
1878 MindSporePlatform._create_group_with_options(group_name, split_rank, pg_options=pg_options)
1879 EXISTING_COMM_GROUPS[group_name] = group_name
1880 return group_name
1881 raise ValueError(f"Split group invalid rank, the Split_ranks {split_ranks} does not contain current rank"
1882 f" {rank_id}")
1884 @staticmethod
1885 def get_group_local_rank(group=None) -> int:
1886 """get group local rank id."""
1887 return dist.get_group_rank(group, MindSporePlatform.get_rank())
1889 @staticmethod
1890 def get_group_rank(group=None) -> int:
1891 return MindSporePlatform.get_group_local_rank(group)
1893 @staticmethod
1894 def no_grad():
1895 return _no_grad()
1897 @staticmethod
1898 def preserve_version_counter(tensor):
1899 from mindspore.common.api import _unsafe_preserve_version_counter # pylint: disable=C0415
1900 return _unsafe_preserve_version_counter(tensor)
1902 @staticmethod
1903 def relu(tensor):
1904 return mint.nn.functional.relu(tensor)
1906 @staticmethod
1907 def cat(tensors, dim=0):
1908 return mint.cat(tensors, dim=dim)
1910 @staticmethod
1911 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1912 return mint.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory)
1914 def get_current_stream(self):
1915 return ms.runtime.current_stream()
1917 def new_event(self):
1918 return ms.runtime.Event()
1920 def tree_map(self, fn, tree):
1921 """
1922 Apply fn to each leaf in a nested structure (list / tuple / dict),
1923 preserving the original structure.
1924 """
1925 if isinstance(tree, dict):
1926 return type(tree)(
1927 (k, self.tree_map(fn, v)) for k, v in tree.items()
1928 )
1930 if isinstance(tree, tuple):
1931 return tuple(self.tree_map(fn, v) for v in tree)
1933 if isinstance(tree, list):
1934 return [self.tree_map(fn, v) for v in tree]
1936 # leaf
1937 return fn(tree)
1939 @staticmethod
1940 def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
1941 return module.register_forward_pre_hook(hook, with_kwargs=with_kwargs)
1943 @staticmethod
1944 def register_full_backward_hook(module, hook, prepend=False):
1945 return module.register_backward_hook(hook)
1947 @staticmethod
1948 def register_full_backward_pre_hook(module, hook, prepend=False):
1949 return module.register_backward_pre_hook(hook)
1951 @property
1952 def checkpoint(self):
1953 return ms.recompute
1955 @staticmethod
1956 def checkpoint_wrapper(module, **checkpoint_kwargs):
1957 # pylint: disable=C0415
1958 from hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_wrapper import ckpt_wrapper
1959 return ckpt_wrapper(module, **checkpoint_kwargs)
1961 @staticmethod
1962 def checkpoint_exclude_wrapper(module: Any, *, save_output: bool = True) -> Any:
1963 """Wrap a Cell or callable whose activations should not be recomputed.
1965 Args:
1966 module: MindSpore Cell or callable to exclude from checkpoint replay.
1967 save_output: Whether to retain the excluded region output for replay.
1969 Returns:
1970 The platform-specific checkpoint exclusion wrapper.
1971 """
1972 # pylint: disable=C0415
1973 from hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_exclude_wrapper import (
1974 checkpoint_exclude_wrapper,
1975 )
1976 return checkpoint_exclude_wrapper(module, save_output=save_output)
1978 @staticmethod
1979 def swap_wrapper(module, policy_fn=None, group_swap=False):
1980 # pylint: disable=C0415
1981 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_wrapper
1982 return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap)
1984 @staticmethod
1985 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1986 # pylint: disable=C0415
1987 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import swap_tensor_wrapper
1988 return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap)
1990 @staticmethod
1991 def get_class_activation_wrapper():
1992 # pylint: disable=C0415
1993 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper
1994 return ActivationWrapper
1996 @property
1997 def noop_context_fn(self):
1998 return null_context_fn
2000 @staticmethod
2001 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
2002 # pylint: disable=C0415
2003 from hyper_parallel.platform.mindspore.activation_checkpoint.sac import create_selective_checkpoint_contexts
2004 return create_selective_checkpoint_contexts(policy_fn_or_list,
2005 allow_cache_entry_mutation=allow_cache_entry_mutation,
2006 group_swap=group_swap)
2008 @staticmethod
2009 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
2010 # pylint: disable=C0415
2011 from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import AsyncSaveOnCpu
2012 return AsyncSaveOnCpu(policy_fn=policy_fn, group_swap=group_swap)
2014 @staticmethod
2015 def recompute_handle_collector_ctx():
2016 # pylint: disable=C0415
2017 from mindspore.common.recompute import _recompute_handle_collector_ctx
2018 return _recompute_handle_collector_ctx()
2020 @staticmethod
2021 def recompute_handle(handle, session_id):
2022 return handle.recompute(session_id)
2024 @staticmethod
2025 def recompute_session_ctx(session_id, retain_on_unpack=False):
2026 if session_id is None:
2027 raise ValueError("session_id must not be None.")
2028 # pylint: disable=C0415
2029 from mindspore.common.recompute import _recompute_session_ctx
2030 return _recompute_session_ctx(session_id=session_id, retain_on_unpack=retain_on_unpack)
2032 @staticmethod
2033 def clear_recompute_session(session_id):
2034 # pylint: disable=C0415
2035 from mindspore.common.recompute import _clear_recompute_session
2036 return _clear_recompute_session(session_id)
2038 _MS_DEVICE_MAP = {
2039 "npu": "Ascend",
2040 "ascend": "Ascend",
2041 "gpu": "GPU",
2042 "cpu": "cpu",
2043 "": "cpu",
2044 }
2046 @staticmethod
2047 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
2048 """Allocate an uninitialized 1-D tensor buffer."""
2049 if pin_memory:
2050 return mint.empty((numel,), dtype=dtype, device="cpu", pin_memory=True)
2051 if device is None:
2052 return mint.empty((numel,), dtype=dtype)
2053 device_type = str(device).split(":", maxsplit=1)[0].lower()
2054 ms_device = MindSporePlatform._MS_DEVICE_MAP.get(device_type)
2055 if ms_device is None:
2056 raise ValueError(
2057 f"Unsupported device type '{device_type}' for MindSpore; "
2058 f"supported: {sorted(MindSporePlatform._MS_DEVICE_MAP)}"
2059 )
2060 if ms_device == "cpu":
2061 return mint.empty((numel,), dtype=dtype, device="cpu")
2062 return mint.empty((numel,), dtype=dtype, device=ms_device)
2064 @staticmethod
2065 def get_element_size(tensor):
2066 """Get Tensor Element Size"""
2067 return tensor.itemsize
2069 @staticmethod
2070 def tensor_to_numpy(tensor) -> np.ndarray:
2071 """Convert MindSpore tensor to numpy array."""
2072 return tensor.asnumpy()
2074 @staticmethod
2075 def from_numpy(np_array):
2076 """Create a host (CPU) MindSpore tensor from a numpy array."""
2077 return ms.from_numpy(np_array)
2079 @staticmethod
2081 def clip_grad_norm_(
2082 parameters, max_norm, norm_type=2.0,
2083 error_if_nonfinite=False, foreach=None,
2084 ):
2085 raise NotImplementedError(
2086 "clip_grad_norm_ is not yet supported on MindSpore"
2087 )
2089 @property
2090 def meta_device(self):
2091 return "meta"
2093 def init_on_device(self, device, include_buffers=False):
2094 return _init_on_device(device, include_buffers=include_buffers)
2096 def cast_fp_tensor(self, dtype, x):
2097 """
2098 Cast floating-point tensor to target dtype if applicable.
2099 """
2100 if (
2101 not isinstance(x, ms.Tensor)
2102 or not ms.ops.is_floating_point(x)
2103 or x.dtype == dtype
2104 ):
2105 return x
2106 return x.to(dtype)
2108 def apply_to_tensors(self, fn, container):
2109 """Recursively apply to all tensor in different kinds of container types."""
2111 def apply(x):
2112 if isinstance(x, ms.Tensor):
2113 return fn(x)
2114 if hasattr(x, "__dataclass_fields__"):
2115 dc = dataclasses.replace(x)
2116 changes = {
2117 f.name: apply(getattr(dc, f.name)) for f in dataclasses.fields(dc)
2118 }
2119 return dataclasses.replace(dc, **changes)
2120 if isinstance(x, OrderedDict):
2121 od = x.__class__()
2122 for key, value in x.items():
2123 od[key] = apply(value)
2124 return od
2125 if isinstance(x, dict):
2126 return {key: apply(value) for key, value in x.items()}
2127 if isinstance(x, tuple) and hasattr(x, "_asdict") and hasattr(x, "_fields"):
2128 res = (apply(el) for el in x)
2129 return type(x)(*res)
2130 if isinstance(x, (list, tuple, set)):
2131 return type(x)(apply(el) for el in x)
2132 return x
2134 return apply(container)
2136 @staticmethod
2137 def profiler_record(name):
2138 """Profiler context manager for recording operations using mindspore.profiler."""
2139 return ms.profiler.common.record_function.RecordFunction(name)
2141 def str_to_dtype(self, dtype_str: str) -> Any:
2142 """Resolve checkpoint dtype strings (``mindspore.*`` or short ``str(Tensor.dtype)`` e.g. ``Float32``)."""
2143 if "." in dtype_str:
2144 prefix, name = dtype_str.split(".", 1)
2145 if prefix == "mindspore":
2146 return getattr(ms, name)
2147 dtype = getattr(ms, dtype_str.lower(), None)
2148 if dtype is not None:
2149 return dtype
2150 raise ValueError(
2151 f"Expected dtype string like 'mindspore.float32' or 'Float32', got {dtype_str!r}."
2152 )
2154 def list_to_size(self, size_list: list[int]) -> tuple[int, ...]:
2155 return tuple(size_list)
2157 @staticmethod
2158 def _maybe_reuse_world_group(rank_list):
2159 """Reuse the default world group for full-world rank lists."""
2160 normalized = tuple(sorted(rank_list))
2161 world_ranks = tuple(range(MindSporePlatform.get_world_size()))
2162 if normalized != world_ranks:
2163 return None
2165 EXISTING_COMM_GROUPS[str(normalized)] = GlobalComm.WORLD_COMM_GROUP
2166 return GlobalComm.WORLD_COMM_GROUP