« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
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]):
177 if callable(module) and not isinstance(module, Cell):
178 _check_and_mark_callable(module)
179 module = FuncCell(module)
180 _mark_wrapped(module)
181 else:
182 _check_and_mark_wrapped(module)
183 super().__init__(auto_prefix=False)
184 self._ckpt_wrapped_module = module
185 self._is_wrapped = True
186 self._wrapped_param_names = {
187 id(param): param.name for _, param in module.parameters_and_names()
188 }
190 @property
191 def _wrapped_module(self):
192 return self._ckpt_wrapped_module
194 @abstractmethod
195 def construct(self, *args, **kwargs):
196 """Abstract construct method — subclasses must override."""
197 raise ValueError("Subclasses should implement construct().")
199 def __getattr__(self, name: str) -> Any:
200 """Forward missing attributes to the wrapped cell.
202 .. warning::
203 Do **not** call ``super().__getattr__(name)`` here.
204 MindSpore's ``Cell.__init__`` calls ``hasattr(self, "bprop")`` at
205 line 252 of ``cell.py`` *after* ``_cells`` is initialised as an
206 empty ``OrderedDict`` but *before* ``ActivationWrapper.__init__``
207 has registered ``_ckpt_wrapped_module`` into ``_cells``. The
208 PyTorch ``nn.Module.__init__`` is pure Python and never calls
209 ``hasattr`` on ``self``, so this issue does not arise there.
211 Using ``super().__getattr__`` here would raise ``AttributeError``
212 (``_ckpt_wrapped_module`` not yet in ``_cells``), the fallback
213 ``getattr(self._ckpt_wrapped_module, name)`` would access
214 ``self._ckpt_wrapped_module`` — triggering another
215 ``__getattr__("_ckpt_wrapped_module")`` — and the cycle repeats
216 as infinite recursion.
218 Instead we replicate ``Cell.__getattr__``'s own dict-probe logic
219 and fall through to the wrapped module only when it is already
220 registered.
221 """
222 for attr_dict in ('_params', '_buffers', '_cells', '_params_list'):
223 d = self.__dict__.get(attr_dict)
224 if d is not None and name in d:
225 return d[name]
226 cells = self.__dict__.get('_cells', {})
227 wrapped = cells.get(_CKPT_WRAPPED_MODULE)
228 if wrapped is not None:
229 return getattr(wrapped, name)
230 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
232 @property
233 def unwrap_cell(self) -> Cell:
234 """Recursively return the innermost wrapped cell."""
235 return self._ckpt_wrapped_module
237 def __getitem__(self, key: int) -> Any:
238 """Forward indexing calls in case the wrapped cell is a SequentialCell."""
239 return self._ckpt_wrapped_module.__getitem__(key) # type: ignore[operator]
241 def cells_and_names(self, cells=None, name_prefix=''):
242 """
243 Return wrapped cells without exposing the wrapper storage prefix.
245 MindSpore registers ``_ckpt_wrapped_module`` as a real child cell, so
246 the default :meth:`Cell.cells_and_names` would expose names such as
247 ``layer._ckpt_wrapped_module.attn``. Strip that implementation detail
248 so downstream code sees the same names as it would for the unwrapped
249 model.
250 """
251 for cell_name, cell in super().cells_and_names(cells, name_prefix):
252 yield _strip_ckpt_wrapped_module_prefix(cell_name), cell
254 def parameters_and_names(
255 self,
256 name_prefix: str = '',
257 expand: bool = True,
258 ) -> Iterator[tuple[str, Parameter]]:
259 """
260 Override :meth:`parameters_and_names` to strip the wrapper prefix.
262 Removes all occurrences of ``_ckpt_wrapped_module.`` from parameter
263 names so that a checkpoint saved from this wrapper is compatible with
264 the unwrapped cell.
266 Args:
267 name_prefix (str): Prefix prepended to every parameter name.
268 expand (bool): Whether to recursively expand sub-cells.
270 Yields:
271 tuple[str, Parameter]: ``(name, parameter)`` pairs with the
272 wrapper prefix removed.
273 """
274 for param_name, param in super().parameters_and_names(name_prefix, expand):
275 yield _strip_ckpt_wrapped_module_prefix(param_name), param
277 def update_parameters_name(self, prefix='', recurse=True):
278 """
279 Update wrapped parameter names without collapsing existing full paths.
281 When a wrapper replaces an already-registered child cell, the wrapped
282 parameters usually already have globally unique names such as
283 ``0.attn.qkv.weight``. MindSpore will still call
284 ``wrapper.update_parameters_name("attn.")`` during reassignment; if we
285 blindly apply that prefix again through the wrapper view, those names
286 are rewritten to ``attn.qkv.weight`` and collide across layers.
288 For parameters that already contain the requested child prefix in their
289 existing full name, keep the current name unchanged. For fresh
290 standalone modules that only have local names like ``qkv.weight``,
291 synthesize the prefixed name as usual.
292 """
293 if prefix is None:
294 prefix = ''
295 for local_name, param in self._ckpt_wrapped_module.parameters_and_names(expand=recurse):
296 original_name = self._wrapped_param_names.get(id(param), param.name)
297 if prefix and (original_name.startswith(prefix) or f".{prefix}" in original_name):
298 new_name = original_name
299 elif prefix:
300 new_name = prefix + local_name
301 else:
302 new_name = local_name
303 if new_name != param.name:
304 param.is_init = False
305 param.name = new_name
306 self._wrapped_param_names[id(param)] = new_name
309def base_check_fn(tensor: Any) -> bool:
310 """
311 Basic eligibility check: returns ``True`` when *tensor* may be offloaded.
313 Skips:
315 * Non-tensor objects.
316 * :class:`~mindspore.common.parameter.Parameter` objects.
317 * Empty tensors (zero elements).
319 Args:
320 tensor: The value to test.
322 Returns:
323 bool: ``True`` if the tensor is eligible for CPU offloading.
324 """
325 if not isinstance(tensor, Tensor):
326 return False
327 if tensor.param_info is not None:
328 return False
329 if tensor.untyped_storage().size() == 0:
330 return False
331 return True
334def _normalize_device(device: str) -> str:
335 if ":" in device:
336 return device.split(":", maxsplit=1)[0]
337 return device
340class AsyncSaveOnCpu(ms.saved_tensors_hooks):
341 """
342 Context manager to offload tensors to CPU during forward pass.
343 """
344 def __init__(self, policy_fn=None, group_swap: bool = False) -> None:
345 # pylint: disable=C0415
346 from hyper_parallel.core.activation_checkpoint.activation_checkpoint import CheckpointPolicy
347 from hyper_parallel.core.activation_checkpoint.swap import Storage, SwapManager, SwapTensor
348 self.add_to_storage = False
349 self.storage = Storage()
350 self.count_idx = 0
351 self.policy_fn = policy_fn
354 # Cache per-context-manager state once to avoid per-tensor singleton lookups.
355 swap_manager = SwapManager()
357 def pack_to_cpu(tensor: ms.Tensor):
358 if not base_check_fn(tensor):
359 return tensor
360 if policy_fn is not None:
361 if policy_fn(tensor) == CheckpointPolicy.MUST_SAVE:
362 return tensor
363 if policy_fn(tensor) != CheckpointPolicy.MUST_SWAP:
364 raise RuntimeError(f"Swap :set an invalid policy {policy_fn(tensor)}")
365 group_name = swap_manager.get_current_group_name()
366 if not group_name:
367 return tensor
368 if not self.add_to_storage:
369 swap_manager.add_storage(group_name, self.storage)
370 self.add_to_storage = True
371 funcname = f"{group_name}::{tensor.shape}"
372 self.storage[self.count_idx].append(
373 SwapTensor(tensor, funcname, group_swap=group_swap)
374 )
375 self.count_idx += 1
376 return tensor
378 def unpack_from_cpu(tensor) -> ms.Tensor:
379 if self.storage is not None:
380 self.storage.clear()
381 self.storage = None
382 return tensor
384 super().__init__(pack_to_cpu, unpack_from_cpu)
387class SwapWrapper(ActivationWrapper):
388 """
389 MindSpore counterpart of :class:`~hyper_parallel.platform.torch
390 .activation_checkpoint.activation_swap.SwapWrapper`.
392 Wraps a :class:`~mindspore.nn.Cell` and applies async activation swap
393 during the forward pass via the platform's ``async_save_on_cpu`` context
394 manager. Falls back to a no-op context when that context is not yet
395 available on the current platform.
397 Args:
398 mod (Cell): The cell whose intermediate activations should be swapped.
399 policy_fn (callable, optional): Per-tensor swap policy; see
400 :class:`AsyncSaveOnCpu`.
402 Example:
403 >>> from hyper_parallel.platform.mindspore.activation_checkpoint import swap_wrapper
404 >>> model.layers[i].attn = swap_wrapper(model.layers[i].attn, policy_fn)
405 """
407 def __init__(
408 self,
409 mod: Union[Cell, Callable],
410 policy_fn: Optional[Callable] = None,
411 group_swap: bool = False,
412 ):
413 super().__init__(mod)
414 self.policy_fn = policy_fn
415 self.group_swap = group_swap
417 def construct(self, *args, **kwargs):
418 """Execute the wrapped module inside an async CPU-swap context."""
419 with AsyncSaveOnCpu(policy_fn=self.policy_fn, group_swap=self.group_swap):
420 return self._ckpt_wrapped_module(*args, **kwargs)
423def swap_wrapper(
424 module: Union[Cell, Callable],
425 policy_fn: Optional[Callable] = None,
426 group_swap: bool = False,
427) -> SwapWrapper:
428 """
429 Wrap *module* with async activation swap.
431 Args:
432 module (Cell or callable): The cell or plain function to wrap.
433 If a plain callable is passed it is automatically wrapped in a
434 :class:`FuncCell` before being stored.
435 policy_fn (callable, optional): Per-tensor swap policy; see
436 :class:`AsyncSaveOnCpu`.
438 Returns:
439 SwapWrapper: The wrapped cell with activation swap enabled.
440 """
441 return SwapWrapper(module, policy_fn, group_swap)
444def swap_tensor_wrapper(target, tag: Optional[str] = None, group_swap: bool = False):
445 """Register selected tensors into the current swap group.
447 This helper is intended to be used inside a forward path that already
448 participates in the existing swap scheduling managed by ``SwapManager``.
449 It preserves the input structure and returns the original tensors.
450 """
451 # pylint: disable=C0415
452 from hyper_parallel.core.activation_checkpoint.swap import Storage, SwapManager, SwapTensor
453 swap_manager = SwapManager()
454 group_name = swap_manager.get_current_group_name()
455 if not group_name:
456 warnings.warn(
457 f"Tensor {tag} cannot be swapped, for its group is unregistered."
458 )
459 return target
460 if swap_manager.is_last_group(group_name):
461 return target
463 storage = Storage()
464 count_idx = 0
466 def _apply(x):
467 nonlocal count_idx
468 if isinstance(x, Tensor) and base_check_fn(x):
469 tensor_tag = tag or f"{group_name}_swap_tensor"
470 funcname = f"{tensor_tag}::{tuple(x.shape)}"
471 storage[count_idx].append(SwapTensor(x, funcname, group_swap=group_swap))
472 count_idx += 1
473 return x
475 def _map(tree):
476 if isinstance(tree, dict):
477 return type(tree)((k, _map(v)) for k, v in tree.items())
478 if isinstance(tree, tuple):
479 return tuple(_map(v) for v in tree)
480 if isinstance(tree, list):
481 return [_map(v) for v in tree]
482 return _apply(tree)
484 wrapped = _map(target)
485 if count_idx > 0:
486 swap_manager.add_storage(group_name, storage)
487 return wrapped