« 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.
15# Adapted from
16# hyper_parallel/platform/torch/activation_checkpoint/activation_swap.py
17# adapted for MindSpore Cell API.
18# ============================================================================
19"""Activation Swap Wrapper implementation for MindSpore."""
20from abc import ABC, abstractmethod
21from collections.abc import Iterator
22from typing import Optional, Callable, Any, Union
23import types
24import warnings
25import mindspore as ms
26from mindspore import Tensor
27from mindspore.common.parameter import Parameter
28from mindspore.nn import Cell
31_CKPT_WRAPPED_MODULE = "_ckpt_wrapped_module"
34def _strip_ckpt_wrapped_module_prefix(name: str) -> str:
35 """Remove the wrapper cell segment from a dotted MindSpore cell name."""
36 return ".".join(part for part in name.split(".") if part != _CKPT_WRAPPED_MODULE)
39class FuncCell(Cell):
40 """
41 Thin :class:`~mindspore.nn.Cell` adapter that wraps a plain callable.
43 Allows ordinary Python functions (or any callable without Cell
44 parameters) to be passed to :func:`checkpoint_wrapper` and
45 :func:`swap_wrapper` in place of a :class:`~mindspore.nn.Cell`.
46 The wrapped function is stored as ``_fn`` and invoked in
47 :meth:`construct`; the cell has no trainable parameters.
49 Args:
50 fn (callable): The function to wrap.
52 Example:
53 >>> wrapped = checkpoint_wrapper(lambda x: x * 2)
54 """
56 def __init__(self, fn: Callable):
57 super().__init__()
58 self._fn = fn
60 def construct(self, *args, **kwargs):
61 """Delegate to the wrapped function."""
62 return self._fn(*args, **kwargs)
65def _is_shared_function_callable(callable_obj: Callable) -> bool:
66 """Return True for stateless function objects commonly shared by modules."""
67 return isinstance(callable_obj, (types.FunctionType, types.BuiltinFunctionType, types.MethodType))
70def _iter_wrappable_callable_attrs(module: Cell) -> Iterator[tuple[str, Callable]]:
71 """Yield public per-instance callable attributes not registered as child cells.
73 Plain functions, builtins and bound methods are skipped: these are stateless
74 module-level utilities shared by reference across many modules (e.g.
75 ``self.reshape = mint.reshape`` / ``self.cast = ops.cast`` repeated in every
76 layer). They are never standalone checkpoint regions, and marking a shared
77 function's ``_is_wrapped`` flag would both mutate a global object and falsely
78 flag every sibling module that references the same function as an overlapping
79 wrap. Only per-instance callables (e.g. MindSpore ``Primitive`` objects)
80 participate in overlap tracking.
81 """
82 for attr_name, attr_value in vars(module).items():
83 if attr_name.startswith("_") or isinstance(attr_value, Cell):
84 continue
85 if _is_shared_function_callable(attr_value):
86 continue
87 if callable(attr_value):
88 yield attr_name, attr_value
91def _mark_wrapped(obj: Any) -> None:
92 try:
93 obj._is_wrapped = True # pylint: disable=W0212
94 except (AttributeError, TypeError):
95 pass
98def _get_wrapped_callable(cell: Cell) -> Optional[Callable]:
99 wrapped_module = getattr(cell, _CKPT_WRAPPED_MODULE, None)
100 if isinstance(wrapped_module, FuncCell):
101 return getattr(wrapped_module, "_fn", None)
102 if isinstance(cell, FuncCell):
103 return getattr(cell, "_fn", None)
104 return None
107def _raise_callable_already_wrapped(callable_obj: Callable) -> None:
108 warnings.warn(
109 f"Callable '{callable_obj.__class__.__name__}' is already wrapped. "
110 "Wrapping overlapping module regions is not allowed."
111 )
114def _check_callable_attr_not_wrapped(owner: Cell, attr_name: str, attr_value: Callable) -> None:
115 del owner, attr_name
116 if getattr(attr_value, '_is_wrapped', False):
117 _raise_callable_already_wrapped(attr_value)
120def _check_and_mark_callable(callable_obj: Callable) -> None:
121 if _is_shared_function_callable(callable_obj):
122 return
123 if getattr(callable_obj, '_is_wrapped', False):
124 warnings.warn(
125 f"Callable '{callable_obj.__class__.__name__}' or one of its ancestors is already wrapped. "
126 "Wrapping overlapping module regions is not allowed."
127 )
128 _mark_wrapped(callable_obj)
131def _check_and_mark_wrapped(module: Cell) -> None:
132 """Validate no wrapping overlap, then mark module and all descendants as wrapped.
134 Raises:
135 ValueError: If ``module`` or any of its descendants is already wrapped.
136 """
137 if getattr(module, '_is_wrapped', False):
138 warnings.warn(
139 f"Module '{module.__class__.__name__}' or one of its ancestors is already wrapped. "
140 "Wrapping overlapping module regions is not allowed."
141 )
142 for _, submodule in module.cells_and_names():
143 if submodule is module:
144 continue
145 wrapped_callable = _get_wrapped_callable(submodule)
146 if wrapped_callable is not None and _is_shared_function_callable(wrapped_callable):
147 continue
148 if getattr(submodule, '_is_wrapped', False):
149 if wrapped_callable is not None:
150 _raise_callable_already_wrapped(wrapped_callable)
151 warnings.warn(
152 f"Submodule '{getattr(submodule, '_ckpt_wrapped_module', submodule).__class__.__name__}' of "
153 f"'{module.__class__.__name__}' is already wrapped. "
154 "Wrapping overlapping module regions is not allowed."
155 )
156 for _, submodule in module.cells_and_names():
157 for attr_name, attr_value in _iter_wrappable_callable_attrs(submodule):
158 _check_callable_attr_not_wrapped(submodule, attr_name, attr_value)
159 for _, submodule in module.cells_and_names():
160 _mark_wrapped(submodule)
161 for _, attr_value in _iter_wrappable_callable_attrs(submodule):
162 _mark_wrapped(attr_value)
165class ActivationWrapper(Cell, ABC):
166 """
167 Base class for Activation Checkpoint Wrapper in MindSpore.
169 Wraps a :class:`mindspore.nn.Cell` and forwards attribute lookups,
170 parameter iteration, and indexing to the inner cell. Concrete
171 sub-classes must implement :meth:`construct`.
173 Not meant to be instantiated directly.
174 """
176 def __init__(self, module: Union[Cell, Callable], *, track_overlaps: bool = True) -> None:
177 """Initialize a wrapper and optionally participate in overlap tracking."""
178 if callable(module) and not isinstance(module, Cell):
179 if track_overlaps:
180 _check_and_mark_callable(module)
181 module = FuncCell(module)
182 if track_overlaps:
183 _mark_wrapped(module)
184 elif track_overlaps:
185 _check_and_mark_wrapped(module)
186 super().__init__(auto_prefix=False)
187 self._ckpt_wrapped_module = module
188 self._is_wrapped = track_overlaps
189 self._wrapped_param_names = {
190 id(param): param.name for _, param in module.parameters_and_names()
191 }
193 @property
194 def _wrapped_module(self) -> Cell:
195 """Return the wrapped callable normalized to a MindSpore Cell."""
196 return self._ckpt_wrapped_module
198 @abstractmethod
199 def construct(self, *args, **kwargs):
200 """Abstract construct method — subclasses must override."""
201 raise ValueError("Subclasses should implement construct().")
203 def __getattr__(self, name: str) -> Any:
204 """Forward missing attributes to the wrapped cell.
206 .. warning::
207 Do **not** call ``super().__getattr__(name)`` here.
208 MindSpore's ``Cell.__init__`` calls ``hasattr(self, "bprop")`` at
209 line 252 of ``cell.py`` *after* ``_cells`` is initialised as an
210 empty ``OrderedDict`` but *before* ``ActivationWrapper.__init__``
211 has registered ``_ckpt_wrapped_module`` into ``_cells``. The
212 PyTorch ``nn.Module.__init__`` is pure Python and never calls
213 ``hasattr`` on ``self``, so this issue does not arise there.
215 Using ``super().__getattr__`` here would raise ``AttributeError``
216 (``_ckpt_wrapped_module`` not yet in ``_cells``), the fallback
217 ``getattr(self._ckpt_wrapped_module, name)`` would access
218 ``self._ckpt_wrapped_module`` — triggering another
219 ``__getattr__("_ckpt_wrapped_module")`` — and the cycle repeats
220 as infinite recursion.
222 Instead we replicate ``Cell.__getattr__``'s own dict-probe logic
223 and fall through to the wrapped module only when it is already
224 registered.
225 """
226 for attr_dict in ('_params', '_buffers', '_cells', '_params_list'):
227 d = self.__dict__.get(attr_dict)
228 if d is not None and name in d:
229 return d[name]
230 cells = self.__dict__.get('_cells', {})
231 wrapped = cells.get(_CKPT_WRAPPED_MODULE)
232 if wrapped is not None:
233 return getattr(wrapped, name)
234 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
236 @property
237 def unwrap_cell(self) -> Cell:
238 """Recursively return the innermost wrapped cell."""
239 return self._ckpt_wrapped_module
241 def __getitem__(self, key: int) -> Any:
242 """Forward indexing calls in case the wrapped cell is a SequentialCell."""
243 return self._ckpt_wrapped_module.__getitem__(key) # type: ignore[operator]
245 def cells_and_names(self, cells=None, name_prefix=''):
246 """
247 Return wrapped cells without exposing the wrapper storage prefix.
249 MindSpore registers ``_ckpt_wrapped_module`` as a real child cell, so
250 the default :meth:`Cell.cells_and_names` would expose names such as
251 ``layer._ckpt_wrapped_module.attn``. Strip that implementation detail
252 so downstream code sees the same names as it would for the unwrapped
253 model.
254 """
255 for cell_name, cell in super().cells_and_names(cells, name_prefix):
256 yield _strip_ckpt_wrapped_module_prefix(cell_name), cell
258 def parameters_and_names(
259 self,
260 name_prefix: str = '',
261 expand: bool = True,
262 ) -> Iterator[tuple[str, Parameter]]:
263 """
264 Override :meth:`parameters_and_names` to strip the wrapper prefix.
266 Removes all occurrences of ``_ckpt_wrapped_module.`` from parameter
267 names so that a checkpoint saved from this wrapper is compatible with
268 the unwrapped cell.
270 Args:
271 name_prefix (str): Prefix prepended to every parameter name.
272 expand (bool): Whether to recursively expand sub-cells.
274 Yields:
275 tuple[str, Parameter]: ``(name, parameter)`` pairs with the
276 wrapper prefix removed.
277 """
278 for param_name, param in super().parameters_and_names(name_prefix, expand):
279 yield _strip_ckpt_wrapped_module_prefix(param_name), param
281 def update_parameters_name(self, prefix='', recurse=True):
282 """
283 Update wrapped parameter names without collapsing existing full paths.
285 When a wrapper replaces an already-registered child cell, the wrapped
286 parameters usually already have globally unique names such as
287 ``0.attn.qkv.weight``. MindSpore will still call
288 ``wrapper.update_parameters_name("attn.")`` during reassignment; if we
289 blindly apply that prefix again through the wrapper view, those names
290 are rewritten to ``attn.qkv.weight`` and collide across layers.
292 For parameters that already contain the requested child prefix in their
293 existing full name, keep the current name unchanged. For fresh
294 standalone modules that only have local names like ``qkv.weight``,
295 synthesize the prefixed name as usual.
296 """
297 if prefix is None:
298 prefix = ''
299 for local_name, param in self._ckpt_wrapped_module.parameters_and_names(expand=recurse):
300 original_name = self._wrapped_param_names.get(id(param), param.name)
301 if prefix and (original_name.startswith(prefix) or f".{prefix}" in original_name):
302 new_name = original_name
303 elif prefix:
304 new_name = prefix + local_name
305 else:
306 new_name = local_name
307 if new_name != param.name:
308 param.is_init = False
309 param.name = new_name
310 self._wrapped_param_names[id(param)] = new_name
313def base_check_fn(tensor: Any) -> bool:
314 """
315 Basic eligibility check: returns ``True`` when *tensor* may be offloaded.
317 Skips:
319 * Non-tensor objects.
320 * :class:`~mindspore.common.parameter.Parameter` objects.
321 * Empty tensors (zero elements).
323 Args:
324 tensor: The value to test.
326 Returns:
327 bool: ``True`` if the tensor is eligible for CPU offloading.
328 """
329 if not isinstance(tensor, Tensor):
330 return False
331 if tensor.param_info is not None:
332 return False
333 if tensor.untyped_storage().size() == 0:
334 return False
335 return True
338def _normalize_device(device: str) -> str:
339 if ":" in device:
340 return device.split(":", maxsplit=1)[0]
341 return device
344class AsyncSaveOnCpu(ms.saved_tensors_hooks):
345 """
346 Context manager to offload tensors to CPU during forward pass.
347 """
348 def __init__(self, policy_fn=None, group_swap: bool = False) -> None:
349 # pylint: disable=C0415
350 from hyper_parallel.core.activation_checkpoint.activation_checkpoint import CheckpointPolicy
351 from hyper_parallel.core.activation_checkpoint.swap import Storage, SwapManager, SwapTensor
352 self.add_to_storage = False
353 self.storage = Storage()
354 self.count_idx = 0
355 self.policy_fn = policy_fn
358 # Cache per-context-manager state once to avoid per-tensor singleton lookups.
359 swap_manager = SwapManager()
361 def pack_to_cpu(tensor: ms.Tensor):
362 if not base_check_fn(tensor):
363 return tensor
364 if policy_fn is not None:
365 if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE:
366 return tensor
367 if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP:
368 raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}")
369 group_name = swap_manager.get_current_group_name()
370 if not group_name:
371 return tensor
372 if not self.add_to_storage:
373 swap_manager.add_storage(group_name, self.storage)
374 self.add_to_storage = True
375 funcname = f"{group_name}::{tensor.shape}"
376 self.storage[self.count_idx].append(
377 SwapTensor(tensor, funcname, group_swap=group_swap)
378 )
379 self.count_idx += 1
380 return tensor
382 def unpack_from_cpu(tensor) -> ms.Tensor:
383 if self.storage is not None:
384 self.storage.clear()
385 self.storage = None
386 return tensor
388 super().__init__(pack_to_cpu, unpack_from_cpu)
391class SwapWrapper(ActivationWrapper):
392 """
393 MindSpore counterpart of :class:`~hyper_parallel.platform.torch
394 .activation_checkpoint.activation_swap.SwapWrapper`.
396 Wraps a :class:`~mindspore.nn.Cell` and applies async activation swap
397 during the forward pass via the platform's ``async_save_on_cpu`` context
398 manager. Falls back to a no-op context when that context is not yet
399 available on the current platform.
401 Args:
402 mod (Cell): The cell whose intermediate activations should be swapped.
403 policy_fn (callable, optional): Per-tensor swap policy; see
404 :class:`AsyncSaveOnCpu`.
406 Example:
407 >>> from hyper_parallel.platform.mindspore.activation_checkpoint import swap_wrapper
408 >>> model.layers[i].attn = swap_wrapper(model.layers[i].attn, policy_fn)
409 """
411 def __init__(
412 self,
413 mod: Union[Cell, Callable],
414 policy_fn: Optional[Callable] = None,
415 group_swap: bool = False,
416 ):
417 super().__init__(mod)
418 self.policy_fn = policy_fn
419 self.group_swap = group_swap
421 def construct(self, *args, **kwargs):
422 """Execute the wrapped module inside an async CPU-swap context."""
423 with AsyncSaveOnCpu(policy_fn=self.policy_fn, group_swap=self.group_swap):
424 return self._ckpt_wrapped_module(*args, **kwargs)
427def swap_wrapper(
428 module: Union[Cell, Callable],
429 policy_fn: Optional[Callable] = None,
430 group_swap: bool = False,
431) -> SwapWrapper:
432 """
433 Wrap *module* with async activation swap.
435 Args:
436 module (Cell or callable): The cell or plain function to wrap.
437 If a plain callable is passed it is automatically wrapped in a
438 :class:`FuncCell` before being stored.
439 policy_fn (callable, optional): Per-tensor swap policy; see
440 :class:`AsyncSaveOnCpu`.
442 Returns:
443 SwapWrapper: The wrapped cell with activation swap enabled.
444 """
445 return SwapWrapper(module, policy_fn, group_swap)
448def swap_tensor_wrapper(target, tag: Optional[str] = None, group_swap: bool = False):
449 """Register selected tensors into the current swap group.
451 This helper is intended to be used inside a forward path that already
452 participates in the existing swap scheduling managed by ``SwapManager``.
453 It preserves the input structure and returns the original tensors.
454 """
455 # pylint: disable=C0415
456 from hyper_parallel.core.activation_checkpoint.swap import Storage, SwapManager, SwapTensor
457 swap_manager = SwapManager()
458 group_name = swap_manager.get_current_group_name()
459 if not group_name:
460 warnings.warn(
461 f"Tensor {tag} cannot be swapped, for its group is unregistered."
462 )
463 return target
464 if swap_manager.is_last_group(group_name):
465 return target
467 storage = Storage()
468 count_idx = 0
470 def _apply(x):
471 nonlocal count_idx
472 if isinstance(x, Tensor) and base_check_fn(x):
473 tensor_tag = tag or f"{group_name}_swap_tensor"
474 funcname = f"{tensor_tag}::{tuple(x.shape)}"
475 storage[count_idx].append(SwapTensor(x, funcname, group_swap=group_swap))
476 count_idx += 1
477 return x
479 def _map(tree):
480 if isinstance(tree, dict):
481 return type(tree)((k, _map(v)) for k, v in tree.items())
482 if isinstance(tree, tuple):
483 return tuple(_map(v) for v in tree)
484 if isinstance(tree, list):
485 return [_map(v) for v in tree]
486 return _apply(tree)
488 wrapped = _map(target)
489 if count_idx > 0:
490 swap_manager.add_storage(group_name, storage)
491 return wrapped