trainer_recompute-hyper-parallel-AtomGit

Mexamples/training_demo/train.yaml+2-2

	@@ -95,8 +95,8 @@ fsdp_config:
95	mixed_precision:	95	mixed_precision:
96	enabled: false	96	enabled: false
97		97
98	-gradient_checkpointing:	98	+activation_checkpoint:
99	- activation_checkpoint: full	99	+ mode: off # off, full, selective
100		100
101	model_assets:	101	model_assets:
102	datasets_type: pretokenized	102	datasets_type: pretokenized

Mhyper_models/_transformers/auto_model.py+9-2

	@@ -82,6 +82,7 @@ class _BaseHyperAutoModelClass:
82	fp8_config=None,	82	fp8_config=None,
83	compile_config=None,	83	compile_config=None,
84	freeze_config=None,	84	freeze_config=None,
		85	+ activation_checkpoint: Optional[str] = None,
85	**kwargs,	86	**kwargs,
86	) -> PreTrainedModel:	87	) -> PreTrainedModel:
87	"""HF-compatible from_pretrained entry point.	88	"""HF-compatible from_pretrained entry point.
	@@ -134,6 +135,7 @@ class _BaseHyperAutoModelClass:
134	fp8_config=fp8_config,	135	fp8_config=fp8_config,
135	compile_config=compile_config,	136	compile_config=compile_config,
136	freeze_config=freeze_config,	137	freeze_config=freeze_config,
		138	+ activation_checkpoint=activation_checkpoint,
137	**kwargs,	139	**kwargs,
138	)	140	)
139		141
	@@ -147,6 +149,7 @@ class _BaseHyperAutoModelClass:
147	backend=None,	149	backend=None,
148	torch_dtype="auto",	150	torch_dtype="auto",
149	attn_implementation="sdpa",	151	attn_implementation="sdpa",
		152	+ activation_checkpoint: Optional[str] = None,
150	**kwargs,	153	**kwargs,
151	) -> PreTrainedModel:	154	) -> PreTrainedModel:
152	"""Build model from PretrainedConfig (no weight loading).	155	"""Build model from PretrainedConfig (no weight loading).
	@@ -181,6 +184,7 @@ class _BaseHyperAutoModelClass:
181	validate_placement=False,	184	validate_placement=False,
182	load_base_model=False,	185	load_base_model=False,
183	distributed_setup=distributed_setup,	186	distributed_setup=distributed_setup,
		187	+ activation_checkpoint=activation_checkpoint,
184	**kwargs,	188	**kwargs,
185	)	189	)
186		190
	@@ -206,6 +210,7 @@ class _BaseHyperAutoModelClass:
206	fp8_config=None,	210	fp8_config=None,
207	compile_config=None,	211	compile_config=None,
208	freeze_config=None,	212	freeze_config=None,
		213	+ activation_checkpoint: Optional[str] = None,
209	**kwargs,	214	**kwargs,
210	) -> PreTrainedModel:	215	) -> PreTrainedModel:
211	"""Core model building orchestration.	216	"""Core model building orchestration.
	@@ -213,7 +218,8 @@ class _BaseHyperAutoModelClass:
213	Following design doc 01 §6.3:	218	Following design doc 01 §6.3:
214	Step 1: Determine meta device	219	Step 1: Determine meta device
215	Step 2: Build model (meta or real device)	220	Step 2: Build model (meta or real device)
216	- Step 3-11: apply_model_infrastructure (PP, PEFT, QAT, ShardingPlan, apply, compile, FSDP2, load)	221	+ Step 3-12: apply_model_infrastructure (PP, PEFT, QAT, ShardingPlan,
		222	+ activation checkpoint, compile, FSDP2, load)
217	"""	223	"""
218	from contextlib import nullcontext	224	from contextlib import nullcontext
219	from transformers.modeling_utils import ContextManagers	225	from transformers.modeling_utils import ContextManagers
	@@ -246,7 +252,7 @@ class _BaseHyperAutoModelClass:
246	**kwargs,	252	**kwargs,
247	)	253	)
248		254
249	- # Step 3-11: Apply infrastructure	255	+ # Step 3-12: Apply infrastructure
250	model = apply_model_infrastructure(	256	model = apply_model_infrastructure(
251	model,	257	model,
252	mesh=mesh,	258	mesh=mesh,
	@@ -265,6 +271,7 @@ class _BaseHyperAutoModelClass:
265	pretrained_path=pretrained_model_name_or_path,	271	pretrained_path=pretrained_model_name_or_path,
266	validate_placement=validate_placement,	272	validate_placement=validate_placement,
267	distributed_setup=distributed_setup,	273	distributed_setup=distributed_setup,
		274	+ activation_checkpoint=activation_checkpoint,
268	)	275	)
269		276
270	model.train()	277	model.train()

Mhyper_models/_transformers/infrastructure.py+30-3

	@@ -31,6 +31,9 @@ from hyper_models.components.distributed.fsdp2 import FSDP2Manager, _instantiate
31	from hyper_models.components.distributed.pipelining import _instantiate_pipeline	31	from hyper_models.components.distributed.pipelining import _instantiate_pipeline
32	from hyper_models.components.distributed.sharding_applier import apply_sharding_plan	32	from hyper_models.components.distributed.sharding_applier import apply_sharding_plan
33	from hyper_models.components.distributed.sharding_planner import ShardingPlanner	33	from hyper_models.components.distributed.sharding_planner import ShardingPlanner
		34	+from hyper_models.components.distributed.activation_checkpointing import (
		35	+ _apply_activation_checkpointing as _apply_activation_checkpointing_impl,
		36	+)
