Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / platform.py: 66%
428 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""framework platform api"""
16import os
17from datetime import timedelta
18from enum import auto, Enum
19from typing import Optional, Any, Union
21import numpy as np
23# Environment variable name used to specify the AI framework platform to use
24HYPER_PARALLEL_PLATFORM = "HYPER_PARALLEL_PLATFORM"
26# Identifier for the MindSpore framework
27HYPER_PARALLEL_PLATFORM_MINDSPORE = "mindspore"
29# Identifier for the PyTorch framework
30HYPER_PARALLEL_PLATFORM_TORCH = "torch"
33class AsyncHandle:
34 """Idempotent wait handle for an async collective operation.
36 Wraps the async tensor returned by
37 :meth:`Platform.differentiable_all_to_all_single_async` and provides a
38 :meth:`wait` method that is safe to call multiple times.
39 """
41 def __init__(self, async_tensor) -> None:
42 self._tensor = async_tensor
43 self._waited = False
45 def wait(self):
46 """Wait for the async collective to complete.
48 Idempotent — the first call blocks until the collective finishes;
49 subsequent calls are no-ops.
51 Returns:
52 The now-materialised result tensor.
53 """
54 if not self._waited:
55 get_platform().wait_async_tensor(self._tensor)
56 self._waited = True
57 return self._tensor
60class PlatformType(Enum):
61 """Enumeration class for AI framework platform types.
63 Used to identify different deep learning framework platform types.
64 """
65 MINDSPORE = auto()
66 PYTORCH = auto()
69# Global platform instance, used to cache the created platform object
70platform = None
73def get_mindspore_platform():
74 """Create and return a MindSpore platform instance.
76 Returns:
77 MindSporePlatform: A MindSpore platform instance.
78 """
79 # pylint: disable=C0415
80 from hyper_parallel.platform.mindspore.platform import MindSporePlatform
81 global platform
82 platform = MindSporePlatform()
83 return platform
86def get_torch_platform():
87 """Create and return a PyTorch platform instance.
89 Returns:
90 TorchPlatform: A PyTorch platform instance.
91 """
92 # pylint: disable=C0415
93 from hyper_parallel.platform.torch.platform import TorchPlatform
94 global platform
95 platform = TorchPlatform()
96 return platform
99def get_platform():
100 """Obtain a framework platform instance.
102 Returns the appropriate AI framework platform instance based on environment variables or a default priority order.
103 The lookup priority is as follows:
104 1. Platform specified by environment variable
105 2. MindSpore platform (default preferred choice)
106 3. PyTorch platform (fallback option)
108 Returns:
109 Platform: An instance of the framework platform
111 Raises:
112 ImportError: Raised when none of the supported frameworks are available
113 """
114 if platform is not None:
115 return platform
116 platform_type = os.environ.get(HYPER_PARALLEL_PLATFORM)
117 if platform_type is not None and isinstance(platform_type, str):
118 platform_type = platform_type.lower()
119 if platform_type == HYPER_PARALLEL_PLATFORM_MINDSPORE:
120 return get_mindspore_platform()
121 if platform_type == HYPER_PARALLEL_PLATFORM_TORCH:
122 return get_torch_platform()
123 try:
124 return get_mindspore_platform()
125 except ImportError:
126 return get_torch_platform()
129EXISTING_COMM_GROUPS = {}
132class Platform:
133 """Platform api"""
134 current_grad_handle = None
135 post_grad_handle_process = None
136 grad_sync_stream = None
138 @property
139 def custom_ops(self):
140 """Return the platform-specific custom ops interface.
142 Subclasses MUST override this property to return an object that
143 exposes the platform-specific custom operator implementations.
145 Returns:
146 object: Platform-specific custom ops class instance.
147 """
148 raise NotImplementedError(
149 "Platform subclasses must implement custom_ops"
150 )
152 @staticmethod
153 def get_rank():
154 """Get the rank of the current process in the default process group.
156 Returns:
157 int: The rank of the current process.
158 """
159 raise NotImplementedError("Platform subclasses must implement get_rank")
161 @staticmethod
162 def get_global_rank(group, group_rank):
163 """Convert a group rank to its global rank.
165 Args:
166 group: The process group to query.
167 group_rank (int): The rank within the group.
169 Returns:
170 int: The global rank corresponding to the group rank.
171 """
172 raise NotImplementedError("Platform subclasses must implement get_global_rank")
174 @staticmethod
175 def get_group_rank(group):
176 """Return this process's rank within *group*."""
177 raise NotImplementedError("Platform subclasses must implement get_group_rank")
179 @staticmethod
180 def get_world_size():
181 """Get the total number of processes in the default process group.
183 Returns:
184 int: The world size (total number of processes).
185 """
186 raise NotImplementedError("Platform subclasses must implement get_world_size")
188 @staticmethod
189 def get_op_name(func):
190 """Get the canonical name of an operator function.
192 Args:
193 func: The operator function to query.
195 Returns:
196 str: The canonical name of the operator.
197 """
198 raise NotImplementedError("Platform subclasses must implement get_op_name")
200 @staticmethod
201 def differentiable_all_gather_concat(data, group, concat_size, concat_dim, rank_list=None):
202 """Perform differentiable all-gather and concatenate tensors along a dimension.
204 Args:
205 data: The input tensor to gather.
206 group: The process group for collective communication.
207 concat_size (int): The size to concatenate along concat_dim.
208 concat_dim (int): The dimension along which to concatenate.
209 rank_list: Optional rank order expected by the logical layout.
211 Returns:
212 The concatenated tensor after all-gather operation.
213 """
214 raise NotImplementedError("Platform subclasses must implement differentiable_all_gather_concat")
216 @staticmethod
217 def chunk(data, split_dim, split_size, index):
218 """Split tensor along a dimension and return the chunk at the given index.
220 Args:
221 data: The input tensor to split.
222 split_dim (int): The dimension along which to split.
223 split_size (int): The size of each split chunk.
224 index (int): The index of the chunk to return.
226 Returns:
227 The tensor chunk at the specified index.
228 """
229 raise NotImplementedError("Platform subclasses must implement chunk")
231 @staticmethod
232 def differentiable_all_to_all(input_data, output_shape, group):
233 """Perform differentiable all-to-all communication.
235 Args:
236 input_data: The input tensor to redistribute.
237 output_shape: The shape of the output tensor.
238 group: The process group for collective communication.
240 Returns:
241 The output tensor after all-to-all operation.
242 """
243 raise NotImplementedError("Platform subclasses must implement differentiable_all_to_all")
245 @staticmethod
246 def tensor_type_cast(input_data, cast_type):
247 """Cast tensor to a specified dtype.
249 Args:
250 input_data: The input tensor to cast.
251 cast_type: The target dtype to cast to.
253 Returns:
254 The tensor cast to the specified dtype.
255 """
256 raise NotImplementedError("Platform subclasses must implement tensor_type_cast")
258 @staticmethod
259 def is_tensor(obj: Any) -> bool:
260 """Return True if ``obj`` is this framework's tensor type."""
261 raise NotImplementedError("Platform subclasses must implement is_tensor")
263 @staticmethod
264 def get_tensor_storage_size(tensor: Any) -> int:
265 """Return serialized byte size (numel * element size) for this framework's tensor."""
266 raise NotImplementedError("Platform subclasses must implement get_tensor_storage_size")
268 @staticmethod
269 def differentiable_all_reduce(data, op, group):
270 """Perform differentiable all-reduce operation.
272 Args:
273 data: The input tensor to reduce.
274 op: The reduction operation (e.g., sum, max, min).
275 group: The process group for collective communication.
277 Returns:
278 The reduced tensor with gradients supported.
279 """
280 raise NotImplementedError("Platform subclasses must implement differentiable_all_reduce")
282 @staticmethod
283 def differentiable_reduce_scatter(data, dev_num, axis, op, group):
284 """Perform differentiable reduce-scatter operation.
286 Args:
287 data: The input tensor to reduce and scatter.
288 dev_num (int): The number of devices to scatter across.
289 axis (int): The axis along which to scatter.
290 op: The reduction operation (e.g., sum, max, min).
291 group: The process group for collective communication.
293 Returns:
294 The scattered tensor chunk with gradients supported.
295 """
296 raise NotImplementedError("Platform subclasses must implement differentiable_reduce_scatter")
298 @staticmethod
299 def init_parameters(module, stage_index):
300 """Initialize parameters for a module at a specific pipeline stage.
302 This method is primarily needed for MindSpore platform which requires
303 explicit parameter initialization interface.
305 Args:
306 module: The module whose parameters need to be initialized.
307 stage_index (int): The pipeline stage index for the module.
309 Raises:
310 ValueError: If module is None or stage_index is negative.
311 """
312 if module is None:
313 raise ValueError("input module must not be none.")
314 if stage_index < 0:
315 raise ValueError("input stage_index must be positive.")
317 @staticmethod
318 def get_cell_construct(cell):
319 """Get the construct (forward) function of a cell/module.
321 Args:
322 cell: The cell or module to get the construct function from.
324 Returns:
325 The construct/forward callable of the cell.
326 """
327 raise NotImplementedError("Platform subclasses must implement get_cell_construct")
329 @staticmethod
330 def get_cells_and_names(cell):
331 """Get all nested cells/modules and their names.
333 Args:
334 cell: The root cell or module to traverse.
336 Returns:
337 list: A list of tuples containing (name, cell) pairs.
338 """
339 raise NotImplementedError("Platform subclasses must implement get_cells_and_names")
341 @staticmethod
342 def get_modules(module):
343 """Return all sub-modules contained in the given module."""
344 raise NotImplementedError("Platform subclasses must implement get_modules")
346 @staticmethod
347 def search_parameter_by_name(cell, param_name: str):
348 """Search for a parameter by name within a cell/module.
350 Args:
351 cell: The cell or module to search in.
352 param_name (str): The name of the parameter to find.
354 Returns:
355 The parameter if found, otherwise None.
356 """
357 raise NotImplementedError("Platform subclasses must implement search_parameter_by_name")
359 @staticmethod
360 def update_parameter_by_name(cell, result: tuple, new_param) -> bool:
361 """Update a parameter by name within a cell/module.
363 Args:
364 cell: The cell or module containing the parameter.
365 result (tuple): A tuple containing (param_name, parameter) to update.
366 new_param: The new parameter value to set.
368 Returns:
369 bool: True if update was successful, False otherwise.
370 """
371 raise NotImplementedError("Platform subclasses must implement update_parameter_by_name")
373 @staticmethod
374 def set_layout_into_parameter(param, layout):
375 """Attach a DTensor layout to a parameter.
377 Args:
378 param: The parameter to attach the layout to.
379 layout: The DTensor layout describing tensor distribution.
380 """
381 raise NotImplementedError("Platform subclasses must implement set_layout_into_parameter")
383 @staticmethod
384 def get_param_local_shape(param):
385 """Get the local shape of a distributed parameter.
387 Args:
388 param: The parameter to query.
390 Returns:
391 tuple: The local shape of the parameter shard.
392 """
393 raise NotImplementedError("Platform subclasses must implement get_param_local_shape")
395 @staticmethod
396 def get_param_local_data(param):
397 """Get the local data tensor of a distributed parameter.
399 Args:
400 param: The parameter to query.
402 Returns:
403 The local tensor data of the parameter shard.
404 """
405 raise NotImplementedError("Platform subclasses must implement get_param_local_data")
407 @staticmethod
408 def update_param_data(param, data):
409 """Update the data of a parameter with new tensor data.
411 Args:
412 param: The parameter to update.
413 data: The new tensor data to assign.
414 """
415 raise NotImplementedError("Platform subclasses must implement update_param_data")
417 @staticmethod
418 def get_param_type_size(param):
419 """Get the size in bytes of a parameter's dtype.
421 Args:
422 param: The parameter to query.
424 Returns:
425 int: The size in bytes of the parameter's data type.
426 """
427 raise NotImplementedError("Platform subclasses must implement get_param_type_size")
429 @staticmethod
430 def new_zero_parameter(param_shape, param_type, requires_grad, device):
431 """Create a new parameter initialized with zeros.
433 Args:
434 param_shape (tuple): The shape of the parameter.
435 param_type: The dtype of the parameter.
436 requires_grad (bool): Whether the parameter requires gradients.
437 device: The device on which to create the parameter.
439 Returns:
440 A new parameter tensor filled with zeros.
441 """
442 raise NotImplementedError("Platform subclasses must implement new_zero_parameter")
444 @staticmethod
445 def new_tensor(tensor_shape, tensor_type, device):
446 """Create a new tensor with the specified shape, dtype, and device.
448 Args:
449 tensor_shape (tuple): The shape of the tensor.
450 tensor_type: The dtype of the tensor.
451 device: The device on which to create the tensor.
453 Returns:
454 A new tensor with uninitialized values.
455 """
456 raise NotImplementedError("Platform subclasses must implement new_tensor")
458 @staticmethod
459 def full_like(tensor, fill_value, dtype=None):
460 """Create a tensor filled with a value, with same shape as input.
462 Args:
463 tensor: The input tensor to copy shape from.
464 fill_value: The value to fill the new tensor with.
465 dtype: Optional dtype for the new tensor. If None, uses input tensor's dtype.
467 Returns:
468 A new tensor filled with the specified value.
469 """
470 raise NotImplementedError("Platform subclasses must implement full_like")
472 @staticmethod
473 def set_tensor_requires_grad(input_tensor):
474 """Enable gradient tracking for a tensor in-place.
476 Args:
477 input_tensor: The tensor to enable gradients for.
479 Returns:
480 The same tensor with requires_grad set to True.
481 """
482 raise NotImplementedError("Platform subclasses must implement set_tensor_requires_grad")
484 @staticmethod
485 def all_gather_into_tensor(data, group_info, async_op=False):
486 """Gather tensors from all ranks into a single output tensor.
488 Args:
489 data: The input tensor to gather.
490 group_info: The process group for collective communication.
491 async_op (bool): If True, returns a work handle for async operation.
493 Returns:
494 The gathered tensor, or a tuple of (tensor, handle) if async_op is True.
495 """
496 raise NotImplementedError("Platform subclasses must implement all_gather_into_tensor")
498 @staticmethod
499 def all_reduce(data, group_info, async_op=False):
500 """Reduce tensors across all ranks using specified operation.
502 Args:
503 data: The input tensor to reduce.
504 group_info: The process group for collective communication.
505 async_op (bool): If True, returns a work handle for async operation.
507 Returns:
508 The reduced tensor, or a tuple of (tensor, handle) if async_op is True.
509 """
510 raise NotImplementedError("Platform subclasses must implement all_reduce")
512 @staticmethod
513 def broadcast(data, src=None, group=None, async_op=False, group_src=None):
514 """Broadcast tensor from source rank to all ranks in group."""
515 raise NotImplementedError("Platform subclasses must implement broadcast")
517 @staticmethod
518 def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
519 """Scatter tensor list from source rank to all ranks in group."""
520 raise NotImplementedError("Platform subclasses must implement scatter")
522 @staticmethod
523 def isend(tensor, dst=None, group=None, tag=0):
524 """Send tensor asynchronously to destination rank.
526 Args:
527 tensor: The tensor to send.
528 dst (int, optional): The destination rank. Defaults to None.
529 group: The process group for communication. Defaults to None.
530 tag (int): A tag to identify the send operation. Defaults to 0.
532 Returns:
533 A work handle that can be waited on.
534 """
535 raise NotImplementedError("Platform subclasses must implement isend")
537 @staticmethod
538 def irecv(tensor, src=None, group=None, tag=0):
539 """Receive tensor asynchronously from source rank.
541 Args:
542 tensor: The tensor buffer to receive data into.
543 src (int, optional): The source rank. Defaults to None.
544 group: The process group for communication. Defaults to None.
545 tag (int): A tag to identify the receive operation. Defaults to 0.
547 Returns:
548 A work handle that can be waited on.
549 """
550 raise NotImplementedError("Platform subclasses must implement irecv")
552 @staticmethod
553 def p2p_op(op_type, tensor, peer, group=None):
554 """Build a batched-P2P descriptor (no launch).
556 Returns an opaque object understood by :meth:`batch_isend_irecv`.
557 Lets callers assemble a mixed send/recv batch that the backend can
558 run concurrently (e.g. TX/RX duplex on one link) in a single op.
560 Args:
561 op_type (str): ``"isend"`` or ``"irecv"``.
562 tensor: Tensor to send, or the buffer to receive into.
563 peer (int): Global rank of the peer.
564 group: Process group. ``None`` uses the default group.
566 Returns:
567 A backend P2P-op descriptor.
568 """
569 raise NotImplementedError("Platform subclasses must implement p2p_op")
571 @staticmethod
572 def batch_isend_irecv(p2p_ops):
573 """Launch a batch of :meth:`p2p_op` descriptors as one async op.
575 The whole batch shares a single completion handle (the backend runs
576 the items concurrently on one comm stream), so a send and a recv to
577 the same peer overlap on the duplex link.
579 Args:
580 p2p_ops (list): Descriptors from :meth:`p2p_op`.
582 Returns:
583 A single work handle covering the whole batch, or ``None`` when
584 ``p2p_ops`` is empty.
585 """
586 raise NotImplementedError("Platform subclasses must implement batch_isend_irecv")
588 @staticmethod
589 def prepare_batch_p2p_group(group: Any = None) -> None:
590 """Prepare a process group before its first batched P2P operation.
592 Backends that require full-group participation before subset batched
593 P2P should synchronize the group here. Other backends may implement
594 this as a no-op.
596 Args:
597 group: The process group used by the batched P2P operations.
598 ``None`` uses the default group.
599 """
600 raise NotImplementedError("Platform subclasses must implement prepare_batch_p2p_group")
602 @staticmethod
603 def p2p_exchange(tensor, peer_rank: int, group=None):
604 """Differentiable symmetric P2P exchange (send local tensor, receive peer's tensor).
606 Sends ``tensor`` to ``peer_rank`` and simultaneously receives the peer's
607 tensor. The operation is differentiable: the backward pass performs the
608 same symmetric exchange on the upstream gradient.
610 Args:
611 tensor: Local tensor to send.
612 peer_rank (int): Global rank of the communication peer.
613 group: Process group. ``None`` uses the default group.
615 Returns:
616 Tensor received from ``peer_rank``, with the same shape and dtype as
617 the input ``tensor``.
618 """
619 raise NotImplementedError("Platform subclasses must implement p2p_exchange")
621 @staticmethod
622 def send_object_list(obj_list, dst=None, group=None):
623 """Send a list of Python objects to destination rank.
625 Args:
626 obj_list (list): The list of Python objects to send.
627 dst (int, optional): The destination rank. Defaults to None.
628 group: The process group for communication. Defaults to None.
629 """
630 raise NotImplementedError("Platform subclasses must implement send_object_list")
632 @staticmethod
633 def recv_object_list(obj_list, src=None, group=None):
634 """Receive a list of Python objects from source rank.
636 Args:
637 obj_list (list): The list buffer to receive objects into.
638 src (int, optional): The source rank. Defaults to None.
639 group: The process group for communication. Defaults to None.
640 """
641 raise NotImplementedError("Platform subclasses must implement recv_object_list")
643 @staticmethod
644 def reduce_scatter_tensor(data, group_info, async_op=False):
645 """Reduce and scatter tensor across all ranks in group.
647 Args:
648 data: The input tensor to reduce and scatter.
649 group_info: The process group for collective communication.
650 async_op (bool): If True, returns a work handle for async operation.
652 Returns:
653 The scattered tensor chunk, or a tuple of (tensor, handle) if async_op is True.
654 """
655 raise NotImplementedError("Platform subclasses must implement reduce_scatter_tensor")
657 @staticmethod
658 def all_gather_single(input_tensor, output_shape, group, async_op=False):
659 """All-gather tensor shards with optional async execution.
661 Args:
662 input_tensor: Input tensor whose leading dimension is gathered.
663 output_shape: Shape of the gathered output tensor.
664 group: Process group (ProcessGroup for torch, group name string for mindspore).
665 async_op: If True, returns an async work handle.
667 Returns:
668 Tuple ``(output, work)`` where *output* is the gathered tensor and
669 *work* is the async handle (``None`` when ``async_op=False``).
670 """
671 raise NotImplementedError("Platform subclasses must implement all_gather_single")
673 @staticmethod
674 def reduce_scatter_single(input_tensor, output_shape, group, async_op=False):
675 """Reduce-scatter a tensor with optional async execution.
677 Args:
678 input_tensor: Input tensor whose leading dimension is split across ranks.
679 output_shape: Shape of the local reduced output tensor.
680 group: Process group (ProcessGroup for torch, group name string for mindspore).
681 async_op: If True, returns an async work handle.
683 Returns:
684 Tuple ``(output, work)`` where *output* is the local shard and
685 *work* is the async handle (``None`` when ``async_op=False``).
686 """
687 raise NotImplementedError("Platform subclasses must implement reduce_scatter_single")
689 @staticmethod
690 def all_to_all_single(input_tensor, output_shape, group, async_op=False):
691 """All-to-all single collective with optional async execution.
693 Args:
694 input_tensor: Input tensor to scatter.
695 output_shape: Shape of the pre-allocated output tensor.
696 group: Process group (ProcessGroup for torch, group name string for mindspore).
697 async_op: If True, returns a work handle; the output tensor is
698 filled only after ``work.wait()`` is called.
700 Returns:
701 Tuple ``(output, work)`` where *output* is the result tensor and
702 *work* is the async handle (``None`` when ``async_op=False``).
704 Raises:
705 NotImplementedError: Must be implemented by platform subclasses.
706 """
707 raise NotImplementedError("Platform subclasses must implement all_to_all_single")
709 @staticmethod
710 def differentiable_async_allgather_wait(x, work, out_perm, group, world_size, gather_dim,
711 handle_box=None):
712 """Differentiable wrapper that waits for a pre-launched async all-gather.
714 Forward waits for the all-gather handle and reconstructs the tensor by
715 moving the gathered leading dimension back to ``gather_dim``.
717 Backward launches the reverse reduce-scatter. If ``handle_box`` is a
718 mutable list, the reduce-scatter handle is appended there and a zero
719 gradient is returned to be replaced by the caller's backward pre-hook.
720 If ``handle_box`` is ``None``, the reduce-scatter is waited immediately
721 and its local result is returned, preserving composability with an
722 upstream autograd communication op.
724 Args:
725 x: Original input tensor; anchors the op in the autograd graph.
726 work: Async work handle from all-gather.
727 out_perm: Output buffer filled by all-gather.
728 group: Communication group for backward reduce-scatter.
729 world_size: Group size.
730 gather_dim: Dimension gathered in forward.
731 handle_box: Optional mutable list for deferred backward wait.
733 Returns:
734 Gathered tensor connected to the autograd graph through *x*.
735 """
736 raise NotImplementedError("Platform subclasses must implement differentiable_async_allgather_wait")
738 @staticmethod
739 def differentiable_async_a2a_wait(x, work, out_perm, group, world_size, concat_dim, split_dim,
740 handle_box=None):
741 """Differentiable wrapper that waits for a pre-launched async A2A.
743 Wraps the wait-and-reconstruct step in the platform autograd mechanism
744 so gradients flow correctly through the all-to-all communication.
746 The A2A direction is seq→head (forward): the output gathers along
747 ``concat_dim`` (sequence grows from S/cp to S) and scatters along
748 ``split_dim`` (heads shrink from H to H/ws).
750 In backward, launches an async head→seq A2A on the incoming gradient
751 and appends ``(work, out_perm)`` to ``handle_box`` so the caller can
752 wait just before the projection GEMM, achieving GEMM–A2A overlap.
754 Args:
755 x: Original projection output tensor; anchors the op
756 in the autograd graph.
757 work: Async work handle from ``all_to_all_single(async_op=True)``.
758 out_perm: Output buffer filled once ``work.wait()`` completes
759 (shape ``[ws, ...]``).
760 group: Process group for the reverse A2A in backward.
761 world_size: CP/Ulysses degree.
762 concat_dim: Dimension that is gathered (concatenated) in forward;
763 typically the sequence dimension.
764 split_dim: Dimension that is scattered (split) in forward;
765 typically the head dimension.
766 handle_box: Optional mutable list ``[]``. In backward, ``(work, out_perm)``
767 for the reverse A2A is appended here so the pre-hook can wait.
769 Returns:
770 Result tensor with ``concat_dim`` gathered and ``split_dim`` split,
771 connected to the autograd graph through *x*.
773 Raises:
774 NotImplementedError: Must be implemented by platform subclasses.
775 """
776 raise NotImplementedError("Platform subclasses must implement differentiable_async_a2a_wait")
778 @staticmethod
779 def differentiable_sync_hook(x, hook_name: str, coordinator):
780 """Identity operation that intercepts both forward and backward to call
781 coordinator rendezvous, enabling deterministic comm/compute overlap.
783 This is the differentiable building block for dual-pipe schedules.
784 In the forward pass the coordinator is invoked with the forward-side
785 roles for ``hook_name``; in the backward pass it is invoked with the
786 backward-side roles. The tensor value and gradient flow through
787 unchanged.
789 Args:
790 x: Input tensor. Returned as-is; gradients flow through.
791 hook_name: One of ``"A"``, ``"B"``, ``"C"``, ``"D"`` identifying
792 the position relative to MoE dispatch/combine.
793 coordinator: A :class:`HookCoordinator` instance shared between the
794 forward and backward threads.
796 Returns:
797 The same tensor *x*, attached to the autograd graph so that the
798 backward hook will fire.
799 """
800 raise NotImplementedError("Platform subclasses must implement differentiable_sync_hook")
802 @staticmethod
803 def differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group):
804 """Variable-split all-to-all single that supports gradient flow.
806 Unlike ``all_to_all_single`` (which is not differentiable), this method
807 wraps the collective in an autograd function so gradients are correctly
808 routed back through the reverse all-to-all in the backward pass.
809 Intended for Expert Parallelism token dispatch / combine.
811 Args:
812 input_tensor: Input tensor to scatter. Shape ``[sum(input_splits), *feature_dims]``.
813 input_splits: Per-rank sizes of data sent from this rank (list of ints,
814 length equal to ep_degree).
815 output_splits: Per-rank sizes of data received by this rank (list of ints,
816 length equal to ep_degree).
817 group: Process group (ProcessGroup for torch, group name str for mindspore).
819 Returns:
820 Output tensor of shape ``[sum(output_splits), *feature_dims]``.
822 Raises:
823 NotImplementedError: Must be implemented by platform subclasses.
824 """
825 raise NotImplementedError("Platform subclasses must implement differentiable_all_to_all_single")
827 @staticmethod
828 def differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group):
829 """Async variant of :meth:`differentiable_all_to_all_single`.
831 Same semantics but launches the collective with ``async_op=True`` and
832 only performs a stream-level ``wait`` — the host returns immediately
833 after dispatching the kernel. Intended for dual-pipe comm/compute
834 overlap paths where the paired COMPUTE side's rendezvous notify must
835 fire right after kernel launch (not after the collective actually
836 completes on device).
838 Args:
839 input_tensor: Input tensor to scatter. Shape ``[sum(input_splits), *feature_dims]``.
840 input_splits: Per-rank sizes of data sent from this rank.
841 output_splits: Per-rank sizes of data received by this rank.
842 group: Process group.
844 Returns:
845 Output tensor of shape ``[sum(output_splits), *feature_dims]``.
847 Raises:
848 NotImplementedError: Must be implemented by platform subclasses.
849 """
850 raise NotImplementedError(
851 "Platform subclasses must implement differentiable_all_to_all_single_async"
852 )
854 @staticmethod
855 def wait_async_tensor(tensor):
856 """Wait for an async collective tensor to become materialised.
858 Intended for use with :class:`AsyncHandle` so that callers can
859 wait on an async all-to-all result without importing framework-specific
860 modules directly. The call is **idempotent** — waiting on an already-
861 completed tensor is a no-op.
863 Args:
864 tensor: An async collective tensor (e.g. PyTorch
865 ``AsyncCollectiveTensor``) whose values have not yet been
866 fully written by the remote ranks.
868 Returns:
869 The same *tensor*, now guaranteed to be fully materialised.
871 Raises:
872 NotImplementedError: Must be implemented by platform subclasses.
873 """
874 raise NotImplementedError(
875 "Platform subclasses must implement wait_async_tensor"
876 )
878 @staticmethod
879 def arange(start, end=None, step=1, dtype=None, device=None):
880 """Create a 1-D tensor with evenly spaced values.
882 Args:
883 start: Start of interval (inclusive). If *end* is ``None``,
884 treated as the stop value and *start* defaults to 0.
885 end: End of interval (exclusive). Defaults to ``None``.
886 step: Step size. Defaults to ``1``.
887 dtype: Data type. ``None`` uses the framework default (int64).
888 device: Target device.
890 Returns:
891 1-D tensor ``[start, start+step, ..., end)``.
893 Raises:
894 NotImplementedError: Must be implemented by platform subclasses.
895 """
896 raise NotImplementedError("Platform subclasses must implement arange")
898 @staticmethod
899 def zeros(size, dtype=None, device=None):
900 """Create a zero-filled tensor of the given shape.
902 Args:
903 size: Shape of the tensor (a single tuple/list).
904 dtype: Desired data type. ``None`` uses the framework default (float32).
905 device: Target device. ``None`` uses the framework default.
907 Returns:
908 Zero-filled tensor of the specified shape.
910 Raises:
911 NotImplementedError: Must be implemented by platform subclasses.
912 """
913 raise NotImplementedError("Platform subclasses must implement zeros")
915 @staticmethod
916 def parameters_dict(cell):
917 """Get the parameters dictionary of a cell/module.
919 Args:
920 cell: The cell or module to get parameters from.
922 Returns:
923 dict: A dictionary mapping parameter names to parameters.
924 """
925 raise NotImplementedError("Platform subclasses must implement parameters_dict")
927 @staticmethod
928 def buffers_dict(cell: Any) -> Any:
929 """Get the named buffers of a cell/module.
931 Args:
932 cell: The cell or module to get buffers from.
934 Returns:
935 An iterable of ``(name, buffer)`` pairs, including non-persistent
936 buffers and buffers registered by child modules.
937 """
938 raise NotImplementedError("Platform subclasses must implement buffers_dict")
940 @staticmethod
941 def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
942 """Get the state dictionary of a model.
944 Args:
945 model: The model to extract state from.
946 options: Optional configuration for state dict extraction.
948 Returns:
949 dict: The state dictionary containing model parameters and buffers.
951 Raises:
952 NotImplementedError: Platform subclasses must implement this method.
953 """
954 raise NotImplementedError(
955 "Platform subclasses must implement get_model_state_dict"
956 )
958 @staticmethod
959 def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
960 """Set the state dictionary of a model.
962 Args:
963 model: The model to load state into.
964 model_state_dict: The state dict to load into the model.
965 options: Optional configuration for state dict loading.
967 Returns:
968 None.
970 Raises:
971 NotImplementedError: Platform subclasses must implement this method.
972 """
973 raise NotImplementedError(
974 "Platform subclasses must implement set_model_state_dict"
975 )
977 @staticmethod
978 def save_checkpoint(cell, file_path: str, ckpt_format: str = "safetensors") -> None:
979 """Save a cell/module checkpoint to file.
981 Args:
982 cell: The cell or module to save.
983 file_path (str): The path to save the checkpoint to.
984 ckpt_format (str): The file format.
985 """
986 raise NotImplementedError("Platform subclasses must implement save_checkpoint")
988 @staticmethod
989 def load_checkpoint(file_path: str, ckpt_format: str = "safetensors") -> dict:
990 """Load a checkpoint from file.
992 Args:
993 file_path (str): The path to load the checkpoint from.
994 ckpt_format (str): The file format.
996 Returns:
997 dict: The loaded checkpoint state dictionary.
998 """
999 raise NotImplementedError("Platform subclasses must implement load_checkpoint")
1001 def _create_group(self, rank_list):
1002 """Create a new process group with the specified ranks.
1004 Internal method to be implemented by subclasses.
1006 Args:
1007 rank_list (list): List of ranks to include in the group.
1009 Returns:
1010 The newly created process group.
1011 """
1012 raise NotImplementedError("Platform subclasses must implement _create_group")
1014 def new_stream(self):
1015 """Create a new compute stream for asynchronous operations.
1017 Returns:
1018 A new stream object for the current device.
1019 """
1020 raise NotImplementedError("Platform subclasses must implement new_stream")
1022 def get_stream_context(self):
1023 """Get a context manager for executing operations on a specific stream.
1025 Returns:
1026 A context manager that can be used with 'with' statement to set stream.
1027 """
1028 raise NotImplementedError("Platform subclasses must implement get_stream_context")
1030 @staticmethod
1031 def get_tensor_transform():
1032 """Get the tensor transformation utilities for the current framework.
1034 Returns:
1035 A module or object containing tensor transformation functions.
1036 """
1037 raise NotImplementedError("Platform subclasses must implement get_tensor_transform")
1039 @staticmethod
1040 def construct_strided_slice(x, begin, end, stride):
1041 """Construct a strided slice operation on a tensor.
1043 Args:
1044 x: The input tensor to slice.
1045 begin: The starting indices for each dimension.
1046 end: The ending indices for each dimension.
1047 stride: The stride for each dimension.
1049 Returns:
1050 The sliced tensor.
1051 """
1052 raise NotImplementedError("Platform subclasses must implement construct_strided_slice")
1054 @staticmethod
1055 def micro_batch(micro_batch_num, args_batch_dim=None, kwargs_batch_dim=None):
1056 """Split inputs into micro-batches for pipeline parallelism.
1058 Args:
1059 micro_batch_num (int): The number of micro-batches to create.
1060 args_batch_dim (list, optional): Batch dimension for each positional arg.
1061 kwargs_batch_dim (dict, optional): Batch dimension for each keyword arg.
1063 Returns:
1064 A decorator that splits function inputs into micro-batches.
1065 """
1066 raise NotImplementedError("Platform subclasses must implement micro_batch")
1068 @staticmethod
1069 def get_symmetric_memory_handler():
1070 """Return a platform-specific symmetric memory handler instance."""
1071 raise NotImplementedError("Platform subclasses must implement get_symmetric_memory_handler")
1073 @staticmethod
1074 def load_into_param(param, data):
1075 """Load data into a parameter, handling framework-specific semantics."""
1076 raise NotImplementedError("Platform subclasses must implement load_into_param")
1078 def create_group(self, rank_list):
1079 """Create or retrieve a communication group with the specified ranks.
1081 If a group with the same rank list already exists, returns the existing
1082 group instead of creating a new one.
1084 Args:
1085 rank_list (list): List of ranks to include in the group.
1087 Returns:
1088 The process group for the specified ranks.
1089 """
1090 group_key = str(tuple(sorted(rank_list)))
1091 if group_key in EXISTING_COMM_GROUPS:
1092 return EXISTING_COMM_GROUPS[group_key]
1094 group = self._create_group(rank_list)
1095 EXISTING_COMM_GROUPS[group_key] = group
1096 return group
1098 @staticmethod
1099 def _process_current_handle():
1100 """Wait for the current gradient handle and execute post-process callback.
1102 Internal method to synchronize pending gradient operations.
1103 """
1104 if Platform.current_grad_handle is None:
1105 return
1107 Platform.current_grad_handle.wait()
1108 if Platform.post_grad_handle_process is None:
1109 return
1110 # pylint: disable=E1102
1111 Platform.post_grad_handle_process()
1113 def set_grad_reduce_handle(self, handle, post_process=None):
1114 """Set a new gradient reduction handle after waiting for the current one.
1116 Waits for any pending gradient handle on the grad sync stream, then
1117 sets the new handle and optional post-process callback.
1119 Args:
1120 handle: The async work handle for gradient reduction.
1121 post_process (callable, optional): Callback to run after handle completes.
1122 """
1123 if Platform.grad_sync_stream is None:
1124 Platform.grad_sync_stream = self.new_stream()
1125 stream_context = self.get_stream_context()
1126 with stream_context(Platform.grad_sync_stream):
1127 Platform._process_current_handle()
1128 Platform.current_grad_handle = handle
1129 Platform.post_grad_handle_process = post_process
1131 def wait_grad_handle(self):
1132 """Wait for the current gradient handle to complete.
1134 Blocks until the current gradient reduction handle completes and
1135 clears the handle state.
1136 """
1137 if Platform.current_grad_handle is None:
1138 return
1139 if Platform.grad_sync_stream is None:
1140 Platform.grad_sync_stream = self.new_stream()
1141 stream_context = self.get_stream_context()
1142 with stream_context(Platform.grad_sync_stream):
1143 Platform._process_current_handle()
1144 sync_event = Platform.grad_sync_stream.record_event()
1145 sync_event.wait()
1146 Platform.current_grad_handle = None
1147 Platform.post_grad_handle_process = None
1149 @staticmethod
1150 def all_gather_object(object_list, obj, group=None) -> None:
1151 """Gather Python objects from all ranks into a list.
1153 Each rank contributes its object, and all ranks receive the complete list.
1155 Args:
1156 object_list (list): List to store gathered objects (output parameter).
1157 obj: The Python object from this rank to contribute.
1158 group: The process group for communication. Defaults to None (default group).
1159 """
1160 raise NotImplementedError("Platform subclasses must implement all_gather_object")
1162 @staticmethod
1163 def barrier(group=None, async_op: bool = False, device_ids=None) -> Any:
1164 """Synchronize all processes in the given process group.
1166 Each rank blocks until every rank in the group enters this collective (when ``async_op``
1167 is False), or returns an async handle that must be completed before proceeding.
1169 Args:
1170 group: The process group or communication group. ``None`` uses the default group.
1171 async_op (bool): If True, returns a backend-specific async work handle. Default: False.
1172 device_ids: Optional device id list; semantics depend on the backend.
1174 Returns:
1175 Async work handle when ``async_op`` is True; otherwise ``None`` (unless the rank
1176 is not in the group, in which case the backend may return ``None``).
1177 """
1178 raise NotImplementedError("Platform subclasses must implement barrier")
1180 @staticmethod
1181 def init_process_group(
1182 backend: Optional[str] = None,
1183 *,
1184 init_method: Optional[str] = None,
1185 timeout: Optional[timedelta] = None,
1186 world_size: int = -1,
1187 rank: int = -1,
1188 store: Any = None,
1189 pg_options: Any = None,
1190 device_id: Any = None
1191 ) -> None:
1192 """
1193 Initialize the default distributed process group.
1195 Args:
1196 backend: The backend to use for distributed communication
1197 init_method: URL specifying how to initialize the process group
1198 timeout: Timeout for operations executed against the process group
1199 world_size: Number of processes participating in the job
1200 rank: Rank of the current process
1201 store: Key/value store for exchanging connection information
1202 pg_options: Process group options for backend-specific configurations
1203 device_id: Specific device this process will work on
1205 Raises:
1206 NotImplementedError: This method must be implemented by subclasses
1207 """
1208 raise NotImplementedError("Platform subclasses must implement init_process_group")
1210 @staticmethod
1211 def destroy_process_group(group=None) -> None:
1212 """
1213 Destroy a given process group.
1215 Args:
1216 group: The process group to be destroyed. If None, destroys the default group.
1218 Raises:
1219 NotImplementedError: This method must be implemented by subclasses
1220 """
1221 raise NotImplementedError("Platform subclasses must implement destroy_process_group")
1223 @staticmethod
1224 def get_process_group_ranks(group=None) -> list[int]:
1225 """
1226 Get rank list of the given process group.
1228 Args:
1229 group: The process group to get ranks from. If None, uses the default group.
1231 Returns:
1232 List of ranks in the specified process group.
1234 Raises:
1235 NotImplementedError: This method must be implemented by subclasses
1236 """
1237 raise NotImplementedError("Platform subclasses must implement get_process_group_ranks")
1239 @staticmethod
1240 def get_backend(group=None):
1241 """
1242 Get the backend of the given process group.
1243 Args:
1244 group: The process group to get backend from. If None, uses the default group.
1246 Returns:
1247 The backend name of the specified process group.
1249 Raises:
1250 NotImplementedError: This method must be implemented by subclasses
1251 """
1252 raise NotImplementedError("Platform subclasses must implement get_backend")
1254 @staticmethod
1255 def split_group(parent_pg: Any = None,
1256 split_ranks: Optional[list] = None,
1257 timeout: Optional[timedelta] = None,
1258 pg_options: Optional[Any] = None,
1259 group_desc: Optional[str] = None,
1260 ) -> Any:
1261 """Create a split group relative to the parent process group.
1263 Args:
1264 parent_pg: The parent process group to split from.
1265 split_ranks (list, optional): Ranks to include in the split group.
1266 timeout (timedelta, optional): Timeout for operations.
1267 pg_options: Process group options for backend-specific configurations.
1268 group_desc (str, optional): Description of the group.
1270 Returns:
1271 The new split process group.
1272 """
1273 raise NotImplementedError("Platform subclasses must implement split_group")
1275 @staticmethod
1276 def get_group_local_rank(group=None) -> int:
1277 """Get the local rank within the given process group.
1279 Args:
1280 group: The process group to query. If None, uses the default group.
1282 Returns:
1283 int: The local rank within the group.
1284 """
1285 raise NotImplementedError("Platform subclasses must implement get_group_local_rank")
1287 @staticmethod
1288 def no_grad():
1289 """Get a context manager to disable gradient computation.
1291 Returns:
1292 A context manager that disables gradient tracking.
1293 """
1294 raise NotImplementedError("Platform subclasses must implement no_grad")
1296 @staticmethod
1297 def preserve_version_counter(tensor):
1298 """Get a context manager that preserves version for an internal tensor update."""
1299 raise NotImplementedError("Platform subclasses must implement preserve_version_counter")
1301 @staticmethod
1302 def relu(tensor):
1303 """Apply ReLU activation element-wise.
1305 Args:
1306 tensor: Input tensor.
1308 Returns:
1309 Tensor with ReLU applied (max(0, x)).
1310 """
1311 raise NotImplementedError("Platform subclasses must implement relu")
1313 @staticmethod
1314 def cat(tensors, dim=0):
1315 """Concatenate tensors along a dimension."""
1316 raise NotImplementedError("Platform subclasses must implement cat")
1318 @staticmethod
1319 def empty_like(tensor, *, dtype=None, device=None, pin_memory=False):
1320 """Create an uninitialized tensor with the same shape as input.
1322 Args:
1323 tensor: The input tensor to copy shape from.
1324 dtype: Optional dtype for the new tensor. If None, uses input tensor's dtype.
1325 device: Optional device for the new tensor. If None, uses input tensor's device.
1326 pin_memory (bool): If True, allocate pinned memory for faster CPU-GPU transfer.
1328 Returns:
1329 An uninitialized tensor with the same shape as input.
1330 """
1331 raise NotImplementedError("Platform subclasses must implement empty_like")
1333 def get_current_stream(self):
1334 """Get the current compute stream for the device.
1336 Returns:
1337 The current stream object.
1338 """
1339 raise NotImplementedError("Platform subclasses must implement get_current_stream")
1341 def new_event(self):
1342 """Create a new event for stream synchronization.
1344 Returns:
1345 A new event object.
1346 """
1347 raise NotImplementedError("Platform subclasses must implement new_event")
1349 def tree_map(self, fn, tree):
1350 """Apply a function to all tensors in a nested structure.
1352 Args:
1353 fn (callable): Function to apply to each tensor.
1354 tree: Nested structure (list, tuple, dict) containing tensors.
1356 Returns:
1357 The same nested structure with fn applied to all tensors.
1358 """
1359 raise NotImplementedError("Platform subclasses must implement tree_map")
1361 @staticmethod
1362 def is_linear_module(module) -> bool:
1363 """Check whether *module* is a linear/dense layer for the current framework.
1365 Args:
1366 module: The module instance to check.
1368 Returns:
1369 True if *module* is the framework's linear layer type.
1370 """
1371 raise NotImplementedError("Platform subclasses must implement is_linear_module")
1373 @staticmethod
1374 def is_embedding_module(module) -> bool:
1375 """Check whether *module* is an embedding layer for the current framework.
1377 Args:
1378 module: The module instance to check.
1380 Returns:
1381 True if *module* is the framework's embedding layer type.
1382 """
1383 raise NotImplementedError("Platform subclasses must implement is_embedding_module")
1385 @staticmethod
1386 def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False):
1387 """Register a forward pre-hook on a module.
1389 Args:
1390 module: The module to register the hook on.
1391 hook (callable): The hook function to register.
1392 prepend (bool): If True, prepend the hook to existing hooks.
1393 with_kwargs (bool): If True, hook receives both args and kwargs.
1395 Returns:
1396 A handle that can be used to remove the hook.
1397 """
1398 return module.register_forward_pre_hook(hook, prepend=prepend, with_kwargs=with_kwargs)
1400 @staticmethod
1401 def register_full_backward_hook(module, hook, prepend=False):
1402 """Register a full backward hook on a module.
1404 Args:
1405 module: The module to register the hook on.
1406 hook (callable): The hook function to register.
1407 prepend (bool): If True, prepend the hook to existing hooks.
1409 Returns:
1410 A handle that can be used to remove the hook.
1411 """
1412 return module.register_full_backward_hook(hook, prepend)
1414 @staticmethod
1415 def register_full_backward_pre_hook(module, hook, prepend=False):
1416 """Register a full backward pre-hook on a module.
1418 Args:
1419 module: The module to register the hook on.
1420 hook (callable): The hook function to register.
1421 prepend (bool): If True, prepend the hook to existing hooks.
1423 Returns:
1424 A handle that can be used to remove the hook.
1425 """
1426 return module.register_full_backward_pre_hook(hook, prepend)
1428 @property
1429 def checkpoint(self):
1430 """Get the checkpoint function for activation checkpointing.
1432 Returns:
1433 The checkpoint function for the current framework.
1434 """
1435 raise NotImplementedError("Platform subclasses must implement checkpoint")
1437 @staticmethod
1438 def checkpoint_wrapper(module, **checkpoint_kwargs):
1439 """Wrap a module with activation checkpointing functionality.
1441 Args:
1442 module: The module or callable to wrap with activation checkpointing.
1443 **checkpoint_kwargs: Keyword arguments forwarded to the framework
1444 checkpoint wrapper implementation.
1446 Returns:
1447 The wrapped module with activation checkpointing enabled.
1448 """
1449 raise NotImplementedError("Platform subclasses must implement checkpoint_wrapper")
1451 @staticmethod
1452 def checkpoint_exclude_wrapper(module: Any) -> Any:
1453 """Wrap a callable whose activations should be saved instead of recomputed.
1455 Args:
1456 module: The module or callable to exclude from activation recomputation.
1458 Returns:
1459 The wrapped module or callable.
1460 """
1461 raise NotImplementedError("checkpoint_exclude_wrapper is currently only supported on MindSpore")
1463 @staticmethod
1464 def swap_wrapper(module, policy_fn=None, group_swap=False):
1465 """Wrap a module with activation swap functionality.
1467 Args:
1468 module: The module to wrap with activation swap.
1469 policy_fn: Optional per-tensor swap policy function.
1470 group_swap (bool, optional): Whether tensors participate in group copy fusion. Default: ``False``.
1472 Returns:
1473 The wrapped module with activation swap enabled.
1474 """
1475 raise NotImplementedError("Platform subclasses must implement swap_wrapper")
1477 @staticmethod
1478 def swap_tensor_wrapper(target, tag=None, group_swap=False):
1479 """Register target tensors into the current swap group.
1481 Args:
1482 target: A tensor or nested container of tensors to register.
1483 tag: Optional debug tag associated with the wrapped tensors.
1484 group_swap (bool, optional): Whether tensors participate in group copy fusion. Default: ``False``.
1486 Returns:
1487 The original target structure, unchanged semantically.
1488 """
1489 raise NotImplementedError("Platform subclasses must implement swap_tensor_wrapper")
1491 @staticmethod
1492 def get_class_activation_wrapper():
1493 """Return the platform-specific activation wrapper class."""
1494 raise NotImplementedError("Platform subclasses must implement get_class_activation_wrapper")
1496 @property
1497 def noop_context_fn(self):
1498 """Get a no-op context function for checkpointing.
1500 Returns:
1501 A context function that performs no operation.
1502 """
1503 raise NotImplementedError("Platform subclasses must implement noop_context_fn")
1505 @staticmethod
1506 def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
1507 """Create contexts for selective activation checkpointing.
1509 Args:
1510 policy_fn_or_list: A policy function or list of layer names to checkpoint.
1511 allow_cache_entry_mutation (bool): Whether to allow cache entry mutation.
1512 group_swap (bool, optional): Whether MUST_SWAP tensors participate in group copy fusion. Default: ``False``.
1514 Returns:
1515 Context functions for selective checkpointing.
1516 """
1517 raise NotImplementedError("Platform subclasses must implement create_selective_checkpoint_contexts")
1519 @staticmethod
1520 def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1521 """Create an async CPU offload context for activation checkpointing.
1523 Args:
1524 policy_fn: Optional policy function to determine which activations to offload.
1525 group_swap (bool): Whether swapped tensors participate in group copy fusion.
1526 Default: ``False``.
1528 Returns:
1529 Context manager for async CPU offloading during checkpointing.
1530 """
1531 raise NotImplementedError("Platform subclasses must implement async_save_on_cpu")
1533 @staticmethod
1534 def recompute_handle_collector_ctx():
1535 """Context manager that collects recompute handles created in its scope.
1537 Yields:
1538 A list populated with one opaque recompute handle per checkpointed
1539 block executed during the forward pass within the context. Each
1540 handle can later be fired via :meth:`recompute_handle`.
1541 """
1542 raise NotImplementedError("Platform subclasses must implement recompute_handle_collector_ctx")
1544 @staticmethod
1545 def recompute_handle(handle, session_id):
1546 """Eagerly fire one checkpointed block's forward re-run.
1548 Materializes and caches the block's activations under ``session_id`` so
1549 a later backward in the same session reuses them instead of re-running.
1551 Args:
1552 handle: An opaque recompute handle from
1553 :meth:`recompute_handle_collector_ctx`.
1554 session_id: Stable key shared by the producing re-run and the
1555 consuming backward.
1556 """
1557 raise NotImplementedError("Platform subclasses must implement recompute_handle")
1559 @staticmethod
1560 def recompute_session_ctx(session_id, retain_on_unpack=False):
1561 """Context manager binding recompute unpack to a caller-provided session.
1563 Args:
1564 session_id: Stable session key. Recompute caches are keyed by this
1565 instead of the transient autodiff engine id, so a re-run fired
1566 under one engine can be reused by another.
1567 retain_on_unpack (bool): When ``True``, unpack returns recomputed
1568 tensors without popping them, so a later backward can consume
1569 them. Default: ``False``.
1571 Returns:
1572 A context manager activating the session for its scope.
1573 """
1574 raise NotImplementedError("Platform subclasses must implement recompute_session_ctx")
1576 @staticmethod
1577 def clear_recompute_session(session_id):
1578 """Release retained recompute data for a session.
1580 Args:
1581 session_id: The session key whose cached recompute data is cleared.
1582 """
1583 raise NotImplementedError("Platform subclasses must implement clear_recompute_session")
1585 @staticmethod
1586 def get_element_size(tensor):
1587 """Get Tensor Element Size"""
1588 raise NotImplementedError("Platform subclasses must implement get_element_size")
1590 @staticmethod
1591 def alloc_tensor_buffer(numel: int, dtype, device, pin_memory: bool = False):
1592 """Allocate an uninitialized 1-D tensor buffer."""
1593 raise NotImplementedError("Platform subclasses must implement alloc_tensor_buffer")
1595 @staticmethod
1596 def tensor_to_numpy(tensor) -> np.ndarray:
1597 """Convert a framework tensor to a NumPy array.
1599 Args:
1600 tensor: The tensor to convert.
1602 Returns:
1603 np.ndarray: The tensor data as a NumPy array.
1604 """
1605 raise NotImplementedError("Platform subclasses must implement tensor_to_numpy")
1607 @staticmethod
1608 def from_numpy(np_array):
1609 """Create a host-resident tensor from a NumPy array (inverse of tensor_to_numpy).
1611 The result stays on the host regardless of the active device context, so it
1612 remains asnumpy-able even when built under ``ms.DeviceCtx("meta")`` (e.g. while
1613 ``fully_shard`` lazily constructs a default device mesh). Use it for rank/mesh
1614 bookkeeping tensors, which are only ever read back via ``tensor_to_numpy``.
1615 """
1616 raise NotImplementedError("Platform subclasses must implement from_numpy")
1618 @staticmethod
1619 def profiler_record(name):
1620 """Record a profiler event with the given name.
1622 Args:
1623 name (str): The name of the profiler event.
1625 Returns:
1626 A context manager or decorator for profiling a code region.
1627 """
1628 raise NotImplementedError("Platform subclasses must implement profiler_record")
1630 def cast_fp_tensor(self, dtype, x):
1631 """Cast floating-point tensor to target dtype if applicable.
1633 Args:
1634 dtype: The target dtype to cast to.
1635 x: The input tensor.
1637 Returns:
1638 The tensor cast to target dtype, or unchanged if not floating-point.
1639 """
1640 raise NotImplementedError("Platform subclasses must implement cast_fp_tensor")
1642 def apply_to_tensors(self, fn, container):
1643 """Recursively apply a function to all tensors in a container.
1645 Supports nested structures including lists, tuples, and dicts.
1647 Args:
1648 fn (callable): Function to apply to each tensor.
1649 container: Nested structure containing tensors.
1651 Returns:
1652 The same structure with fn applied to all tensors.
1653 """
1654 raise NotImplementedError("Platform subclasses must implement apply_to_tensors")
1656 @staticmethod
1657 def clip_grad_norm_(
1658 parameters, max_norm: float, norm_type: float = 2.0,
1659 error_if_nonfinite: bool = False, foreach=None,
1660 ):
1661 """Compute and clip gradient norms for distributed models.
1663 Communication is derived from each parameter's DTensor spec.
1664 Subclasses must implement this method.
1666 Args:
1667 parameters: An ``nn.Module``, a single ``Tensor``, or an
1668 iterable of ``Tensor`` s whose gradients to clip.
1669 max_norm: Maximum allowed gradient norm.
1670 norm_type: Type of the norm (default ``2.0``).
1671 error_if_nonfinite: If ``True``, raise when total norm is
1672 non-finite. Default ``False``.
1673 foreach: Unused, accepted for API compatibility.
1675 Returns:
1676 The total (unclipped) gradient norm.
1677 """
1678 raise NotImplementedError(
1679 "Platform subclasses must implement clip_grad_norm_"
1680 )
1682 @staticmethod
1683 def get_created_group(rank_list: Union[list[int], tuple[int]]):
1684 """Get an existing process group by rank list.
1686 Args:
1687 rank_list (Union[list[int], tuple[int]]): Tuple or list of ranks.
1689 Returns:
1690 The process group corresponding to the rank list if it exists, else None.
1691 """
1692 group_key = str(tuple(sorted(rank_list)))
1693 if group_key in EXISTING_COMM_GROUPS:
1694 return EXISTING_COMM_GROUPS[group_key]
1695 return None
1697 @classmethod
1698 def mark_created_groups(cls, process_group: Union[Any, list[Any]]) -> None:
1699 """Register process groups in the global cache for reuse.
1701 Args:
1702 process_group (Union[Any, list[Any]]): A process group or a list of process groups.
1703 """
1704 if not isinstance(process_group, list):
1705 process_group = [process_group]
1706 for group in process_group:
1707 rank_list = cls.get_process_group_ranks(group)
1708 group_key = str(tuple(sorted(rank_list)))
1709 EXISTING_COMM_GROUPS[group_key] = group
1711 @property
1712 def meta_device(self):
1713 """Get the framework-specific meta device for tensor shape inference.
1715 The meta device allows creating tensors without allocating actual storage,
1716 useful for shape inference and model initialization.
1718 Returns:
1719 The meta device object for the current framework.
1720 """
1721 raise NotImplementedError("Platform subclasses must implement meta_device")
1723 def init_on_device(self, device, include_buffers=False):
1724 """Get a context manager for initializing module parameters on a device.
1726 Args:
1727 device: The target device for parameter initialization.
1728 include_buffers (bool): If True, also initialize buffers on the device.
1730 Returns:
1731 A context manager for device-specific initialization.
1732 """
1733 raise NotImplementedError("Platform subclasses must implement init_on_device")
1735 def str_to_dtype(self, dtype_str: str) -> Any:
1736 """
1737 Map a framework-style dtype string (e.g. ``torch.float32``) to the backend dtype object.
1739 Args:
1740 dtype_str (str): Serialized dtype identifier produced by checkpoint metadata.
1742 Returns:
1743 Framework dtype object (e.g. ``torch.dtype`` or MindSpore dtype).
1744 """
1745 raise NotImplementedError("Platform subclasses must implement str_to_dtype")
1747 def list_to_size(self, size_list: list[int]) -> Any:
1748 """
1749 Convert a shape list from checkpoint metadata to the framework's size type (e.g. ``torch.Size``).
1751 Args:
1752 size_list (list[int]): Tensor global shape as a list of ints.
1754 Returns:
1755 Framework-specific size object.
1756 """
1757 raise NotImplementedError("Platform subclasses must implement list_to_size")