Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / platform.py: 66%
434 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"""framework platform api"""
16# Backend platform modules intentionally import this abstraction to register
17# their implementations; the resulting import cycle is architectural.
18# pylint: disable=cyclic-import
19import os
20from datetime import timedelta
21from enum import auto, Enum
22from typing import Any, Optional, Sequence, Union
24import numpy as np
26# Environment variable name used to specify the AI framework platform to use
27HYPER_PARALLEL_PLATFORM = "HYPER_PARALLEL_PLATFORM"
29# Identifier for the MindSpore framework
30HYPER_PARALLEL_PLATFORM_MINDSPORE = "mindspore"
32# Identifier for the PyTorch framework
33HYPER_PARALLEL_PLATFORM_TORCH = "torch"
36class AsyncHandle:
37 """Idempotent wait handle for an async collective operation.
39 Wraps the async tensor returned by
40 :meth:`Platform.differentiable_all_to_all_single_async` and provides a
41 :meth:`wait` method that is safe to call multiple times.
42 """
44 def __init__(self, async_tensor) -> None:
45 self._tensor = async_tensor
46 self._waited = False
48 def wait(self):
49 """Wait for the async collective to complete.
51 Idempotent — the first call blocks until the collective finishes;
52 subsequent calls are no-ops.
54 Returns:
55 The now-materialised result tensor.
56 """
57 if not self._waited:
58 get_platform().wait_async_tensor(self._tensor)
59 self._waited = True
60 return self._tensor
63class PlatformType(Enum):
64 """Enumeration class for AI framework platform types.
66 Used to identify different deep learning framework platform types.
67 """
68 MINDSPORE = auto()
69 PYTORCH = auto()
72# Global platform instance, used to cache the created platform object
73platform = None
76def get_mindspore_platform():
77 """Create and return a MindSpore platform instance.
79 Returns:
80 MindSporePlatform: A MindSpore platform instance.
81 """
82 # pylint: disable=C0415
83 from hyper_parallel.platform.mindspore.platform import MindSporePlatform
84 global platform
85 platform = MindSporePlatform()
86 return platform
89def get_torch_platform():
90 """Create and return a PyTorch platform instance.
92 Returns:
93 TorchPlatform: A PyTorch platform instance.
94 """
95 # pylint: disable=C0415
96 from hyper_parallel.platform.torch.platform import TorchPlatform
97 global platform
98 platform = TorchPlatform()
99 return platform
102def get_platform():
103 """Obtain a framework platform instance.
105 Returns the appropriate AI framework platform instance based on environment variables or a default priority order.
106 The lookup priority is as follows:
107 1. Platform specified by environment variable
108 2. MindSpore platform (default preferred choice)
109 3. PyTorch platform (fallback option)
111 Returns:
112 Platform: An instance of the framework platform
114 Raises:
115 ImportError: Raised when none of the supported frameworks are available
116 """
117 if platform is not None:
118 return platform
119 platform_type = os.environ.get(HYPER_PARALLEL_PLATFORM)
120 if platform_type is not None and isinstance(platform_type, str):
121 platform_type = platform_type.lower()
122 if platform_type == HYPER_PARALLEL_PLATFORM_MINDSPORE:
123 return get_mindspore_platform()
124 if platform_type == HYPER_PARALLEL_PLATFORM_TORCH:
125 return get_torch_platform()
126 try:
127 return get_mindspore_platform()
128 except ImportError:
129 return get_torch_platform()
132EXISTING_COMM_GROUPS = {}
135class Platform:
136 """Platform api"""
137 current_grad_handle = None
138 post_grad_handle_process = None
139 grad_sync_stream = None
141 @property
142 def custom_ops(self):
143 """Return the platform-specific custom ops interface.
145 Subclasses MUST override this property to return an object that
146 exposes the platform-specific custom operator implementations.
148 Returns:
149 object: Platform-specific custom ops class instance.
150 """
151 raise NotImplementedError(
152 "Platform subclasses must implement custom_ops"
153 )
155 @staticmethod
156 def get_swap_optimizer():
157 """Return the active backend's optimizer-state swap wrapper class."""
158 raise NotImplementedError("Platform subclasses must implement get_swap_optimizer")
160 @staticmethod
161 def get_rank():
162 """Get the rank of the current process in the default process group.
164 Returns:
165 int: The rank of the current process.
166 """
167 raise NotImplementedError("Platform subclasses must implement get_rank")
169 @staticmethod
170 def get_global_rank(group, group_rank):
171 """Convert a group rank to its global rank.
173 Args:
174 group: The process group to query.
175 group_rank (int): The rank within the group.
177 Returns:
178 int: The global rank corresponding to the group rank.
179 """
180 raise NotImplementedError("Platform subclasses must implement get_global_rank")
182 @staticmethod
183 def get_group_rank(group):
184 """Return this process's rank within *group*."""
185 raise NotImplementedError("Platform subclasses must implement get_group_rank")
187 @staticmethod
188 def get_world_size():
189 """Get the total number of processes in the default process group.
191 Returns:
192 int: The world size (total number of processes).
193 """
194 raise NotImplementedError("Platform subclasses must implement get_world_size")
196 @staticmethod
197 def get_op_name(func):
198 """Get the canonical name of an operator function.
200 Args:
201 func: The operator function to query.
203 Returns:
204 str: The canonical name of the operator.
205 """
206 raise NotImplementedError("Platform subclasses must implement get_op_name")
208 @staticmethod
209 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
210 """Perform differentiable all-gather and concatenate tensors along a dimension.
212 Args:
213 data: The input tensor to gather.
214 group: The process group for collective communication.
215 concat_size (int): The size to concatenate along concat_dim.
216 concat_dim (int): The dimension along which to concatenate.
217 rank_list: Optional rank order expected by the logical layout.
219 Returns:
220 The concatenated tensor after all-gather operation.
221 """
222 raise NotImplementedError("Platform subclasses must implement differentiable_all_gather_concat")
224 @staticmethod
225 def chunk(data, split_dim, split_size, index):
226 """Split tensor along a dimension and return the chunk at the given index.
228 Args:
229 data: The input tensor to split.
230 split_dim (int): The dimension along which to split.
231 split_size (int): The size of each split chunk.
232 index (int): The index of the chunk to return.
234 Returns:
235 The tensor chunk at the specified index.
236 """
237 raise NotImplementedError("Platform subclasses must implement chunk")
239 @staticmethod
240 def differentiable_all_to_all(input_data, output_shape, group):
241 """Perform differentiable all-to-all communication.
243 Args:
244 input_data: The input tensor to redistribute.
245 output_shape: The shape of the output tensor.
246 group: The process group for collective communication.
248 Returns:
249 The output tensor after all-to-all operation.
250 """
251 raise NotImplementedError("Platform subclasses must implement differentiable_all_to_all")
253 @staticmethod
254 def tensor_type_cast(input_data, cast_type):
255 """Cast tensor to a specified dtype.
257 Args:
258 input_data: The input tensor to cast.
259 cast_type: The target dtype to cast to.
261 Returns:
262 The tensor cast to the specified dtype.
263 """
264 raise NotImplementedError("Platform subclasses must implement tensor_type_cast")
266 @staticmethod
267 def is_tensor(obj: Any) -> bool:
268 """Return True if ``obj`` is this framework's tensor type."""
269 raise NotImplementedError("Platform subclasses must implement is_tensor")
271 @staticmethod
272 def get_tensor_storage_size(tensor: Any) -> int:
273 """Return serialized byte size (numel * element size) for this framework's tensor."""
274 raise NotImplementedError("Platform subclasses must implement get_tensor_storage_size")
276 @staticmethod
277 def differentiable_all_reduce(data, op, group):
278 """Perform differentiable all-reduce operation.
280 Args:
281 data: The input tensor to reduce.
282 op: The reduction operation (e.g., sum, max, min).
283 group: The process group for collective communication.
285 Returns:
286 The reduced tensor with gradients supported.
287 """
288 raise NotImplementedError("Platform subclasses must implement differentiable_all_reduce")
290 @staticmethod
291 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
292 """Perform differentiable reduce-scatter operation.
294 Args:
295 data: The input tensor to reduce and scatter.
296 dev_num (int): The number of devices to scatter across.
297 axis (int): The axis along which to scatter.
298 op: The reduction operation (e.g., sum, max, min).
299 group: The process group for collective communication.
301 Returns:
302 The scattered tensor chunk with gradients supported.
303 """
304 raise NotImplementedError("Platform subclasses must implement differentiable_reduce_scatter")
306 @staticmethod
307 def init_parameters(module, stage_index):
308 """Initialize parameters for a module at a specific pipeline stage.
310 This method is primarily needed for MindSpore platform which requires
311 explicit parameter initialization interface.
313 Args:
314 module: The module whose parameters need to be initialized.
315 stage_index (int): The pipeline stage index for the module.
317 Raises:
318 ValueError: If module is None or stage_index is negative.
319 """
320 if module is None:
321 raise ValueError("input module must not be none.")
322 if stage_index < 0:
323 raise ValueError("input stage_index must be positive.")
325 @staticmethod
326 def get_cell_construct(cell):
327 """Get the construct (forward) function of a cell/module.
329 Args:
330 cell: The cell or module to get the construct function from.
332 Returns:
333 The construct/forward callable of the cell.
334 """
335 raise NotImplementedError("Platform subclasses must implement get_cell_construct")
337 @staticmethod
338 def get_cells_and_names(cell):
339 """Get all nested cells/modules and their names.
341 Args:
342 cell: The root cell or module to traverse.
344 Returns:
345 list: A list of tuples containing (name, cell) pairs.
346 """
347 raise NotImplementedError("Platform subclasses must implement get_cells_and_names")
349 @staticmethod
350 def get_modules(module):
351 """Return all sub-modules contained in the given module."""
352 raise NotImplementedError("Platform subclasses must implement get_modules")
354 @staticmethod
355 def search_parameter_by_name(cell, param_name: str):
356 """Search for a parameter by name within a cell/module.
358 Args:
359 cell: The cell or module to search in.
360 param_name (str): The name of the parameter to find.
362 Returns:
363 The parameter if found, otherwise None.
364 """
365 raise NotImplementedError("Platform subclasses must implement search_parameter_by_name")
367 @staticmethod
368 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
369 """Update a parameter by name within a cell/module.
371 Args:
372 cell: The cell or module containing the parameter.
373 result (tuple): A tuple containing (param_name, parameter) to update.
374 new_param: The new parameter value to set.
376 Returns:
377 bool: True if update was successful, False otherwise.
378 """
379 raise NotImplementedError("Platform subclasses must implement update_parameter_by_name")
381 @staticmethod
382 def set_layout_into_parameter(param, layout):
383 """Attach a DTensor layout to a parameter.
385 Args:
386 param: The parameter to attach the layout to.
387 layout: The DTensor layout describing tensor distribution.
388 """
389 raise NotImplementedError("Platform subclasses must implement set_layout_into_parameter")
391 @staticmethod
392 def get_param_local_shape(param):
393 """Get the local shape of a distributed parameter.
395 Args:
396 param: The parameter to query.
398 Returns:
399 tuple: The local shape of the parameter shard.
400 """
401 raise NotImplementedError("Platform subclasses must implement get_param_local_shape")
403 @staticmethod
404 def get_param_local_data(param):
405 """Get the local data tensor of a distributed parameter.
407 Args:
408 param: The parameter to query.
410 Returns:
411 The local tensor data of the parameter shard.
412 """
413 raise NotImplementedError("Platform subclasses must implement get_param_local_data")
415 @staticmethod
416 def update_param_data(param, data):
417 """Update the data of a parameter with new tensor data.
419 Args:
420 param: The parameter to update.
421 data: The new tensor data to assign.
422 """
423 raise NotImplementedError("Platform subclasses must implement update_param_data")
425 @staticmethod
426 def get_param_type_size(param):
427 """Get the size in bytes of a parameter's dtype.
429 Args:
430 param: The parameter to query.
432 Returns:
433 int: The size in bytes of the parameter's data type.
434 """
435 raise NotImplementedError("Platform subclasses must implement get_param_type_size")
437 @staticmethod
438 def new_zero_parameter(param_shape, param_type, requires_grad, device):
439 """Create a new parameter initialized with zeros.
441 Args:
442 param_shape (tuple): The shape of the parameter.
443 param_type: The dtype of the parameter.
444 requires_grad (bool): Whether the parameter requires gradients.
445 device: The device on which to create the parameter.
447 Returns:
448 A new parameter tensor filled with zeros.
449 """
450 raise NotImplementedError("Platform subclasses must implement new_zero_parameter")
452 @staticmethod
453 def new_tensor(tensor_shape, tensor_type, device):
454 """Create a new tensor with the specified shape, dtype, and device.
456 Args:
457 tensor_shape (tuple): The shape of the tensor.
458 tensor_type: The dtype of the tensor.
459 device: The device on which to create the tensor.
461 Returns:
462 A new tensor with uninitialized values.
463 """
464 raise NotImplementedError("Platform subclasses must implement new_tensor")
466 @staticmethod
467 def full_like(tensor, fill_value, dtype=None):
468 """Create a tensor filled with a value, with same shape as input.
470 Args:
471 tensor: The input tensor to copy shape from.
472 fill_value: The value to fill the new tensor with.
473 dtype: Optional dtype for the new tensor. If None, uses input tensor's dtype.
475 Returns:
476 A new tensor filled with the specified value.
477 """
478 raise NotImplementedError("Platform subclasses must implement full_like")
480 @staticmethod
481 def set_tensor_requires_grad(input_tensor):
482 """Enable gradient tracking for a tensor in-place.
484 Args:
485 input_tensor: The tensor to enable gradients for.
487 Returns:
488 The same tensor with requires_grad set to True.
489 """
490 raise NotImplementedError("Platform subclasses must implement set_tensor_requires_grad")
492 @staticmethod
493 def all_gather_into_tensor(data, group_info, async_op=False):
494 """Gather tensors from all ranks into a single output tensor.
496 Args:
497 data: The input tensor to gather.
498 group_info: The process group for collective communication.
499 async_op (bool): If True, returns a work handle for async operation.
501 Returns:
502 The gathered tensor, or a tuple of (tensor, handle) if async_op is True.
503 """
504 raise NotImplementedError("Platform subclasses must implement all_gather_into_tensor")
506 @staticmethod
507 def all_reduce(data, group_info, async_op=False):
508 """Reduce tensors across all ranks using specified operation.
510 Args:
511 data: The input tensor to reduce.
512 group_info: The process group for collective communication.
513 async_op (bool): If True, returns a work handle for async operation.
515 Returns:
516 The reduced tensor, or a tuple of (tensor, handle) if async_op is True.
517 """
518 raise NotImplementedError("Platform subclasses must implement all_reduce")
520 @staticmethod
521 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
522 """Broadcast tensor from source rank to all ranks in group."""
523 raise NotImplementedError("Platform subclasses must implement broadcast")
525 @staticmethod
526 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
527 """Scatter tensor list from source rank to all ranks in group."""
528 raise NotImplementedError("Platform subclasses must implement scatter")
530 @staticmethod
531 def isend(tensor, dst=None, group=None, tag=0):
532 """Send tensor asynchronously to destination rank.
534 Args:
535 tensor: The tensor to send.
536 dst (int, optional): The destination rank. Defaults to None.
537 group: The process group for communication. Defaults to None.
538 tag (int): A tag to identify the send operation. Defaults to 0.
540 Returns:
541 A work handle that can be waited on.
542 """
543 raise NotImplementedError("Platform subclasses must implement isend")
545 @staticmethod
546 def irecv(tensor, src=None, group=None, tag=0):
547 """Receive tensor asynchronously from source rank.
549 Args:
550 tensor: The tensor buffer to receive data into.
551 src (int, optional): The source rank. Defaults to None.
552 group: The process group for communication. Defaults to None.
553 tag (int): A tag to identify the receive operation. Defaults to 0.
555 Returns:
556 A work handle that can be waited on.
557 """
558 raise NotImplementedError("Platform subclasses must implement irecv")
560 @staticmethod
561 def p2p_op(op_type, tensor, peer, group=None):
562 """Build a batched-P2P descriptor (no launch).
564 Returns an opaque object understood by :meth:`batch_isend_irecv`.
565 Lets callers assemble a mixed send/recv batch that the backend can
566 run concurrently (e.g. TX/RX duplex on one link) in a single op.
568 Args:
569 op_type (str): ``"isend"`` or ``"irecv"``.
570 tensor: Tensor to send, or the buffer to receive into.
571 peer (int): Global rank of the peer.
572 group: Process group. ``None`` uses the default group.
574 Returns:
575 A backend P2P-op descriptor.
576 """
577 raise NotImplementedError("Platform subclasses must implement p2p_op")
579 @staticmethod
580 def batch_isend_irecv(p2p_ops):
581 """Launch a batch of :meth:`p2p_op` descriptors as one async op.
583 The whole batch shares a single completion handle (the backend runs
584 the items concurrently on one comm stream), so a send and a recv to
585 the same peer overlap on the duplex link.
587 Args:
588 p2p_ops (list): Descriptors from :meth:`p2p_op`.
590 Returns:
591 A single work handle covering the whole batch, or ``None`` when
592 ``p2p_ops`` is empty.
593 """
594 raise NotImplementedError("Platform subclasses must implement batch_isend_irecv")
596 @staticmethod
597 def prepare_batch_p2p_group(group: Any = None) -> None:
598 """Prepare a process group before its first batched P2P operation.
600 Backends that require full-group participation before subset batched
601 P2P should synchronize the group here. Other backends may implement
602 this as a no-op.
604 Args:
605 group: The process group used by the batched P2P operations.
606 ``None`` uses the default group.
607 """
608 raise NotImplementedError("Platform subclasses must implement prepare_batch_p2p_group")
610 @staticmethod
611 def p2p_exchange(tensor, peer_rank: int, group=None):
612 """Differentiable symmetric P2P exchange (send local tensor, receive peer's tensor).
614 Sends ``tensor`` to ``peer_rank`` and simultaneously receives the peer's
615 tensor. The operation is differentiable: the backward pass performs the
616 same symmetric exchange on the upstream gradient.
618 Args:
619 tensor: Local tensor to send.
620 peer_rank (int): Global rank of the communication peer.
621 group: Process group. ``None`` uses the default group.
623 Returns:
624 Tensor received from ``peer_rank``, with the same shape and dtype as
625 the input ``tensor``.
626 """
627 raise NotImplementedError("Platform subclasses must implement p2p_exchange")
629 @staticmethod
630 def send_object_list(obj_list, dst=None, group=None):
631 """Send a list of Python objects to destination rank.
633 Args:
634 obj_list (list): The list of Python objects to send.
635 dst (int, optional): The destination rank. Defaults to None.
636 group: The process group for communication. Defaults to None.
637 """
638 raise NotImplementedError("Platform subclasses must implement send_object_list")
640 @staticmethod
641 def recv_object_list(obj_list, src=None, group=None):
642 """Receive a list of Python objects from source rank.
644 Args:
645 obj_list (list): The list buffer to receive objects into.
646 src (int, optional): The source rank. Defaults to None.
647 group: The process group for communication. Defaults to None.
648 """
649 raise NotImplementedError("Platform subclasses must implement recv_object_list")
651 @staticmethod
652 def reduce_scatter_tensor(data, group_info, async_op=False):
653 """Reduce and scatter tensor across all ranks in group.
655 Args:
656 data: The input tensor to reduce and scatter.
657 group_info: The process group for collective communication.
658 async_op (bool): If True, returns a work handle for async operation.
660 Returns:
661 The scattered tensor chunk, or a tuple of (tensor, handle) if async_op is True.
662 """
663 raise NotImplementedError("Platform subclasses must implement reduce_scatter_tensor")
665 @staticmethod
666 def all_gather_single(input_tensor, output_shape, group, async_op=False):
667 """All-gather tensor shards with optional async execution.
669 Args:
670 input_tensor: Input tensor whose leading dimension is gathered.
671 output_shape: Shape of the gathered output tensor.
672 group: Process group (ProcessGroup for torch, group name string for mindspore).
673 async_op: If True, returns an async work handle.
675 Returns:
676 Tuple ``(output, work)`` where *output* is the gathered tensor and
677 *work* is the async handle (``None`` when ``async_op=False``).
678 """
679 raise NotImplementedError("Platform subclasses must implement all_gather_single")
681 @staticmethod
682 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
683 """Reduce-scatter a tensor with optional async execution.
685 Args:
686 input_tensor: Input tensor whose leading dimension is split across ranks.
687 output_shape: Shape of the local reduced output tensor.
688 group: Process group (ProcessGroup for torch, group name string for mindspore).
689 async_op: If True, returns an async work handle.
691 Returns:
692 Tuple ``(output, work)`` where *output* is the local shard and
693 *work* is the async handle (``None`` when ``async_op=False``).
694 """
695 raise NotImplementedError("Platform subclasses must implement reduce_scatter_single")
697 @staticmethod
698 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
699 """All-to-all single collective with optional async execution.
701 Args:
702 input_tensor: Input tensor to scatter.
703 output_shape: Shape of the pre-allocated output tensor.
704 group: Process group (ProcessGroup for torch, group name string for mindspore).
705 async_op: If True, returns a work handle; the output tensor is
706 filled only after ``work.wait()`` is called.
708 Returns:
709 Tuple ``(output, work)`` where *output* is the result tensor and
710 *work* is the async handle (``None`` when ``async_op=False``).
712 Raises:
713 NotImplementedError: Must be implemented by platform subclasses.
714 """
715 raise NotImplementedError("Platform subclasses must implement all_to_all_single")
717 @staticmethod
718 def differentiable_variable_all_gather(
719 input_tensor: Any, output_splits: Sequence[int], group: Any) -> Any:
720 """Gather variable dim-zero shards on every rank with gradient support.
722 Args:
723 input_tensor: Local input shaped ``[local_rows, *feature_dims]``.
724 output_splits: Dim-zero rows contributed by each group rank.
725 group: Raw platform process group.
727 Returns:
728 Tensor concatenated in group-rank order along dim zero.
730 Raises:
731 NotImplementedError: Must be implemented by platform subclasses.
732 """
733 raise NotImplementedError(
734 "Platform subclasses must implement differentiable_variable_all_gather"
735 )
737 @staticmethod
738 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
739 handle_box=None):
740 """Differentiable wrapper that waits for a pre-launched async all-gather.
742 Forward waits for the all-gather handle and reconstructs the tensor by
743 moving the gathered leading dimension back to ``gather_dim``.
745 Backward launches the reverse reduce-scatter. If ``handle_box`` is a
746 mutable list, the reduce-scatter handle is appended there and a zero
747 gradient is returned to be replaced by the caller's backward pre-hook.
748 If ``handle_box`` is ``None``, the reduce-scatter is waited immediately
749 and its local result is returned, preserving composability with an
750 upstream autograd communication op.
752 Args:
753 x: Original input tensor; anchors the op in the autograd graph.
754 work: Async work handle from all-gather.
755 out_perm: Output buffer filled by all-gather.
756 group: Communication group for backward reduce-scatter.
757 world_size: Group size.
758 gather_dim: Dimension gathered in forward.
759 handle_box: Optional mutable list for deferred backward wait.
761 Returns:
762 Gathered tensor connected to the autograd graph through *x*.
763 """
764 raise NotImplementedError("Platform subclasses must implement differentiable_async_allgather_wait")
766 @staticmethod
767 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim,
768 handle_box=None):
769 """Differentiable wrapper that waits for a pre-launched async A2A.
771 Wraps the wait-and-reconstruct step in the platform autograd mechanism
772 so gradients flow correctly through the all-to-all communication.
774 The A2A direction is seq→head (forward): the output gathers along
775 ``concat_dim`` (sequence grows from S/cp to S) and scatters along
776 ``split_dim`` (heads shrink from H to H/ws).
778 In backward, launches an async head→seq A2A on the incoming gradient
779 and appends ``(work, out_perm)`` to ``handle_box`` so the caller can
780 wait just before the projection GEMM, achieving GEMM–A2A overlap.
782 Args:
783 x: Original projection output tensor; anchors the op
784 in the autograd graph.
785 work: Async work handle from ``all_to_all_single(async_op=True)``.
786 out_perm: Output buffer filled once ``work.wait()`` completes
787 (shape ``[ws, ...]``).
788 group: Process group for the reverse A2A in backward.
789 world_size: CP/Ulysses degree.
790 concat_dim: Dimension that is gathered (concatenated) in forward;
791 typically the sequence dimension.
792 split_dim: Dimension that is scattered (split) in forward;
793 typically the head dimension.
794 handle_box: Optional mutable list ``[]``. In backward, ``(work, out_perm)``
795 for the reverse A2A is appended here so the pre-hook can wait.
797 Returns:
798 Result tensor with ``concat_dim`` gathered and ``split_dim`` split,
799 connected to the autograd graph through *x*.
801 Raises:
802 NotImplementedError: Must be implemented by platform subclasses.
803 """
804 raise NotImplementedError("Platform subclasses must implement differentiable_async_a2a_wait")
806 @staticmethod
807 def differentiable_sync_hook(x, hook_name: str, coordinator):
808 """Identity operation that intercepts both forward and backward to call
809 coordinator rendezvous, enabling deterministic comm/compute overlap.
811 This is the differentiable building block for dual-pipe schedules.
812 In the forward pass the coordinator is invoked with the forward-side
813 roles for ``hook_name``; in the backward pass it is invoked with the
814 backward-side roles. The tensor value and gradient flow through
815 unchanged.
817 Args:
818 x: Input tensor. Returned as-is; gradients flow through.
819 hook_name: One of ``"A"``, ``"B"``, ``"C"``, ``"D"`` identifying
820 the position relative to MoE dispatch/combine.
821 coordinator: A :class:`HookCoordinator` instance shared between the
822 forward and backward threads.
824 Returns:
825 The same tensor *x*, attached to the autograd graph so that the
826 backward hook will fire.
827 """
828 raise NotImplementedError("Platform subclasses must implement differentiable_sync_hook")
830 @staticmethod
831 def differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group):
832 """Variable-split all-to-all single that supports gradient flow.
834 Unlike ``all_to_all_single`` (which is not differentiable), this method
835 wraps the collective in an autograd function so gradients are correctly
836 routed back through the reverse all-to-all in the backward pass.
837 Intended for Expert Parallelism token dispatch / combine.
839 Args:
840 input_tensor: Input tensor to scatter. Shape ``[sum(input_splits), *feature_dims]``.
841 input_splits: Per-rank sizes of data sent from this rank (list of ints,
842 length equal to ep_degree).
843 output_splits: Per-rank sizes of data received by this rank (list of ints,
844 length equal to ep_degree).
845 group: Process group (ProcessGroup for torch, group name str for mindspore).
847 Returns:
848 Output tensor of shape ``[sum(output_splits), *feature_dims]``.
850 Raises:
851 NotImplementedError: Must be implemented by platform subclasses.
852 """
853 raise NotImplementedError("Platform subclasses must implement differentiable_all_to_all_single")
855 @staticmethod
856 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
857 """Async variant of :meth:`differentiable_all_to_all_single`.
859 Same semantics but launches the collective with ``async_op=True`` and
860 only performs a stream-level ``wait`` — the host returns immediately
861 after dispatching the kernel. Intended for dual-pipe comm/compute
862 overlap paths where the paired COMPUTE side's rendezvous notify must
863 fire right after kernel launch (not after the collective actually
864 completes on device).
866 Args:
867 input_tensor: Input tensor to scatter. Shape ``[sum(input_splits), *feature_dims]``.
868 input_splits: Per-rank sizes of data sent from this rank.
869 output_splits: Per-rank sizes of data received by this rank.
870 group: Process group.
872 Returns:
873 Output tensor of shape ``[sum(output_splits), *feature_dims]``.
875 Raises:
876 NotImplementedError: Must be implemented by platform subclasses.
877 """
878 raise NotImplementedError(
879 "Platform subclasses must implement differentiable_all_to_all_single_async"
880 )
882 @staticmethod
883 def wait_async_tensor(tensor):
884 """Wait for an async collective tensor to become materialised.
886 Intended for use with :class:`AsyncHandle` so that callers can
887 wait on an async all-to-all result without importing framework-specific
888 modules directly. The call is **idempotent** — waiting on an already-
889 completed tensor is a no-op.
891 Args:
892 tensor: An async collective tensor (e.g. PyTorch
893 ``AsyncCollectiveTensor``) whose values have not yet been
894 fully written by the remote ranks.
896 Returns:
897 The same *tensor*, now guaranteed to be fully materialised.
899 Raises:
900 NotImplementedError: Must be implemented by platform subclasses.
901 """
902 raise NotImplementedError(
903 "Platform subclasses must implement wait_async_tensor"
904 )
906 @staticmethod
907 def arange(start, end=None, step=1, dtype=None, device=None):
908 """Create a 1-D tensor with evenly spaced values.
910 Args:
911 start: Start of interval (inclusive). If *end* is ``None``,
912 treated as the stop value and *start* defaults to 0.
913 end: End of interval (exclusive). Defaults to ``None``.
914 step: Step size. Defaults to ``1``.
915 dtype: Data type. ``None`` uses the framework default (int64).
916 device: Target device.
918 Returns:
919 1-D tensor ``[start, start+step, ..., end)``.
921 Raises:
922 NotImplementedError: Must be implemented by platform subclasses.
923 """
924 raise NotImplementedError("Platform subclasses must implement arange")
926 @staticmethod
927 def zeros(size, dtype=None, device=None):
928 """Create a zero-filled tensor of the given shape.
930 Args:
931 size: Shape of the tensor (a single tuple/list).
932 dtype: Desired data type. ``None`` uses the framework default (float32).
933 device: Target device. ``None`` uses the framework default.
935 Returns:
936 Zero-filled tensor of the specified shape.
938 Raises:
939 NotImplementedError: Must be implemented by platform subclasses.
940 """
941 raise NotImplementedError("Platform subclasses must implement zeros")
943 @staticmethod
944 def parameters_dict(cell):
945 """Get the parameters dictionary of a cell/module.
947 Args:
948 cell: The cell or module to get parameters from.
950 Returns:
951 dict: A dictionary mapping parameter names to parameters.
952 """
953 raise NotImplementedError("Platform subclasses must implement parameters_dict")
955 @staticmethod
956 def buffers_dict(cell: Any) -> Any:
957 """Get the named buffers of a cell/module.
959 Args:
960 cell: The cell or module to get buffers from.
962 Returns:
963 An iterable of ``(name, buffer)`` pairs, including non-persistent
964 buffers and buffers registered by child modules.
965 """
966 raise NotImplementedError("Platform subclasses must implement buffers_dict")
968 @staticmethod
969 def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
970 """Get the state dictionary of a model.
972 Args:
973 model: The model to extract state from.
974 options: Optional configuration for state dict extraction.
976 Returns:
977 dict: The state dictionary containing model parameters and buffers.
979 Raises:
980 NotImplementedError: Platform subclasses must implement this method.
981 """
982 raise NotImplementedError(
983 "Platform subclasses must implement get_model_state_dict"
984 )
986 @staticmethod
987 def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
988 """Set the state dictionary of a model.
990 Args:
991 model: The model to load state into.
992 model_state_dict: The state dict to load into the model.
993 options: Optional configuration for state dict loading.
995 Returns:
996 None.
998 Raises:
999 NotImplementedError: Platform subclasses must implement this method.
1000 """
1001 raise NotImplementedError(
1002 "Platform subclasses must implement set_model_state_dict"
1003 )
1005 @staticmethod
1006 def save_checkpoint(cell, file_path: str, ckpt_format: str = "safetensors") -> None:
1007 """Save a cell/module checkpoint to file.
1009 Args:
1010 cell: The cell or module to save.
1011 file_path (str): The path to save the checkpoint to.
1012 ckpt_format (str): The file format.
1013 """
1014 raise NotImplementedError("Platform subclasses must implement save_checkpoint")
1016 @staticmethod
1017 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
1018 """Load a checkpoint from file.
1020 Args:
1021 file_path (str): The path to load the checkpoint from.
1022 ckpt_format (str): The file format.
1024 Returns:
1025 dict: The loaded checkpoint state dictionary.
1026 """
1027 raise NotImplementedError("Platform subclasses must implement load_checkpoint")
1029 def _create_group(self, rank_list):
1030 """Create a new process group with the specified ranks.
1032 Internal method to be implemented by subclasses.
1034 Args:
1035 rank_list (list): List of ranks to include in the group.
1037 Returns:
1038 The newly created process group.
1039 """
1040 raise NotImplementedError("Platform subclasses must implement _create_group")
1042 def new_stream(self):
1043 """Create a new compute stream for asynchronous operations.
1045 Returns:
1046 A new stream object for the current device.
1047 """
1048 raise NotImplementedError("Platform subclasses must implement new_stream")
1050 def get_stream_context(self):
1051 """Get a context manager for executing operations on a specific stream.
1053 Returns:
1054 A context manager that can be used with 'with' statement to set stream.
1055 """
1056 raise NotImplementedError("Platform subclasses must implement get_stream_context")
1058 @staticmethod
1059 def get_tensor_transform():
1060 """Get the tensor transformation utilities for the current framework.
1062 Returns:
1063 A module or object containing tensor transformation functions.
1064 """
1065 raise NotImplementedError("Platform subclasses must implement get_tensor_transform")
1067 @staticmethod
1068 def construct_strided_slice(x, begin, end, stride):
1069 """Construct a strided slice operation on a tensor.
1071 Args:
1072 x: The input tensor to slice.
1073 begin: The starting indices for each dimension.
1074 end: The ending indices for each dimension.
1075 stride: The stride for each dimension.
1077 Returns:
1078 The sliced tensor.
1079 """
1080 raise NotImplementedError("Platform subclasses must implement construct_strided_slice")
1082 @staticmethod
1083 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1084 """Split inputs into micro-batches for pipeline parallelism.
1086 Args:
1087 micro_batch_num (int): The number of micro-batches to create.
1088 args_batch_dim (list, optional): Batch dimension for each positional arg.
1089 kwargs_batch_dim (dict, optional): Batch dimension for each keyword arg.
1091 Returns:
1092 A decorator that splits function inputs into micro-batches.
1093 """
1094 raise NotImplementedError("Platform subclasses must implement micro_batch")
1096 @staticmethod
1097 def get_symmetric_memory_handler():
1098 """Return a platform-specific symmetric memory handler instance."""
1099 raise NotImplementedError("Platform subclasses must implement get_symmetric_memory_handler")
1101 @staticmethod
1102 def load_into_param(param, data):
1103 """Load data into a parameter, handling framework-specific semantics."""
1104 raise NotImplementedError("Platform subclasses must implement load_into_param")
1106 def create_group(self, rank_list):
1107 """Create or retrieve a communication group with the specified ranks.
1109 If a group with the same rank list already exists, returns the existing
1110 group instead of creating a new one.
1112 Args:
1113 rank_list (list): List of ranks to include in the group.
1115 Returns:
1116 The process group for the specified ranks.
1117 """
1118 group_key = str(tuple(sorted(rank_list)))
1119 if group_key in EXISTING_COMM_GROUPS:
1120 return EXISTING_COMM_GROUPS[group_key]
1122 group = self._create_group(rank_list)
1123 EXISTING_COMM_GROUPS[group_key] = group
1124 return group
1126 @staticmethod
1127 def _process_current_handle():
1128 """Wait for the current gradient handle and execute post-process callback.
1130 Internal method to synchronize pending gradient operations.
1131 """
1132 if Platform.current_grad_handle is None:
1133 return
1135 Platform.current_grad_handle.wait()
1136 if Platform.post_grad_handle_process is None:
1137 return
1138 # pylint: disable=E1102
1139 Platform.post_grad_handle_process()
1141 def set_grad_reduce_handle(self, handle, post_process=None):
1142 """Set a new gradient reduction handle after waiting for the current one.
1144 Waits for any pending gradient handle on the grad sync stream, then
1145 sets the new handle and optional post-process callback.
1147 Args:
1148 handle: The async work handle for gradient reduction.
1149 post_process (callable, optional): Callback to run after handle completes.
1150 """
1151 if Platform.grad_sync_stream is None:
1152 Platform.grad_sync_stream = self.new_stream()
1153 stream_context = self.get_stream_context()
1154 with stream_context(Platform.grad_sync_stream):
1155 Platform._process_current_handle()
1156 Platform.current_grad_handle = handle
1157 Platform.post_grad_handle_process = post_process
1159 def wait_grad_handle(self):
1160 """Wait for the current gradient handle to complete.
1162 Blocks until the current gradient reduction handle completes and
1163 clears the handle state.
1164 """
1165 if Platform.current_grad_handle is None:
1166 return
1167 if Platform.grad_sync_stream is None:
1168 Platform.grad_sync_stream = self.new_stream()
1169 stream_context = self.get_stream_context()
1170 with stream_context(Platform.grad_sync_stream):
1171 Platform._process_current_handle()
1172 sync_event = Platform.grad_sync_stream.record_event()
1173 sync_event.wait()
1174 Platform.current_grad_handle = None
1175 Platform.post_grad_handle_process = None
1177 @staticmethod
1178 def all_gather_object(object_list, obj, group=None) -> None:
1179 """Gather Python objects from all ranks into a list.
1181 Each rank contributes its object, and all ranks receive the complete list.
1183 Args:
1184 object_list (list): List to store gathered objects (output parameter).
1185 obj: The Python object from this rank to contribute.
1186 group: The process group for communication. Defaults to None (default group).
1187 """
1188 raise NotImplementedError("Platform subclasses must implement all_gather_object")
1190 @staticmethod
1191 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1192 """Synchronize all processes in the given process group.
1194 Each rank blocks until every rank in the group enters this collective (when ``async_op``
1195 is False), or returns an async handle that must be completed before proceeding.
1197 Args:
1198 group: The process group or communication group. ``None`` uses the default group.
1199 async_op (bool): If True, returns a backend-specific async work handle. Default: False.
1200 device_ids: Optional device id list; semantics depend on the backend.
1202 Returns:
1203 Async work handle when ``async_op`` is True; otherwise ``None`` (unless the rank
1204 is not in the group, in which case the backend may return ``None``).
1205 """
1206 raise NotImplementedError("Platform subclasses must implement barrier")
1208 @staticmethod
1209 def init_process_group(
1210 backend: Optional[str] = None,
1211 *,
1212 init_method: Optional[str] = None,
1213 timeout: Optional[timedelta] = None,
1214 world_size: int = -1,
1215 rank: int = -1,
1216 store: Any = None,
1217 pg_options: Any = None,
1218 device_id: Any = None
1219 ) -> None:
1220 """
1221 Initialize the default distributed process group.
1223 Args:
1224 backend: The backend to use for distributed communication
1225 init_method: URL specifying how to initialize the process group
1226 timeout: Timeout for operations executed against the process group
1227 world_size: Number of processes participating in the job
1228 rank: Rank of the current process
1229 store: Key/value store for exchanging connection information
1230 pg_options: Process group options for backend-specific configurations
1231 device_id: Specific device this process will work on
1233 Raises:
1234 NotImplementedError: This method must be implemented by subclasses
1235 """
1236 raise NotImplementedError("Platform subclasses must implement init_process_group")
1238 @staticmethod
1239 def destroy_process_group(group=None) -> None:
1240 """
1241 Destroy a given process group.
1243 Args:
1244 group: The process group to be destroyed. If None, destroys the default group.
1246 Raises:
1247 NotImplementedError: This method must be implemented by subclasses
1248 """
1249 raise NotImplementedError("Platform subclasses must implement destroy_process_group")
1251 @staticmethod
1252 def get_process_group_ranks(group=None) -> list[int]:
1253 """
1254 Get rank list of the given process group.
1256 Args:
1257 group: The process group to get ranks from. If None, uses the default group.
1259 Returns:
1260 List of ranks in the specified process group.
1262 Raises:
1263 NotImplementedError: This method must be implemented by subclasses
1264 """
1265 raise NotImplementedError("Platform subclasses must implement get_process_group_ranks")
1267 @staticmethod
1268 def get_backend(group=None):
1269 """
1270 Get the backend of the given process group.
1271 Args:
1272 group: The process group to get backend from. If None, uses the default group.
1274 Returns:
1275 The backend name of the specified process group.
1277 Raises:
1278 NotImplementedError: This method must be implemented by subclasses
1279 """
1280 raise NotImplementedError("Platform subclasses must implement get_backend")
1282 @staticmethod
1283 def split_group(parent_pg: Any = None,
1284 split_ranks: Optional[list] = None,
1285 timeout: Optional[timedelta] = None,
1286 pg_options: Optional[Any] = None,
1287 group_desc: Optional[str] = None,
1288 ) -> Any:
1289 """Create a split group relative to the parent process group.
1291 Args:
1292 parent_pg: The parent process group to split from.
1293 split_ranks (list, optional): Ranks to include in the split group.
1294 timeout (timedelta, optional): Timeout for operations.
1295 pg_options: Process group options for backend-specific configurations.
1296 group_desc (str, optional): Description of the group.
1298 Returns:
1299 The new split process group.
1300 """
1301 raise NotImplementedError("Platform subclasses must implement split_group")
1303 @staticmethod
1304 def get_group_local_rank(group=None) -> int:
1305 """Get the local rank within the given process group.
1307 Args:
1308 group: The process group to query. If None, uses the default group.
1310 Returns:
1311 int: The local rank within the group.
1312 """
1313 raise NotImplementedError("Platform subclasses must implement get_group_local_rank")
1315 @staticmethod
1316 def no_grad():
1317 """Get a context manager to disable gradient computation.
1319 Returns:
1320 A context manager that disables gradient tracking.
1321 """
1322 raise NotImplementedError("Platform subclasses must implement no_grad")
1324 @staticmethod
1325 def preserve_version_counter(tensor):
1326 """Get a context manager that preserves version for an internal tensor update."""
1327 raise NotImplementedError("Platform subclasses must implement preserve_version_counter")
1329 @staticmethod
1330 def relu(tensor):
1331 """Apply ReLU activation element-wise.
1333 Args:
1334 tensor: Input tensor.
1336 Returns:
1337 Tensor with ReLU applied (max(0, x)).
1338 """
1339 raise NotImplementedError("Platform subclasses must implement relu")
1341 @staticmethod
1342 def cat(tensors, dim=0):
1343 """Concatenate tensors along a dimension."""
1344 raise NotImplementedError("Platform subclasses must implement cat")
1346 @staticmethod
1347 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1348 """Create an uninitialized tensor with the same shape as input.
1350 Args:
1351 tensor: The input tensor to copy shape from.
1352 dtype: Optional dtype for the new tensor. If None, uses input tensor's dtype.
1353 device: Optional device for the new tensor. If None, uses input tensor's device.
1354 pin_memory (bool): If True, allocate pinned memory for faster CPU-GPU transfer.
1356 Returns:
1357 An uninitialized tensor with the same shape as input.
1358 """
1359 raise NotImplementedError("Platform subclasses must implement empty_like")
1361 def get_current_stream(self):
1362 """Get the current compute stream for the device.
1364 Returns:
1365 The current stream object.
1366 """
1367 raise NotImplementedError("Platform subclasses must implement get_current_stream")
1369 def new_event(self):
1370 """Create a new event for stream synchronization.
1372 Returns:
1373 A new event object.
1374 """
1375 raise NotImplementedError("Platform subclasses must implement new_event")
1377 def tree_map(self, fn, tree):
1378 """Apply a function to all tensors in a nested structure.
1380 Args:
1381 fn (callable): Function to apply to each tensor.
1382 tree: Nested structure (list, tuple, dict) containing tensors.
1384 Returns:
1385 The same nested structure with fn applied to all tensors.
1386 """
1387 raise NotImplementedError("Platform subclasses must implement tree_map")
1389 @staticmethod
1390 def is_linear_module(module) -> bool:
1391 """Check whether *module* is a linear/dense layer for the current framework.
1393 Args:
1394 module: The module instance to check.
1396 Returns:
1397 True if *module* is the framework's linear layer type.
1398 """
1399 raise NotImplementedError("Platform subclasses must implement is_linear_module")
1401 @staticmethod
1402 def is_embedding_module(module) -> bool:
1403 """Check whether *module* is an embedding layer for the current framework.
1405 Args:
1406 module: The module instance to check.
1408 Returns:
1409 True if *module* is the framework's embedding layer type.
1410 """
1411 raise NotImplementedError("Platform subclasses must implement is_embedding_module")
1413 @staticmethod
1414 def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
1415 """Register a forward pre-hook on a module.
1417 Args:
1418 module: The module to register the hook on.
1419 hook (callable): The hook function to register.
1420 prepend (bool): If True, prepend the hook to existing hooks.
1421 with_kwargs (bool): If True, hook receives both args and kwargs.
1423 Returns:
1424 A handle that can be used to remove the hook.
1425 """
1426 return module.register_forward_pre_hook(hook, prepend=prepend, with_kwargs=with_kwargs)
1428 @staticmethod
1429 def register_full_backward_hook(module, hook, prepend=False):
1430 """Register a full backward hook on a module.
1432 Args:
1433 module: The module to register the hook on.
1434 hook (callable): The hook function to register.
1435 prepend (bool): If True, prepend the hook to existing hooks.
1437 Returns:
1438 A handle that can be used to remove the hook.
1439 """
1440 return module.register_full_backward_hook(hook, prepend)
1442 @staticmethod
1443 def register_full_backward_pre_hook(module, hook, prepend=False):
1444 """Register a full backward pre-hook on a module.
1446 Args:
1447 module: The module to register the hook on.
1448 hook (callable): The hook function to register.
1449 prepend (bool): If True, prepend the hook to existing hooks.
1451 Returns:
1452 A handle that can be used to remove the hook.
1453 """
1454 return module.register_full_backward_pre_hook(hook, prepend)
1456 @property
1457 def checkpoint(self):
1458 """Get the checkpoint function for activation checkpointing.
1460 Returns:
1461 The checkpoint function for the current framework.
1462 """
1463 raise NotImplementedError("Platform subclasses must implement checkpoint")
1465 @staticmethod
1466 def checkpoint_wrapper(module, **checkpoint_kwargs):
1467 """Wrap a module with activation checkpointing functionality.
1469 Args:
1470 module: The module or callable to wrap with activation checkpointing.
1471 **checkpoint_kwargs: Keyword arguments forwarded to the framework
1472 checkpoint wrapper implementation.
1474 Returns:
1475 The wrapped module with activation checkpointing enabled.
1476 """
1477 raise NotImplementedError("Platform subclasses must implement checkpoint_wrapper")
1479 @staticmethod
1480 def checkpoint_exclude_wrapper(module: Any, *, save_output: bool = True) -> Any:
1481 """Wrap a callable whose activations should be saved instead of recomputed.
1483 Args:
1484 module: The module or callable to exclude from activation recomputation.
1485 save_output: Whether to retain the excluded region output for replay.
1487 Returns:
1488 The wrapped module or callable.
1489 """
1490 raise NotImplementedError("Platform subclasses must implement checkpoint_exclude_wrapper")
1492 @staticmethod
1493 def swap_wrapper(module, policy_fn=None, group_swap=False):
1494 """Wrap a module with activation swap functionality.
1496 Args:
1497 module: The module to wrap with activation swap.
1498 policy_fn: Optional per-tensor swap policy function.
1499 group_swap (bool, optional): Whether tensors participate in group copy fusion. Default: ``False``.
1501 Returns:
1502 The wrapped module with activation swap enabled.
1503 """
1504 raise NotImplementedError("Platform subclasses must implement swap_wrapper")
1506 @staticmethod
1507 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1508 """Register target tensors into the current swap group.
1510 Args:
1511 target: A tensor or nested container of tensors to register.
1512 tag: Optional debug tag associated with the wrapped tensors.
1513 group_swap (bool, optional): Whether tensors participate in group copy fusion. Default: ``False``.
1515 Returns:
1516 The original target structure, unchanged semantically.
1517 """
1518 raise NotImplementedError("Platform subclasses must implement swap_tensor_wrapper")
1520 @staticmethod
1521 def get_class_activation_wrapper():
1522 """Return the platform-specific activation wrapper class."""
1523 raise NotImplementedError("Platform subclasses must implement get_class_activation_wrapper")
1525 @property
1526 def noop_context_fn(self):
1527 """Get a no-op context function for checkpointing.
1529 Returns:
1530 A context function that performs no operation.
1531 """
1532 raise NotImplementedError("Platform subclasses must implement noop_context_fn")
1534 @staticmethod
1535 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1536 """Create contexts for selective activation checkpointing.
1538 Args:
1539 policy_fn_or_list: A policy function or list of layer names to checkpoint.
1540 allow_cache_entry_mutation (bool): Whether to allow cache entry mutation.
1541 group_swap (bool, optional): Whether MUST_SWAP tensors participate in group copy fusion. Default: ``False``.
1543 Returns:
1544 Context functions for selective checkpointing.
1545 """
1546 raise NotImplementedError("Platform subclasses must implement create_selective_checkpoint_contexts")
1548 @staticmethod
1549 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1550 """Create an async CPU offload context for activation checkpointing.
1552 Args:
1553 policy_fn: Optional policy function to determine which activations to offload.
1554 group_swap (bool): Whether swapped tensors participate in group copy fusion.
1555 Default: ``False``.
1557 Returns:
1558 Context manager for async CPU offloading during checkpointing.
1559 """
1560 raise NotImplementedError("Platform subclasses must implement async_save_on_cpu")
1562 @staticmethod
1563 def recompute_handle_collector_ctx():
1564 """Context manager that collects recompute handles created in its scope.
1566 Yields:
1567 A list populated with one opaque recompute handle per checkpointed
1568 block executed during the forward pass within the context. Each
1569 handle can later be fired via :meth:`recompute_handle`.
1570 """
1571 raise NotImplementedError("Platform subclasses must implement recompute_handle_collector_ctx")
1573 @staticmethod
1574 def recompute_handle(handle, session_id):
1575 """Eagerly fire one checkpointed block's forward re-run.
1577 Materializes and caches the block's activations under ``session_id`` so
1578 a later backward in the same session reuses them instead of re-running.
1580 Args:
1581 handle: An opaque recompute handle from
1582 :meth:`recompute_handle_collector_ctx`.
1583 session_id: Stable key shared by the producing re-run and the
1584 consuming backward.
1585 """
1586 raise NotImplementedError("Platform subclasses must implement recompute_handle")
1588 @staticmethod
1589 def recompute_session_ctx(session_id, retain_on_unpack=False):
1590 """Context manager binding recompute unpack to a caller-provided session.
1592 Args:
1593 session_id: Required stable session key. Recompute caches are keyed
1594 by this instead of the transient autodiff engine id, so a re-run
1595 fired under one engine can be reused by another. Must not be
1596 ``None``.
1597 retain_on_unpack (bool): When ``True``, unpack returns recomputed
1598 tensors without popping them, so a later backward can consume
1599 them. Default: ``False``.
1601 Returns:
1602 A context manager activating the session for its scope.
1604 Yields:
1605 The supplied session id.
1606 """
1607 raise NotImplementedError("Platform subclasses must implement recompute_session_ctx")
1609 @staticmethod
1610 def clear_recompute_session(session_id):
1611 """Release retained recompute data for a session.
1613 Args:
1614 session_id: The session key whose cached recompute data is cleared.
1615 """
1616 raise NotImplementedError("Platform subclasses must implement clear_recompute_session")
1618 @staticmethod
1619 def get_element_size(tensor):
1620 """Get Tensor Element Size"""
1621 raise NotImplementedError("Platform subclasses must implement get_element_size")
1623 @staticmethod
1624 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1625 """Allocate an uninitialized 1-D tensor buffer."""
1626 raise NotImplementedError("Platform subclasses must implement alloc_tensor_buffer")
1628 @staticmethod
1629 def tensor_to_numpy(tensor) -> np.ndarray:
1630 """Convert a framework tensor to a NumPy array.
1632 Args:
1633 tensor: The tensor to convert.
1635 Returns:
1636 np.ndarray: The tensor data as a NumPy array.
1637 """
1638 raise NotImplementedError("Platform subclasses must implement tensor_to_numpy")
1640 @staticmethod
1641 def from_numpy(np_array):
1642 """Create a host-resident tensor from a NumPy array (inverse of tensor_to_numpy).
1644 The result stays on the host regardless of the active device context, so it
1645 remains asnumpy-able even when built under ``ms.DeviceCtx("meta")`` (e.g. while
1646 ``fully_shard`` lazily constructs a default device mesh). Use it for rank/mesh
1647 bookkeeping tensors, which are only ever read back via ``tensor_to_numpy``.
1648 """
1649 raise NotImplementedError("Platform subclasses must implement from_numpy")
1651 @staticmethod
1652 def profiler_record(name):
1653 """Record a profiler event with the given name.
1655 Args:
1656 name (str): The name of the profiler event.
1658 Returns:
1659 A context manager or decorator for profiling a code region.
1660 """
1661 raise NotImplementedError("Platform subclasses must implement profiler_record")
1663 def cast_fp_tensor(self, dtype, x):
1664 """Cast floating-point tensor to target dtype if applicable.
1666 Args:
1667 dtype: The target dtype to cast to.
1668 x: The input tensor.
1670 Returns:
1671 The tensor cast to target dtype, or unchanged if not floating-point.
1672 """
1673 raise NotImplementedError("Platform subclasses must implement cast_fp_tensor")
1675 def apply_to_tensors(self, fn, container):
1676 """Recursively apply a function to all tensors in a container.
1678 Supports nested structures including lists, tuples, and dicts.
1680 Args:
1681 fn (callable): Function to apply to each tensor.
1682 container: Nested structure containing tensors.
1684 Returns:
1685 The same structure with fn applied to all tensors.
1686 """
1687 raise NotImplementedError("Platform subclasses must implement apply_to_tensors")
1689 @staticmethod
1690 def clip_grad_norm_(
1691 parameters, max_norm: float, norm_type: float = 2.0,
1692 error_if_nonfinite: bool = False, foreach=None,
1693 ):
1694 """Compute and clip gradient norms for distributed models.
1696 Communication is derived from each parameter's DTensor spec.
1697 Subclasses must implement this method.
1699 Args:
1700 parameters: An ``nn.Module``, a single ``Tensor``, or an
1701 iterable of ``Tensor`` s whose gradients to clip.
1702 max_norm: Maximum allowed gradient norm.
1703 norm_type: Type of the norm (default ``2.0``).
1704 error_if_nonfinite: If ``True``, raise when total norm is
1705 non-finite. Default ``False``.
1706 foreach: Unused, accepted for API compatibility.
1708 Returns:
1709 The total (unclipped) gradient norm.
1710 """
1711 raise NotImplementedError(
1712 "Platform subclasses must implement clip_grad_norm_"
1713 )
1715 @staticmethod
1716 def get_created_group(rank_list: Union[list[int], tuple[int]]):
1717 """Get an existing process group by rank list.
1719 Args:
1720 rank_list (Union[list[int], tuple[int]]): Tuple or list of ranks.
1722 Returns:
1723 The process group corresponding to the rank list if it exists, else None.
1724 """
1725 group_key = str(tuple(sorted(rank_list)))
1726 if group_key in EXISTING_COMM_GROUPS:
1727 return EXISTING_COMM_GROUPS[group_key]
1728 return None
1730 @classmethod
1731 def mark_created_groups(cls, process_group: Union[Any, list[Any]]) -> None:
1732 """Register process groups in the global cache for reuse.
1734 Args:
1735 process_group (Union[Any, list[Any]]): A process group or a list of process groups.
1736 """
1737 if not isinstance(process_group, list):
1738 process_group = [process_group]
1739 for group in process_group:
1740 rank_list = cls.get_process_group_ranks(group)
1741 group_key = str(tuple(sorted(rank_list)))
1742 EXISTING_COMM_GROUPS[group_key] = group
1744 @property
1745 def meta_device(self):
1746 """Get the framework-specific meta device for tensor shape inference.
1748 The meta device allows creating tensors without allocating actual storage,
1749 useful for shape inference and model initialization.
1751 Returns:
1752 The meta device object for the current framework.
1753 """
1754 raise NotImplementedError("Platform subclasses must implement meta_device")
1756 def init_on_device(self, device, include_buffers=False):
1757 """Get a context manager for initializing module parameters on a device.
1759 Args:
1760 device: The target device for parameter initialization.
1761 include_buffers (bool): If True, also initialize buffers on the device.
1763 Returns:
1764 A context manager for device-specific initialization.
1765 """
1766 raise NotImplementedError("Platform subclasses must implement init_on_device")
1768 def str_to_dtype(self, dtype_str: str) -> Any:
1769 """
1770 Map a framework-style dtype string (e.g. ``torch.float32``) to the backend dtype object.
1772 Args:
1773 dtype_str (str): Serialized dtype identifier produced by checkpoint metadata.
1775 Returns:
1776 Framework dtype object (e.g. ``torch.dtype`` or MindSpore dtype).
1777 """
1778 raise NotImplementedError("Platform subclasses must implement str_to_dtype")
1780 def list_to_size(self, size_list: list[int]) -> Any:
1781 """
1782 Convert a shape list from checkpoint metadata to the framework's size type (e.g. ``torch.Size``).
1784 Args:
1785 size_list (list[int]): Tensor global shape as a list of ints.
1787 Returns:
1788 Framework-specific size object.
1789 """
1790 raise NotImplementedError("Platform subclasses must implement list_to_size")