34	from hyper_models.trainer.config import (	37	from hyper_models.trainer.config import (
35	entries_to_module_replacements,	38	entries_to_module_replacements,
36	entries_to_plan_overrides,	39	entries_to_plan_overrides,
	@@ -63,6 +66,19 @@ class _TargetSnapshot:
63	layout_id: int \| None	66	layout_id: int \| None
64		67
65		68
		69	+def _apply_activation_checkpointing(
		70	+ model: nn.Module,
		71	+ activation_checkpoint: Optional[str],
		72	+ enable_compile: bool = False,
		73	+) -> nn.Module:
		74	+ """Compatibility wrapper for the distributed activation checkpoint helper."""
		75	+ return _apply_activation_checkpointing_impl(
		76	+ model,
		77	+ activation_checkpoint,
		78	+ enable_compile=enable_compile,
		79	+ )
		80	+
		81	+
66	def instantiate_infrastructure(	82	def instantiate_infrastructure(
67	distributed_setup=None,	83	distributed_setup=None,
68	device=None,	84	device=None,
	@@ -509,6 +525,7 @@ def apply_model_infrastructure(
509	fp8_config=None,	525	fp8_config=None,
510	freeze_config=None,	526	freeze_config=None,
511	compile_config=None,	527	compile_config=None,
		528	+ activation_checkpoint: Optional[str] = None,
512	is_meta_device: bool = False,	529	is_meta_device: bool = False,
513	is_hf_model: bool = False,	530	is_hf_model: bool = False,
514	device=None,	531	device=None,
	@@ -521,7 +538,8 @@ def apply_model_infrastructure(
521		538
522	Following design doc 01 §8.3 canonical order:	539	Following design doc 01 §8.3 canonical order:
523	PP split → PEFT → QAT/FP8 → freeze → plan → apply_sharding_plan	540	PP split → PEFT → QAT/FP8 → freeze → plan → apply_sharding_plan
524	- → torch.compile → FSDP2 wrap → to_empty + load_base_model	541	+ → activation checkpoint → torch.compile → FSDP2 wrap
		542	+ → to_empty + load_base_model
525		543
526	D-01'': CP K/V all-gather is injected at apply_sharding_plan time;	544	D-01'': CP K/V all-gather is injected at apply_sharding_plan time;
527	no extra CP hooks are registered here.	545	no extra CP hooks are registered here.
	@@ -561,11 +579,20 @@ def apply_model_infrastructure(
561	validate_placement,	579	validate_placement,
562	)	580	)
563		581
564	- # Step 9: torch.compile	582	+ # Step 9: activation checkpointing. The whole-layer wrapper must enclose
		583	+ # TP/EP transformations and remain inside the FSDP boundary.
		584	+ if activation_checkpoint not in (None, "off"):
		585	+ model = _apply_activation_checkpointing(
		586	+ model,
		587	+ activation_checkpoint,
		588	+ enable_compile=compile_config is not None,
		589	+ )
		590	+
		591	+ # Step 10: torch.compile
565	if compile_config is not None:	592	if compile_config is not None:
566	model = torch.compile(model, **compile_config)	593	model = torch.compile(model, **compile_config)
567		594
568	- # Steps 10-11: FSDP2 wrap, then materialize/move model storage.	595	+ # Steps 11-12: FSDP2 wrap, then materialize/move model storage.
569	model = _apply_fsdp2(model, fsdp2_manager, tp_grad_info)	596	model = _apply_fsdp2(model, fsdp2_manager, tp_grad_info)
570		597
571	# init model weights (either from base model or from model's own init)	598	# init model weights (either from base model or from model's own init)

Ahyper_models/components/distributed/activation_checkpointing.py+760-0

Mhyper_models/trainer/base.py+1-3

	@@ -363,15 +363,13 @@ class BaseTrainer(Stateful, ABC):
363	if self.global_rank == 0:	363	if self.global_rank == 0:
364	save_configs(self.config, self.config.checkpoint.checkpoint_dir)	364	save_configs(self.config, self.config.checkpoint.checkpoint_dir)
365		365
366	- # # Gradient checkpointing debug
367	- # set_checkpoint_debug_enabled(self.config.gradient_checkpointing.debug)
368	-
369	def _build_model(self) -> None:	366	def _build_model(self) -> None:
370	"""Build the model and derive Trainer-owned runtime state."""	367	"""Build the model and derive Trainer-owned runtime state."""
371	self.peft_config = self.config.peft	368	self.peft_config = self.config.peft
372	self.model = self.config.model.build(	369	self.model = self.config.model.build(
373	distributed_setup=self.distributed_setup,	370	distributed_setup=self.distributed_setup,
374	peft_config=self.peft_config,	371	peft_config=self.peft_config,
		372	+ activation_checkpoint=self.config.activation_checkpoint.mode,
375	)	373	)
376	self.model_config = self.model.config	374	self.model_config = self.model.config
377	model_parts = getattr(self.model, "parts", None)	375	model_parts = getattr(self.model, "parts", None)

Mhyper_models/trainer/config.py+5-5

	@@ -79,10 +79,10 @@ class MixedPrecisionConfig:
79		79
80		80
81	@dataclass	81	@dataclass
82	-class GradientCheckpointingConfig:	82	+class ActivationCheckpointConfig:
83	"""Activation-checkpoint mode exposed by the initial YAML schema."""	83	"""Activation-checkpoint mode exposed by the initial YAML schema."""
84		84
85	- activation_checkpoint: Literal["off", "none", "full", "selective"] = "off"	85	+ mode: Optional[Literal["off", "full", "selective"]] = "off"
86		86
87		87
88	@dataclass	88	@dataclass
	@@ -579,8 +579,8 @@ class TrainerConfig:
579	mixed_precision: MixedPrecisionConfig = field(	579	mixed_precision: MixedPrecisionConfig = field(
580	default_factory=MixedPrecisionConfig	580	default_factory=MixedPrecisionConfig
581	)	581	)
582	- gradient_checkpointing: GradientCheckpointingConfig = field(	582	+ activation_checkpoint: ActivationCheckpointConfig = field(
583	- default_factory=GradientCheckpointingConfig	583	+ default_factory=ActivationCheckpointConfig
584	)	584	)
585		585
586	# data	586	# data
	@@ -621,9 +621,9 @@ def save_configs(config: TrainerConfig, output_dir: str) -> None:
621		621
622	__all__ = [	622	__all__ = [
623	"AcceleratorConfig",	623	"AcceleratorConfig",
		624	+ "ActivationCheckpointConfig",
624	"DebugConfig",	625	"DebugConfig",
625	"FSDP2Config",	626	"FSDP2Config",
626	- "GradientCheckpointingConfig",
627	"MixedPrecisionConfig",	627	"MixedPrecisionConfig",
628	"Target",	628	"Target",
629	"TrainerConfig",	629	"TrainerConfig",

Mhyper_parallel/platform/mindspore/activation_checkpoint/sac.py+10-0

	@@ -68,6 +68,16 @@ class SelectiveCheckpointContext:
68	SAC_IGNORED_OPS = {"StopGradient"}	68	SAC_IGNORED_OPS = {"StopGradient"}
69		69
70		70
		71	+def ignore_sac_ops(ops: List[Optional[object]]) -> None:
		72	+ """Add available operator names to the selective-checkpoint ignore set.
		73	+
		74	+ Args:
		75	+ ops: MindSpore operator names to execute without selective-AC replay
		76	+ accounting. ``None`` entries are ignored.
		77	+ """
		78	+ SAC_IGNORED_OPS.update(op for op in ops if op is not None)
		79	+
		80	+
71	class _CachingMindSporeDispatchMode(MsDispatchMode):	81	class _CachingMindSporeDispatchMode(MsDispatchMode):
72	def __init__(self, policy_fn, swap_storage, storage, group_swap=False):	82	def __init__(self, policy_fn, swap_storage, storage, group_swap=False):
73	self.policy_fn = policy_fn	83	self.policy_fn = policy_fn

Mhyper_parallel/platform/mindspore/platform.py+6-0

	@@ -1844,6 +1844,12 @@ class MindSporePlatform(Platform):
1844	allow_cache_entry_mutation=allow_cache_entry_mutation,	1844	allow_cache_entry_mutation=allow_cache_entry_mutation,
1845	group_swap=group_swap)	1845	group_swap=group_swap)
1846		1846
		1847	+ @staticmethod
		1848	+ def ignore_sac_ops(ops: list[object \| None]) -> None:
		1849	+ # pylint: disable=C0415
		1850	+ from hyper_parallel.platform.mindspore.activation_checkpoint.sac import ignore_sac_ops
		1851	+ ignore_sac_ops(ops)
		1852	+
1847	@staticmethod	1853	@staticmethod
1848	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):	1854	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1849	# pylint: disable=C0415	1855	# pylint: disable=C0415

Mhyper_parallel/platform/mindspore/platform_graph.py+4-0

	@@ -52,6 +52,10 @@ class MindSporeGraphPlatform(MindSporePlatform):
52	def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):	52	def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False):
53	raise NotImplementedError("create_selective_checkpoint_contexts is not supported on MindSpore Graph platform")	53	raise NotImplementedError("create_selective_checkpoint_contexts is not supported on MindSpore Graph platform")
54		54
		55	+ @staticmethod
		56	+ def ignore_sac_ops(ops: list[object \| None]) -> None:
		57	+ raise NotImplementedError("ignore_sac_ops is not supported on MindSpore Graph platform")
		58	+
55	@staticmethod	59	@staticmethod
56	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):	60	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
57	raise NotImplementedError("async_save_on_cpu is not supported on MindSpore Graph platform")	61	raise NotImplementedError("async_save_on_cpu is not supported on MindSpore Graph platform")

Mhyper_parallel/platform/platform.py+10-0

	@@ -1503,6 +1503,16 @@ class Platform:
1503	"""	1503	"""
1504	raise NotImplementedError("Platform subclasses must implement create_selective_checkpoint_contexts")	1504	raise NotImplementedError("Platform subclasses must implement create_selective_checkpoint_contexts")
1505		1505
		1506	+ @staticmethod
		1507	+ def ignore_sac_ops(ops: list[object \| None]) -> None:
		1508	+ """Exclude backend operators from selective-checkpoint replay accounting.
		1509	+
		1510	+ Args:
		1511	+ ops: Iterable of backend-native operator identifiers. Unavailable
		1512	+ optional operators may be represented by ``None``.
		1513	+ """
		1514	+ raise NotImplementedError("Platform subclasses must implement ignore_sac_ops")
		1515	+
1506	@staticmethod	1516	@staticmethod
1507	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):	1517	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1508	"""Create an async CPU offload context for activation checkpointing.	1518	"""Create an async CPU offload context for activation checkpointing.

Mhyper_parallel/platform/torch/activation_checkpoint/sac.py+10-0

	@@ -130,6 +130,16 @@ SAC_IGNORED_OPS = {
130	} \| set(torch._subclasses.functional_tensor.FunctionalTensor.metadata_fns)	130	} \| set(torch._subclasses.functional_tensor.FunctionalTensor.metadata_fns)
131		131
132		132
		133	+def ignore_sac_ops(ops: List[Optional[object]]) -> None:
		134	+ """Add available operators to the selective-checkpoint ignore set.
		135	+
		136	+ Args:
		137	+ ops: Operators to execute without selective-AC replay accounting.
		138	+ ``None`` entries are ignored for optional-version compatibility.
		139	+ """
		140	+ SAC_IGNORED_OPS.update(op for op in ops if op is not None)
		141	+
		142	+
133	class _CachingTorchDispatchMode(TorchDispatchMode):	143	class _CachingTorchDispatchMode(TorchDispatchMode):
134	# Used together with _CachedTorchDispatchMode to implement SAC.	144	# Used together with _CachedTorchDispatchMode to implement SAC.
135	def __init__(self, policy_fn, swap_storage, storage, group_swap=False):	145	def __init__(self, policy_fn, swap_storage, storage, group_swap=False):

Mhyper_parallel/platform/torch/platform.py+6-0

	@@ -1509,6 +1509,12 @@ class TorchPlatform(Platform):
1509	from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts	1509	from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts
1510	return create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation, group_swap)	1510	return create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation, group_swap)
1511		1511
		1512	+ @staticmethod
		1513	+ def ignore_sac_ops(ops: list[object \| None]) -> None:
		1514	+ # pylint: disable=C0415
		1515	+ from hyper_parallel.platform.torch.activation_checkpoint.sac import ignore_sac_ops
		1516	+ ignore_sac_ops(ops)
		1517	+
1512	@staticmethod	1518	@staticmethod
1513	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):	1519	def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
1514	# pylint: disable=C0415	1520	# pylint: disable=C0415

Atests/hyper_models/trainer/test_activation_checkpointing.py+458-0

@@ -0,0 +1,458 @@
	1	+# Copyright 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	+"""Tests for activation checkpointing in model infrastructure."""
	16	+
	17	+from types import MethodType, SimpleNamespace
	18	+from unittest.mock import patch
	19	+
	20	+import pytest
	21	+import torch
	22	+import torch.nn as nn
	23	+import yaml
	24	+from torch.utils.checkpoint import checkpoint
	25	+from transformers.modeling_layers import GradientCheckpointingLayer
	26	+
	27	+from hyper_models._transformers.infrastructure import (
	28	+ _apply_activation_checkpointing,
	29	+ apply_model_infrastructure,
	30	+)
	31	+from hyper_models.components.distributed.activation_checkpointing import (
	32	+ _find_transformer_layer_container_infos,
	33	+ make_selective_checkpoint_context_fn,
	34	+)
	35	+from hyper_models.components.distributed.fsdp2 import FSDP2Manager
	36	+from hyper_models.config.resolver import resolve_component
	37	+from hyper_models.trainer.config import ActivationCheckpointConfig
	38	+
	39	+
	40	+class _CountingLinear(nn.Linear):
	41	+ def __init__(self) -> None:
	42	+ super().__init__(4, 4)
	43	+ self.forward_calls = 0
	44	+
	45	+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	46	+ self.forward_calls += 1
	47	+ return super().forward(hidden_states)
	48	+
	49	+
	50	+class _Layer(nn.Module):
	51	+ def __init__(self) -> None:
	52	+ super().__init__()
	53	+ self.mlp = _CountingLinear()
	54	+ self.forward_calls = 0
	55	+
	56	+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	57	+ self.forward_calls += 1
	58	+ return self.mlp(hidden_states)
	59	+
	60	+
	61	+class _Backbone(nn.Module):
	62	+ def __init__(self, num_layers: int = 2) -> None:
	63	+ super().__init__()
	64	+ self.layers = nn.ModuleList(_Layer() for _ in range(num_layers))
	65	+
	66	+
	67	+class _Model(nn.Module):
	68	+ def __init__(self, num_layers: int = 2) -> None:
	69	+ super().__init__()
	70	+ self.model = _Backbone(num_layers)
	71	+ self.config = SimpleNamespace(use_cache=True)
	72	+
	73	+
	74	+def test_activation_checkpoint_config_accepts_off_full_or_selective() -> None:
	75	+ disabled = resolve_component(
	76	+ yaml.safe_load("mode: off"),
	77	+ expected_type=ActivationCheckpointConfig,
	78	+ path="$.activation_checkpoint",
	79	+ )
	80	+ null_disabled = resolve_component(
	81	+ {"mode": None},
	82	+ expected_type=ActivationCheckpointConfig,
	83	+ path="$.activation_checkpoint",
	84	+ )
	85	+ enabled = resolve_component(
	86	+ {"mode": "full"},
	87	+ expected_type=ActivationCheckpointConfig,
	88	+ path="$.activation_checkpoint",
	89	+ )
	90	+
	91	+ assert disabled.mode == "off"
	92	+ assert null_disabled.mode is None
	93	+ assert enabled.mode == "full"
	94	+ selective = resolve_component(
	95	+ {"mode": "selective"},
	96	+ expected_type=ActivationCheckpointConfig,
	97	+ path="$.activation_checkpoint",
	98	+ )
	99	+ assert selective.mode == "selective"
	100	+
	101	+
	102	+@pytest.mark.parametrize("mode", [None, "off"])
	103	+def test_disabled_activation_checkpoint_does_not_apply(mode: str \| None) -> None:
	104	+ model = _Model()
	105	+ original_layers = list(model.model.layers)
	106	+
	107	+ with patch(
	108	+ "hyper_models._transformers.infrastructure._apply_activation_checkpointing"
	109	+ ) as apply_checkpoint:
	110	+ result = apply_model_infrastructure(model, activation_checkpoint=mode)
	111	+
	112	+ assert result is model
	113	+ apply_checkpoint.assert_not_called()
	114	+ assert list(model.model.layers) == original_layers
	115	+ assert model.config.use_cache is True
	116	+
	117	+
	118	+def test_full_activation_checkpoint_wraps_layer_submodules() -> None:
	119	+ model = _Model(num_layers=3)
	120	+ original_layers = list(model.model.layers)
	121	+ original_mlps = [layer.mlp for layer in original_layers]
	122	+
	123	+ result = _apply_activation_checkpointing(model, "full")
	124	+
	125	+ assert result is model
	126	+ assert model.config.use_cache is False
	127	+ assert list(model.model.layers) == original_layers
	128	+ for layer, original_mlp in zip(model.model.layers, original_mlps):
	129	+ assert layer.mlp._wrapped_module is original_mlp
	130	+
	131	+
	132	+def test_full_activation_checkpoint_recomputes_submodules_during_backward() -> None:
	133	+ model = _Model(num_layers=2)
	134	+ original_layers = list(model.model.layers)
	135	+ original_mlps = [layer.mlp for layer in original_layers]
	136	+ _apply_activation_checkpointing(model, "full")
	137	+
	138	+ hidden_states = torch.randn(2, 4, requires_grad=True)
	139	+ for layer in model.model.layers:
	140	+ hidden_states = layer(hidden_states)
	141	+ hidden_states.sum().backward()
	142	+
	143	+ assert [layer.forward_calls for layer in original_layers] == [1, 1]
	144	+ assert [mlp.forward_calls for mlp in original_mlps] == [2, 2]
	145	+
	146	+
	147	+def test_selective_activation_checkpoint_wraps_layers_with_context_fn() -> None:
	148	+ model = _Model(num_layers=2)
	149	+ original_layers = list(model.model.layers)
	150	+
	151	+ _apply_activation_checkpointing(model, "selective")
	152	+
	153	+ assert model.config.use_cache is False
	154	+ for wrapped, original in zip(model.model.layers, original_layers):
	155	+ assert wrapped._wrapped_module is original
	156	+ assert callable(wrapped.checkpoint_kwargs["context_fn"])
	157	+
	158	+ hidden_states = torch.randn(2, 4, requires_grad=True)
	159	+ for layer in model.model.layers:
	160	+ hidden_states = layer(hidden_states)
	161	+ hidden_states.sum().backward()
	162	+
	163	+ assert [layer.forward_calls for layer in original_layers] == [2, 2]
	164	+
	165	+
	166	+def test_selective_checkpoint_recomputes_mutated_topk_output() -> None:
	167	+ def router_loss(hidden_states: torch.Tensor) -> torch.Tensor:
	168	+ router_probs = torch.softmax(hidden_states, dim=-1)
	169	+ routing_weights, selected_experts = torch.topk(router_probs, 2, dim=-1)
	170	+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
	171	+ selected_states = torch.gather(hidden_states, -1, selected_experts)
	172	+ return (routing_weights * selected_states).sum()
	173	+
	174	+ baseline_input = torch.randn(3, 8, requires_grad=True)
	175	+ checkpoint_input = baseline_input.detach().clone().requires_grad_(True)
	176	+
	177	+ baseline_loss = router_loss(baseline_input)
	178	+ baseline_loss.backward()
	179	+ checkpoint_loss = checkpoint(
	180	+ router_loss,
	181	+ checkpoint_input,
	182	+ use_reentrant=False,
	183	+ context_fn=make_selective_checkpoint_context_fn(),
	184	+ )
	185	+ checkpoint_loss.backward()
	186	+
	187	+ torch.testing.assert_close(checkpoint_loss, baseline_loss)
	188	+ torch.testing.assert_close(checkpoint_input.grad, baseline_input.grad)
	189	+
	190	+
	191	+def test_selective_context_registers_profiler_and_fsdp_runtime_ops() -> None:
	192	+ with (
	193	+ patch(
	194	+ "hyper_models.components.distributed.activation_checkpointing."
	195	+ "ensure_profiler_ops_sac_ignored"
	196	+ ) as ensure_profiler_ops,
	197	+ patch(
	198	+ "hyper_models.components.distributed.activation_checkpointing."
	199	+ "ensure_fsdp_ops_sac_ignored"
	200	+ ) as ensure_fsdp_ops,
	201	+ ):
	202	+ context_fn = make_selective_checkpoint_context_fn()
	203	+
	204	+ ensure_profiler_ops.assert_called_once_with()
	205	+ ensure_fsdp_ops.assert_called_once_with()
	206	+ assert callable(context_fn)
	207	+
	208	+
	209	+def test_selective_checkpointing_falls_back_for_kv_sharing(caplog) -> None:
	210	+ class _KVLayer(nn.Module):
	211	+ def __init__(self) -> None:
	212	+ super().__init__()
	213	+ self.self_attn = nn.Linear(4, 4)
	214	+ self.mlp = nn.Linear(4, 4)
	215	+ self.input_layernorm = nn.LayerNorm(4)
	216	+ self.post_attention_layernorm = nn.LayerNorm(4)
	217	+
	218	+ model = _Model(num_layers=0)
	219	+ model.model.layers = nn.ModuleList([_KVLayer()])
	220	+ model.config = SimpleNamespace(
	221	+ use_cache=True,
	222	+ text_config=SimpleNamespace(num_kv_shared_layers=1, use_cache=True),
	223	+ )
	224	+ layer = model.model.layers[0]
	225	+ original_attention = layer.self_attn
	226	+ wrapped_submodules = []
	227	+
	228	+ with (
	229	+ caplog.at_level("WARNING"),
	230	+ patch(
	231	+ "hyper_models.components.distributed.activation_checkpointing.checkpoint_wrapper",
	232	+ side_effect=lambda submodule: wrapped_submodules.append(submodule)
	233	+ or submodule,
	234	+ ),
	235	+ ):
	236	+ _apply_activation_checkpointing(model, "selective")
	237	+
	238	+ assert model.model.layers[0] is layer
	239	+ assert layer.self_attn is original_attention
	240	+ assert wrapped_submodules == [
	241	+ layer.mlp,
	242	+ layer.input_layernorm,
	243	+ layer.post_attention_layernorm,
	244	+ ]
	245	+ assert model.config.use_cache is True
	246	+ assert model.config.text_config.use_cache is True
	247	+ assert "falling back to submodule activation checkpointing" in caplog.text
	248	+
	249	+
	250	+def test_layer_container_discovery_finds_all_known_vlm_towers_once() -> None:
	251	+ class Qwen2VLForConditionalGeneration(nn.Module):
	252	+ def __init__(self) -> None:
	253	+ super().__init__()
	254	+ self.model = nn.Module()
	255	+ self.model.language_model = nn.Module()
	256	+ self.model.language_model.layers = nn.ModuleList([_Layer(), _Layer()])
	257	+ self.model.visual = nn.Module()
	258	+ self.model.visual.blocks = nn.ModuleList([_Layer()])
	259	+
	260	+ # Deprecated transformers aliases point to the same modules. The
	261	+ # canonical candidates must win without double-counting containers.
	262	+ self.language_model = self.model.language_model
	263	+ self.visual = self.model.visual
	264	+
	265	+ model = Qwen2VLForConditionalGeneration()
	266	+
	267	+ groups = _find_transformer_layer_container_infos(model)
	268	+ assert {group_name: len(containers[0][2]) for group_name, containers in groups.items()} == {
	269	+ "language": 2,
	270	+ "vision": 1,
	271	+ }
	272	+ assert groups["language"][0][2] is model.model.language_model.layers
	273	+ assert groups["vision"][0][2] is model.model.visual.blocks
	274	+
	275	+
	276	+def test_known_model_with_stale_paths_does_not_use_heuristic(caplog) -> None:
	277	+ class Qwen2VLForConditionalGeneration(nn.Module):
	278	+ def __init__(self) -> None:
	279	+ super().__init__()
	280	+ self.unrelated_modules = nn.ModuleList([_Layer(), _Layer()])
	281	+
	282	+ model = Qwen2VLForConditionalGeneration()
	283	+
	284	+ with caplog.at_level("WARNING"):
	285	+ containers = _find_transformer_layer_container_infos(model)
	286	+
	287	+ assert containers == {}
	288	+ assert "resolved no modules" in caplog.text
	289	+
	290	+
	291	+def test_unknown_model_uses_numeric_module_dict_heuristic() -> None:
	292	+ class UnknownModel(nn.Module):
	293	+ def __init__(self) -> None:
	294	+ super().__init__()
	295	+ self.adapters = nn.ModuleDict({"default": _Layer()})
	296	+ self.blocks = nn.ModuleDict({"0": _Layer(), "2": _Layer()})
	297	+
	298	+ model = UnknownModel()
	299	+ original_layers = list(model.blocks.values())
	300	+
	301	+ groups = _find_transformer_layer_container_infos(model)
	302	+ assert list(groups) == ["unknown"]
	303	+ assert len(groups["unknown"]) == 1
	304	+ assert list(groups["unknown"][0][2].values()) == original_layers
	305	+
	306	+ _apply_activation_checkpointing(model, "full")
	307	+
	308	+ assert list(model.blocks.values()) == original_layers
	309	+ for layer in original_layers:
	310	+ assert hasattr(layer.mlp, "_wrapped_module")
	311	+
	312	+
	313	+class _HFGradientCheckpointingLayer(GradientCheckpointingLayer):
	314	+ def __init__(self) -> None:
	315	+ super().__init__()
	316	+ self.mlp = nn.Linear(4, 4)
	317	+
	318	+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	319	+ return self.mlp(hidden_states)
	320	+
	321	+
	322	+class _HFGradientCheckpointingModel(nn.Module):
	323	+ supports_gradient_checkpointing = True
	324	+
	325	+ def __init__(self) -> None:
	326	+ super().__init__()
	327	+ self.model = nn.Module()
	328	+ self.model.layers = nn.ModuleList([_HFGradientCheckpointingLayer()])
	329	+ self.config = SimpleNamespace(use_cache=True)
	330	+ self.native_gradient_checkpointing_kwargs = None
	331	+
	332	+ def gradient_checkpointing_enable(self, **kwargs) -> None:
	333	+ self.native_gradient_checkpointing_kwargs = kwargs
	334	+
	335	+
	336	+def test_full_checkpointing_uses_hf_native_api_for_eligible_language_model() -> None:
	337	+ model = _HFGradientCheckpointingModel()
	338	+
	339	+ with patch(
	340	+ "hyper_models.components.distributed.activation_checkpointing.checkpoint_wrapper"
	341	+ ) as wrapper:
	342	+ _apply_activation_checkpointing(model, "full")
	343	+
	344	+ assert model.native_gradient_checkpointing_kwargs == {
	345	+ "gradient_checkpointing_kwargs": {"use_reentrant": True}
	346	+ }
	347	+ wrapper.assert_not_called()
	348	+
	349	+
	350	+def test_compile_disables_hf_native_checkpointing() -> None:
	351	+ model = _HFGradientCheckpointingModel()
	352	+ wrapper_calls = []
	353	+ original_mlp = model.model.layers[0].mlp
	354	+
	355	+ with patch(
	356	+ "hyper_models.components.distributed.activation_checkpointing.checkpoint_wrapper",
	357	+ side_effect=lambda layer, **kwargs: wrapper_calls.append((layer, kwargs)) or layer,
	358	+ ):
	359	+ _apply_activation_checkpointing(model, "full", enable_compile=True)
	360	+
	361	+ assert model.native_gradient_checkpointing_kwargs is None
	362	+ assert wrapper_calls == [(original_mlp, {})]
	363	+
	364	+
	365	+def test_full_submodule_checkpointing_skips_attention_for_kv_sharing() -> None:
	366	+ class _KVLayer(nn.Module):
	367	+ def __init__(self) -> None:
	368	+ super().__init__()
	369	+ self.self_attn = nn.Linear(4, 4)
	370	+ self.mlp = nn.Linear(4, 4)
	371	+ self.input_layernorm = nn.LayerNorm(4)
	372	+ self.post_attention_layernorm = nn.LayerNorm(4)
	373	+
	374	+ model = _Model(num_layers=0)
	375	+ model.model.layers = nn.ModuleList([_KVLayer()])
	376	+ model.config = SimpleNamespace(
	377	+ use_cache=True,
	378	+ text_config=SimpleNamespace(num_kv_shared_layers=1, use_cache=True),
	379	+ )
	380	+ layer = model.model.layers[0]
	381	+ original_attention = layer.self_attn
	382	+ wrapped_submodules = []
	383	+
	384	+ with patch(
	385	+ "hyper_models.components.distributed.activation_checkpointing.checkpoint_wrapper",
	386	+ side_effect=lambda submodule: wrapped_submodules.append(submodule) or submodule,
	387	+ ):
	388	+ _apply_activation_checkpointing(model, "full")
	389	+
	390	+ assert layer.self_attn is original_attention
	391	+ assert wrapped_submodules == [
	392	+ layer.mlp,
	393	+ layer.input_layernorm,
	394	+ layer.post_attention_layernorm,
	395	+ ]
	396	+ assert model.config.use_cache is True
	397	+ assert model.config.text_config.use_cache is True
	398	+
	399	+
	400	+def test_activation_checkpoint_is_after_sharding_and_before_fsdp() -> None:
	401	+ model = _Model()
	402	+ events = []
	403	+ plan = SimpleNamespace(modules={"model.layers.0": object()})
	404	+ planner = SimpleNamespace(
	405	+ plan=lambda args, *kwargs: (events.append("sharding_plan") or plan)
	406	+ )
	407	+ mesh = SimpleNamespace(
	408	+ device_mesh=object(),
	409	+ tp_size=2,
	410	+ cp_size=1,
	411	+ ep_size=1,
	412	+ sequence_parallel=False,
	413	+ loss_parallel=False,
	414	+ )
	415	+ fsdp_manager = object.__new__(FSDP2Manager)
	416	+
	417	+ def _parallelize(self, wrapped_model, **kwargs):
	418	+ del self, kwargs
	419	+ events.append("fsdp")
	420	+ return wrapped_model
	421	+
	422	+ fsdp_manager.parallelize = MethodType(_parallelize, fsdp_manager)
	423	+
	424	+ def _apply_sharding(sharded_model, args, *kwargs):
	425	+ del args, kwargs
	426	+ events.append("sharding_apply")
	427	+ return sharded_model, {}
	428	+
	429	+ def _checkpoint_wrapper(layer):
	430	+ events.append("checkpoint")
	431	+ return layer
	432	+
	433	+ with (
	434	+ patch(
	435	+ "hyper_models.components.distributed.sharding_applier."
	436	+ "apply_sharding_plan",
	437	+ side_effect=_apply_sharding,
	438	+ ),
	439	+ patch(
	440	+ "hyper_models.components.distributed.activation_checkpointing.checkpoint_wrapper",
	441	+ side_effect=_checkpoint_wrapper,
	442	+ ),
	443	+ ):
	444	+ apply_model_infrastructure(
	445	+ model,
	446	+ mesh=mesh,
	447	+ sharding_planner=planner,
	448	+ fsdp2_manager=fsdp_manager,
	449	+ activation_checkpoint="full",
	450	+ )
	451	+
	452	+ assert events == [
	453	+ "sharding_plan",
	454	+ "sharding_apply",
	455	+ "checkpoint",
	456	+ "checkpoint",
	457	+ "fsdp",
	458	+ ]

Mtests/hyper_models/trainer/test_data_build.py+25-1

	@@ -26,7 +26,12 @@ from hyper_models.components.data import (
26	)	26	)
27	from hyper_models.components.distributed.infrastructure import MeshContext	27	from hyper_models.components.distributed.infrastructure import MeshContext
28	from hyper_models.trainer.base import BaseTrainer	28	from hyper_models.trainer.base import BaseTrainer
29	-from hyper_models.trainer.config import Target, TrainerConfig, TrainingConfig	29	+from hyper_models.trainer.config import (
		30	+ ActivationCheckpointConfig,
		31	+ Target,
		32	+ TrainerConfig,
		33	+ TrainingConfig,
		34	+)
30		35
31		36
32	def _unused_target() -> Target:	37	def _unused_target() -> Target:
	@@ -43,12 +48,14 @@ def _model_target(
43	*,	48	*,
44	distributed_setup: object,	49	distributed_setup: object,
45	peft_config: object,	50	peft_config: object,
		51	+ activation_checkpoint: object,
46	) -> SimpleNamespace:	52	) -> SimpleNamespace:
47	"""Return a model object for delegation testing."""	53	"""Return a model object for delegation testing."""
48	return SimpleNamespace(	54	return SimpleNamespace(
49	config=SimpleNamespace(model_type="fake"),	55	config=SimpleNamespace(model_type="fake"),
50	distributed_setup=distributed_setup,	56	distributed_setup=distributed_setup,
51	peft_config=peft_config,	57	peft_config=peft_config,
		58	+ activation_checkpoint=activation_checkpoint,
52	)	59	)
53		60
54		61
	@@ -75,6 +82,7 @@ def test_trainer_model_stage_delegates_to_config_target() -> None:
75	model=Target(_model_target, target_path="tests.model_target"),	82	model=Target(_model_target, target_path="tests.model_target"),
76	optimizer=_unused_target(),	83	optimizer=_unused_target(),
77	peft=peft_config,	84	peft=peft_config,
		85	+ activation_checkpoint=ActivationCheckpointConfig(mode="full"),
78	)	86	)
79	trainer.distributed_setup = distributed_setup	87	trainer.distributed_setup = distributed_setup
80		88
	@@ -83,11 +91,27 @@ def test_trainer_model_stage_delegates_to_config_target() -> None:
83	assert trainer.peft_config is peft_config	91	assert trainer.peft_config is peft_config
84	assert trainer.model.distributed_setup is distributed_setup	92	assert trainer.model.distributed_setup is distributed_setup
85	assert trainer.model.peft_config is peft_config	93	assert trainer.model.peft_config is peft_config
		94	+ assert trainer.model.activation_checkpoint == "full"
86	assert trainer.model_parts == [trainer.model]	95	assert trainer.model_parts == [trainer.model]
87	assert trainer.model_config is trainer.model.config	96	assert trainer.model_config is trainer.model.config
88	assert trainer.hsdp_model_parts == []	97	assert trainer.hsdp_model_parts == []
89		98
90		99
		100	+def test_trainer_model_stage_passes_off_activation_checkpointing_mode() -> None:
		101	+ """Pass the disabled activation-checkpoint mode through unchanged."""
		102	+ trainer = BaseTrainer.__new__(BaseTrainer)
		103	+ trainer.config = TrainerConfig(
		104	+ model=Target(_model_target, target_path="tests.model_target"),
		105	+ optimizer=_unused_target(),
		106	+ activation_checkpoint=ActivationCheckpointConfig(mode="off"),
		107	+ )
		108	+ trainer.distributed_setup = object()
		109	+
		110	+ trainer._build_model()
		111	+
		112	+ assert trainer.model.activation_checkpoint == "off"
		113	+
		114	+
91	def test_trainer_optimizer_stage_passes_runtime_context() -> None:	115	def test_trainer_optimizer_stage_passes_runtime_context() -> None:
92	"""Pass model, mesh, and PEFT state to the configured optimizer target."""	116	"""Pass model, mesh, and PEFT state to the configured optimizer target."""
93	model = object()	117	model = object()

Mtests/ut/platform/mindspore/activation_checkpoint/test_sac.py+12-0

	@@ -45,6 +45,7 @@ from hyper_parallel.platform.mindspore.activation_checkpoint.sac import (
45	_VersionWrapper,	45	_VersionWrapper,
46	_maybe_detach,	46	_maybe_detach,
47	create_selective_checkpoint_contexts,	47	create_selective_checkpoint_contexts,
		48	+ ignore_sac_ops,
48	)	49	)
49		50
50		51
	@@ -99,6 +100,17 @@ class TestSacHelpers(unittest.TestCase):
99	self.assertIs(entry.swap.val, cached)	100	self.assertIs(entry.swap.val, cached)
100	self.assertTrue(entry.swap.group_swap)	101	self.assertTrue(entry.swap.group_swap)
101		102
		103	+ def test_ignore_sac_ops_adds_available_operator_names(self):
		104	+ """Runtime operator names should be added while unavailable entries are omitted."""
		105	+ original_ignored_ops = set(sac.SAC_IGNORED_OPS)
		106	+ self.addCleanup(sac.SAC_IGNORED_OPS.update, original_ignored_ops)
		107	+ self.addCleanup(sac.SAC_IGNORED_OPS.intersection_update, original_ignored_ops)
		108	+
		109	+ ignore_sac_ops(["AllGather", None])
		110	+
		111	+ self.assertIn("AllGather", sac.SAC_IGNORED_OPS)
		112	+ self.assertNotIn(None, sac.SAC_IGNORED_OPS)
		113	+
102		114
103	class TestCreateSelectiveCheckpointContexts(unittest.TestCase):	115	class TestCreateSelectiveCheckpointContexts(unittest.TestCase):
104	"""Unit tests for create_selective_checkpoint_contexts()."""	116	"""Unit tests for create_selective_checkpoint_contexts()."""

Mtests/ut/platform/torch/activation_checkpoint/test_sac.py+20-0

	@@ -34,6 +34,7 @@ from hyper_parallel.platform.torch.activation_checkpoint.sac import (
34	_maybe_detach,	34	_maybe_detach,
35	_policy_from_bool,	35	_policy_from_bool,
36	create_selective_checkpoint_contexts,	36	create_selective_checkpoint_contexts,
		37	+ ignore_sac_ops,
37	)	38	)
38		39
39		40
	@@ -81,6 +82,25 @@ class TestSacHelpers(unittest.TestCase):
81	self.assertIs(entry.swap.val, cached)	82	self.assertIs(entry.swap.val, cached)
82	self.assertTrue(entry.swap.group_swap)	83	self.assertTrue(entry.swap.group_swap)
83		84
		85	+ def test_ignore_sac_ops_adds_available_operators(self):
		86	+ """Added runtime operators should execute without policy evaluation."""
		87	+ op = torch.ops.aten.add.Tensor
		88	+ original_ignored_ops = set(sac.SAC_IGNORED_OPS)
		89	+ self.addCleanup(sac.SAC_IGNORED_OPS.update, original_ignored_ops)
		90	+ self.addCleanup(sac.SAC_IGNORED_OPS.intersection_update, original_ignored_ops)
		91	+
		92	+ ignore_sac_ops([op, None])
		93	+ policy_fn = MagicMock(return_value=CheckpointPolicy.MUST_SAVE)
		94	+ caching, _ = create_selective_checkpoint_contexts(policy_fn)
		95	+ x = torch.tensor([1.0])
		96	+ with caching:
		97	+ result = torch.add(x, x)
		98	+
		99	+ self.assertIn(op, sac.SAC_IGNORED_OPS)
		100	+ self.assertNotIn(None, sac.SAC_IGNORED_OPS)
		101	+ self.assertTrue(torch.equal(result, torch.tensor([2.0])))
		102	+ policy_fn.assert_not_called()
		103	+
84		104
85	class TestCreateSelectiveCheckpointContexts(unittest.TestCase):	105	class TestCreateSelectiveCheckpointContexts(unittest.TestCase):
86	"""Unit tests for create_selective_checkpoint_contexts()."""	106	"""Unit tests for create_selective_checkpoint_contexts()."""