已开启
feat: 优化连续 SAVE 区域的输出显存 #1135
DavidFFFan创建于 8月9日
feat: 优化连续 SAVE 区域的输出显存 #1135
已开启
共 5 个文件变更+907-28
| @@ -14,55 +14,452 @@ | |||
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | """MindSpore wrapper for regions that should be saved instead of recomputed.""" | 15 | """MindSpore wrapper for regions that should be saved instead of recomputed.""" |
| 16 | from collections import defaultdict, deque | 16 | from collections import defaultdict, deque |
| 17 | -from typing import Any, Callable, Deque, Dict | 17 | +from contextvars import ContextVar |
| 18 | +from dataclasses import dataclass | ||
| 19 | +from functools import lru_cache | ||
| 20 | +from typing import Any, Callable, Deque, Dict, List, Tuple | ||
| 21 | + | ||
| 22 | +import mindspore as ms | ||
| 23 | +from mindspore.common._grad_function import _Function | ||
| 18 | 24 | ||
| 19 | from hyper_parallel.core.activation_checkpoint.recompute_state import get_recompute_state | 25 | from hyper_parallel.core.activation_checkpoint.recompute_state import get_recompute_state |
| 26 | +from hyper_parallel.core.dtensor.dtensor import DTensor | ||
| 20 | from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper | 27 | from hyper_parallel.platform.mindspore.activation_checkpoint.activation_swap import ActivationWrapper |
| 21 | 28 | ||
| 22 | 29 | ||
| 23 | -class _CheckpointExcludeCache: | 30 | +_InputPath = Tuple[Tuple[str, Any], ...] |
| 24 | - """Store excluded-region outputs for one checkpoint invocation.""" | 31 | +_TensorInput = Tuple[_InputPath, Any] |
| 32 | + | ||
| 33 | + | ||
| 34 | +@dataclass(frozen=True) | ||
| 35 | +class _TensorMetadata: | ||
| 36 | + """Alias-relevant metadata captured while an excluded call is running.""" | ||
| 37 | + | ||
| 38 | + tensor_id: int | ||
| 39 | + storage_ptr: int | ||
| 40 | + storage_nbytes: int | ||
| 41 | + dtype: Any | ||
| 42 | + shape: Tuple[int, ...] | ||
| 43 | + stride: Tuple[int, ...] | ||
| 44 | + storage_offset: int | ||
| 45 | + version: int | ||
| 46 | + itemsize: int | ||
| 47 | + numel: int | ||
| 48 | + is_contiguous: bool | ||
| 49 | + | ||
| 50 | + | ||
| 51 | +@dataclass(frozen=True) | ||
| 52 | +class _InputInfo: | ||
| 53 | + """Describe one canonical local input without retaining its tensor/storage.""" | ||
| 54 | + | ||
| 55 | + path: _InputPath | ||
| 56 | + metadata: _TensorMetadata | ||
| 57 | + | ||
| 58 | + | ||
| 59 | +@dataclass(frozen=True) | ||
| 60 | +class _ViewRecipe: | ||
| 61 | + """Rebuild one saved input alias from its matching replay input.""" | ||
| 62 | + | ||
| 63 | + shape: Tuple[int, ...] | ||
| 64 | + stride: Tuple[int, ...] | ||
| 65 | + relative_offset: int | ||
| 66 | + dtype: Any | ||
| 67 | + base_shape: Tuple[int, ...] | ||
| 68 | + base_stride: Tuple[int, ...] | ||
| 69 | + base_version: int | ||
| 70 | + base_is_contiguous: bool | ||
| 71 | + exact_input: bool | ||
| 72 | + | ||
| 73 | + | ||
| 74 | +class _RecomputedInputHandle: | ||
| 75 | + """Defer one excluded-region saved alias until checkpoint replay.""" | ||
| 76 | + | ||
| 77 | + def __init__(self, recipe: _ViewRecipe) -> None: | ||
| 78 | + """Initialize an unused and unresolved handle.""" | ||
| 79 | + self.recipe = recipe | ||
| 80 | + self._tensor = None | ||
| 81 | + | ||
| 82 | + def materialize(self, tensor: Any) -> None: | ||
| 83 | + """Bind the handle to the matching input produced during replay.""" | ||
| 84 | + self._tensor = tensor | ||
| 85 | + | ||
| 86 | + def get_recomputed_tensor(self) -> Any: | ||
| 87 | + """Return the replay-produced input for backward.""" | ||
| 88 | + if self._tensor is None: | ||
| 89 | + raise RuntimeError("Checkpoint-excluded input was requested before recomputation") | ||
| 90 | + return self._tensor | ||
| 91 | + | ||
| 92 | + | ||
| 93 | +@dataclass(frozen=True) | ||
| 94 | +class _InputBinding: | ||
| 95 | + """Map one replay input path to a deferred saved-alias handle.""" | ||
| 96 | + | ||
| 97 | + path: _InputPath | ||
| 98 | + handle: _RecomputedInputHandle | ||
| 99 | + | ||
| 100 | + | ||
| 101 | +@dataclass | ||
| 102 | +class _PackState: | ||
| 103 | + """Inputs and matched saved aliases for the active excluded call.""" | ||
| 104 | + | ||
| 105 | + inputs: Tuple[_InputInfo, ...] | ||
| 106 | + bindings: List[_InputBinding] | ||
| 107 | + | ||
| 108 | + | ||
| 109 | +_ACTIVE_PACK_STATE: ContextVar[Any] = ContextVar( | ||
| 110 | + "hyper_parallel_checkpoint_exclude_pack_state", default=None | ||
| 111 | +) | ||
| 112 | + | ||
| 113 | + | ||
| @@ -0,0 +1,201 @@ | |||
| 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 | +"""Validate checkpoint exclusion memory with stacked RMSNorm, MatMul, and SiLU.""" | ||
| 16 | +import importlib | ||
| 17 | +import json | ||
| 18 | +from pathlib import Path | ||
| 19 | +import subprocess | ||
| 20 | +import sys | ||
| 21 | +from typing import Any, Dict | ||
| 22 | + | ||
| 23 | +import mindspore as ms | ||
| 24 | +from mindspore import nn, Tensor | ||
| 25 | +import numpy as np | ||
| 26 | + | ||
| 27 | +from hyper_parallel.core.activation_checkpoint import checkpoint_exclude_wrapper, checkpoint_wrapper | ||
| 28 | +from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +enable_mindspore_backward_compat() | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +_TOKEN_NUM = 16384 | ||
| 35 | +_HIDDEN_SIZE = 2048 | ||
| 36 | +_LAYER_NUM = 20 | ||
| 37 | +_RESULT_MARKER = "__RMSNORM_MATMUL_RESULT__" | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +class _CountedMatmul(nn.Cell): | ||
| 41 | + """MatMul that records whether checkpoint replay executes it.""" | ||
| 42 | + | ||
| 43 | + def __init__(self, calls: Dict[str, int]) -> None: | ||
| 44 | + """Initialize a deterministic square projection.""" | ||
| 45 | + super().__init__() | ||
| 46 | + self.calls = calls | ||
| 47 | + rng = np.random.default_rng(2026) | ||
| 48 | + weight = rng.normal(0.0, 0.02, (_HIDDEN_SIZE, _HIDDEN_SIZE)).astype(np.float32) | ||
| 49 | + self.weight = ms.Parameter(Tensor(weight, ms.bfloat16), name="matmul_weight") | ||
| 50 | + | ||
| 51 | + def construct(self, tensor: Tensor) -> Tensor: | ||
| 52 | + """Apply the projection.""" | ||
| 53 | + self.calls["matmul"] += 1 | ||
| 54 | + return ms.ops.matmul(tensor, self.weight) | ||
| 55 | + | ||
| 56 | + | ||
| 57 | +class _RmsNormMatmulBlock(nn.Cell): | ||
| 58 | + """Apply RMSNorm, an excluded MatMul, and SiLU.""" | ||
| 59 | + | ||
| 60 | + def __init__(self, calls: Dict[str, int]) -> None: | ||
| 61 | + """Initialize one checkpointed layer.""" | ||
| 62 | + super().__init__() | ||
| 63 | + self.rms_norm = ms.ops.RmsNorm(epsilon=1e-6) | ||
| 64 | + self.norm_weight = ms.Parameter(ms.ops.ones((_HIDDEN_SIZE,), ms.bfloat16), name="norm_weight") | ||
| 65 | + self.matmul = checkpoint_exclude_wrapper(_CountedMatmul(calls)) | ||
| 66 | + | ||
| 67 | + def construct(self, tensor: Tensor) -> Tensor: | ||
| 68 | + """Run the layer.""" | ||
| 69 | + normalized = self.rms_norm(tensor, self.norm_weight)[0] | ||
| 70 | + return ms.ops.silu(self.matmul(normalized)) | ||
| 71 | + | ||
| 72 | + | ||
| 73 | +class _RmsNormMatmulNet(nn.Cell): | ||
| 74 | + """Stack individually checkpointed layers.""" | ||
| 75 | + | ||
| 76 | + def __init__(self, calls: Dict[str, int]) -> None: | ||
| 77 | + """Create one checkpoint for each layer.""" | ||
| 78 | + super().__init__() | ||
| 79 | + self.layers = nn.CellList([ | ||
| 80 | + checkpoint_wrapper(_RmsNormMatmulBlock(calls)) | ||
| 81 | + for _ in range(_LAYER_NUM) | ||
| 82 | + ]) | ||
| 83 | + | ||
| 84 | + def construct(self, tensor: Tensor) -> Tensor: | ||
| 85 | + """Run all layers and return a scalar loss.""" | ||
| 86 | + for layer in self.layers: | ||
| 87 | + tensor = layer(tensor) | ||
| 88 | + return ms.ops.sum(ms.ops.square(tensor)) | ||
| 89 | + | ||
| 90 | + | ||
| 91 | +def _run_exclude() -> Dict[str, Any]: | ||
| 92 | + """Run one optimized or legacy exclusion step.""" | ||
| 93 | + ms.set_deterministic(True) | ||
| 94 | + ms.set_context(mode=ms.PYNATIVE_MODE) | ||
| 95 | + calls = {"matmul": 0} | ||
| 96 | + net = _RmsNormMatmulNet(calls) | ||
| 97 | + rng = np.random.default_rng(2027) | ||
| 98 | + input_data = rng.normal(0.0, 0.5, (_TOKEN_NUM, _HIDDEN_SIZE)).astype(np.float32) | ||
| 99 | + | ||
| 100 | + warmup_input = Tensor(input_data, ms.bfloat16) | ||
| 101 | + warmup_input.requires_grad = True | ||
| 102 | + warmup_loss = net(warmup_input) | ||
| 103 | + warmup_loss.backward() | ||
| 104 | + ms.runtime.synchronize() | ||
| 105 | + for parameter in net.trainable_params(): | ||
| 106 | + parameter.grad = None | ||
| 107 | + del warmup_input, warmup_loss | ||
| 108 | + calls["matmul"] = 0 | ||
| 109 | + ms.runtime.empty_cache() | ||
| 110 | + | ||
| 111 | + baseline_bytes = ms.runtime.memory_allocated() | ||
| 112 | + ms.runtime.reset_peak_memory_stats() | ||
| 113 | + tensor = Tensor(input_data, ms.bfloat16) | ||
| 114 | + tensor.requires_grad = True | ||
| 115 | + loss = net(tensor) | ||
| 116 | + ms.runtime.synchronize() | ||
| 117 | + forward_bytes = ms.runtime.memory_allocated() - baseline_bytes | ||
| 118 | + loss.backward() | ||
| 119 | + ms.runtime.synchronize() | ||
| 120 | + | ||
| 121 | + return { | ||
| 122 | + "loss": float(loss.asnumpy()), | ||
| 123 | + "matmul_calls": calls["matmul"], | ||
| 124 | + "forward_bytes": int(forward_bytes), | ||
| 125 | + "peak_bytes": int(ms.runtime.max_memory_allocated() - baseline_bytes), | ||
| 126 | + } | ||
| 127 | + | ||
| 128 | + | ||
| 129 | +def _disable_input_rematerialization(args: Any, kwargs: Any) -> tuple: | ||
| 130 | + """Reproduce legacy exclusion, which retains saved input storage.""" | ||
| 131 | + del args, kwargs | ||
| 132 | + return () | ||
| 133 | + | ||
| 134 | + | ||
| 135 | +def _run_mode(mode: str) -> Dict[str, Any]: | ||
| 136 | + """Run optimized or legacy checkpoint exclusion.""" | ||
| 137 | + if mode == "exclude": | ||
| 138 | + return _run_exclude() | ||
| 139 | + if mode != "legacy_exclude": | ||
| 140 | + raise ValueError(f"Unsupported RMSNorm/MatMul mode: {mode}") | ||
| 141 | + | ||
| 142 | + exclude_module = importlib.import_module( | ||
| 143 | + "hyper_parallel.platform.mindspore.activation_checkpoint.checkpoint_exclude_wrapper" | ||
| 144 | + ) | ||
| 145 | + original_capture = exclude_module._capture_recompute_inputs # pylint: disable=protected-access | ||
| 146 | + exclude_module._capture_recompute_inputs = _disable_input_rematerialization # pylint: disable=protected-access | ||
| 147 | + try: | ||
| 148 | + return _run_exclude() | ||
| 149 | + finally: | ||
| 150 | + exclude_module._capture_recompute_inputs = original_capture # pylint: disable=protected-access | ||
| 151 | + | ||
| 152 | + | ||
| 153 | +def _run_mode_in_subprocess(mode: str) -> Dict[str, Any]: | ||
| 154 | + """Run one mode in an isolated allocator process.""" | ||
| 155 | + project_root = Path(__file__).resolve().parents[4] | ||
| 156 | + command = [ | ||
| 157 | + sys.executable, | ||
| 158 | + "-c", | ||
| 159 | + ( | ||
| 160 | + "import json; " | ||
| 161 | + "from tests.mindspore.st.activation_checkpoint.checkpoint_exclude_matmul import _run_mode; " | ||
| 162 | + f"print({_RESULT_MARKER!r} + json.dumps(_run_mode({mode!r})))" | ||
| 163 | + ), | ||
| 164 | + ] | ||
| 165 | + completed = subprocess.run( | ||
| 166 | + command, | ||
| 167 | + cwd=project_root, | ||
| 168 | + capture_output=True, | ||
| 169 | + text=True, | ||
| 170 | + check=False, | ||
| 171 | + timeout=1200, | ||
| 172 | + ) | ||
| 173 | + if completed.returncode != 0: | ||
| 174 | + raise RuntimeError( | ||
| 175 | + f"RMSNorm/MatMul subprocess exited with code {completed.returncode}.\n" | ||
| 176 | + f"STDOUT:\n{completed.stdout}\nSTDERR:\n{completed.stderr}" | ||
| 177 | + ) | ||
| 178 | + for line in reversed(completed.stdout.splitlines()): | ||
| 179 | + if _RESULT_MARKER in line: | ||
| 180 | + return json.loads(line.split(_RESULT_MARKER, maxsplit=1)[1]) | ||
| 181 | + raise RuntimeError( | ||
| 182 | + "RMSNorm/MatMul subprocess did not produce a result marker.\n" | ||
| 183 | + f"STDOUT:\n{completed.stdout}\nSTDERR:\n{completed.stderr}" | ||
| 184 | + ) | ||
| 185 | + | ||
| 186 | + | ||
| 187 | +def test_rmsnorm_matmul_checkpoint_exclude_memory() -> None: | ||
| 188 | + """Optimized exclusion should retain less memory than legacy exclusion.""" | ||
| 189 | + excluded = _run_mode_in_subprocess("exclude") | ||
| 190 | + legacy = _run_mode_in_subprocess("legacy_exclude") | ||
| 191 | + | ||
| 192 | + assert excluded["loss"] == legacy["loss"] | ||
| 193 | + assert excluded["matmul_calls"] == _LAYER_NUM | ||
| 194 | + assert legacy["matmul_calls"] == _LAYER_NUM | ||
| 195 | + | ||
| 196 | + expected_gap = _LAYER_NUM * _TOKEN_NUM * _HIDDEN_SIZE * 2 | ||
| 197 | + tolerance = expected_gap // 4 | ||
| 198 | + forward_gap = legacy["forward_bytes"] - excluded["forward_bytes"] | ||
| 199 | + peak_gap = legacy["peak_bytes"] - excluded["peak_bytes"] | ||
| 200 | + assert abs(forward_gap - expected_gap) <= tolerance | ||
| 201 | + assert abs(peak_gap - expected_gap) <= tolerance | ||
| @@ -1,3 +1,8 @@ | |||
| 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 | ||
| 1 | # | 6 | # |
| 2 | # http://www.apache.org/licenses/LICENSE-2.0 | 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 3 | # | 8 | # |
| @@ -30,5 +35,6 @@ def test_sac_group(): | |||
| 30 | MindSporeCase(BASE_SHARD, "test_swap_manager_manual_group_api", 11640, 1), | 35 | MindSporeCase(BASE_SHARD, "test_swap_manager_manual_group_api", 11640, 1), |
| 31 | MindSporeCase(BASE_SHARD, "test_inplace_modification", 11641, 1), | 36 | MindSporeCase(BASE_SHARD, "test_inplace_modification", 11641, 1), |
| 32 | MindSporeCase(BASE_SHARD, "test_wrapper_overlap_detection_cases", 11642, 1), | 37 | MindSporeCase(BASE_SHARD, "test_wrapper_overlap_detection_cases", 11642, 1), |
| 33 | - MindSporeCase(BASE_SHARD, "test_wrapper_non_overlapping_allowed_cases", 11643, 1) | 38 | + MindSporeCase(BASE_SHARD, "test_wrapper_non_overlapping_allowed_cases", 11643, 1), |
| 39 | + MindSporeCase("checkpoint_exclude_matmul.py", "test_rmsnorm_matmul_checkpoint_exclude_memory", 11644, 1), | ||
| 34 | ]) | 40 | ]) |
| @@ -134,6 +134,8 @@ def restore_torch_platform_for_ut() -> None: | |||
| 134 | "hyper_parallel.core.expert_parallel.expert_parallel", | 134 | "hyper_parallel.core.expert_parallel.expert_parallel", |
| 135 | *_FULLY_SHARD_PLATFORM_BIND_MODULES, | 135 | *_FULLY_SHARD_PLATFORM_BIND_MODULES, |
| 136 | "hyper_parallel.core.context_parallel.async_context_parallel", | 136 | "hyper_parallel.core.context_parallel.async_context_parallel", |
| 137 | + "hyper_parallel.core.activation_checkpoint", | ||
| 138 | + "hyper_parallel.core.activation_checkpoint.activation_checkpoint", | ||
| 137 | ) | 139 | ) |
| 138 | ) | 140 | ) |
| 139 | try: | 141 | try: |
| @@ -14,9 +14,11 @@ | |||
| 14 | # ============================================================================ | 14 | # ============================================================================ |
| 15 | """Tests for the MindSpore checkpoint exclusion wrapper.""" | 15 | """Tests for the MindSpore checkpoint exclusion wrapper.""" |
| 16 | import contextlib | 16 | import contextlib |
| 17 | +import gc | ||
| 17 | import importlib | 18 | import importlib |
| 18 | -from typing import Iterator, Tuple | 19 | +from typing import Any, Iterator, Tuple |
| 19 | import unittest | 20 | import unittest |
| 21 | +import weakref | ||
| 20 | 22 | ||
| 21 | import numpy as np | 23 | import numpy as np |
| 22 | import pytest | 24 | import pytest |
| @@ -25,12 +27,17 @@ from tests.ut.platform.mindspore._ensure_mindspore_platform import ( | |||
| 25 | ) | 27 | ) |
| 26 | 28 | ||
| 27 | ms = pytest.importorskip("mindspore") | 29 | ms = pytest.importorskip("mindspore") |
| 30 | +from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat | ||
| 31 | + | ||
| 28 | ParameterTuple = ms.ParameterTuple | 32 | ParameterTuple = ms.ParameterTuple |
| 33 | +Parameter = ms.Parameter | ||
| 29 | Tensor = ms.Tensor | 34 | Tensor = ms.Tensor |
| 30 | nn = ms.nn | 35 | nn = ms.nn |
| 31 | ops = ms.ops | 36 | ops = ms.ops |
| 37 | +_Function = importlib.import_module("mindspore.common._grad_function")._Function | ||
| 32 | _pynative_executor = importlib.import_module("mindspore.graph.api")._pynative_executor | 38 | _pynative_executor = importlib.import_module("mindspore.graph.api")._pynative_executor |
| 33 | 39 | ||
| 40 | +enable_mindspore_backward_compat() | ||
| 34 | ensure_mindspore_platform_default() | 41 | ensure_mindspore_platform_default() |
| 35 | 42 | ||
| 36 | activation_checkpoint = importlib.import_module("hyper_parallel.core.activation_checkpoint") | 43 | activation_checkpoint = importlib.import_module("hyper_parallel.core.activation_checkpoint") |
| @@ -47,6 +54,39 @@ checkpoint_exclude_wrapper_module = importlib.import_module( | |||
| 47 | CheckpointExcludeWrapper = checkpoint_exclude_wrapper_module.CheckpointExcludeWrapper | 54 | CheckpointExcludeWrapper = checkpoint_exclude_wrapper_module.CheckpointExcludeWrapper |
| 48 | 55 | ||
| 49 | 56 | ||
| 57 | +def _pack_with_inputs(saved_tensor, args, kwargs=None): | ||
| 58 | + """Run the exclude pack hook with storage metadata captured from inputs.""" | ||
| 59 | + kwargs = {} if kwargs is None else kwargs | ||
| 60 | + inputs = checkpoint_exclude_wrapper_module._capture_recompute_inputs(args, kwargs) | ||
| 61 | + bindings = [] | ||
| 62 | + state = checkpoint_exclude_wrapper_module._PackState(inputs, bindings) | ||
| 63 | + token = checkpoint_exclude_wrapper_module._ACTIVE_PACK_STATE.set(state) | ||
| 64 | + try: | ||
| 65 | + packed = checkpoint_exclude_wrapper_module._pack_saved_tensor(saved_tensor) | ||
| 66 | + finally: | ||
| 67 | + checkpoint_exclude_wrapper_module._ACTIVE_PACK_STATE.reset(token) | ||
| 68 | + return packed, bindings | ||
| 69 | + | ||
| 70 | + | ||
| 71 | +class _SaveExactInput(_Function): | ||
| 72 | + """Autograd function that saves its exact input for backward.""" | ||
| 73 | + | ||
| 74 | + calls = 0 | ||
| 75 | + | ||
| 76 | + @staticmethod | ||
| 77 | + def forward(ctx: Any, tensor: Tensor) -> Tensor: | ||
| 78 | + """Save the input and return its square.""" | ||
| 79 | + _SaveExactInput.calls += 1 | ||
| 80 | + ctx.save_for_backward(tensor) | ||
| 81 | + return tensor * tensor | ||
| 82 | + | ||
| 83 | + @staticmethod | ||
| 84 | + def backward(ctx: Any, grad_output: Tensor) -> Tensor: | ||
| 85 | + """Use the saved input to compute the square gradient.""" | ||
| 86 | + (tensor,) = ctx.saved_tensors | ||
| 87 | + return grad_output * 2 * tensor | ||
| 88 | + | ||
| 89 | + | ||
| 50 | class _FunctionBlock(nn.Cell): | 90 | class _FunctionBlock(nn.Cell): |
| 51 | """Checkpointed block with a plain function excluded from recomputation.""" | 91 | """Checkpointed block with a plain function excluded from recomputation.""" |
| 52 | 92 | ||
| @@ -92,6 +132,27 @@ class _StateBlock(nn.Cell): | |||
| 92 | return (output * x).sum() | 132 | return (output * x).sum() |
| 93 | 133 | ||
| 94 | 134 | ||
| 135 | +class _SavedInputBlock(nn.Cell): | ||
| 136 | + """Checkpointed block whose excluded tail saves its exact input.""" | ||
| 137 | + | ||
| 138 | + def __init__(self, calls: dict) -> None: | ||
| 139 | + """Wrap a custom saved-input function and retain the shared counter.""" | ||
| 140 | + super().__init__() | ||
| 141 | + self.calls = calls | ||
| 142 | + | ||
| 143 | + def excluded(tensor: Tensor) -> Tensor: | ||
| 144 | + """Save the exact recomputed input without rerunning on replay.""" | ||
| 145 | + self.calls["excluded"] += 1 | ||
| 146 | + return _SaveExactInput.apply(tensor) | ||
| 147 | + | ||
| 148 | + self.excluded = checkpoint_exclude_wrapper(excluded) | ||
| 149 | + | ||
| 150 | + def construct(self, tensor: Tensor) -> Tensor: | ||
| 151 | + """End the checkpoint body at the excluded region's scalar output.""" | ||
| 152 | + self.calls["block"] += 1 | ||
| 153 | + return self.excluded(tensor * 3).sum() | ||
| 154 | + | ||
| 155 | + | ||
| 95 | class _PrecisionMiddle(nn.Cell): | 156 | class _PrecisionMiddle(nn.Cell): |
| 96 | """Parameterized nonlinear region used by the precision test.""" | 157 | """Parameterized nonlinear region used by the precision test.""" |
| 97 | 158 | ||
| @@ -217,6 +278,206 @@ class TestCheckpointExcludeWrapper(unittest.TestCase): | |||
| 217 | np.testing.assert_array_equal(packed.asnumpy(), tensor.asnumpy()) | 278 | np.testing.assert_array_equal(packed.asnumpy(), tensor.asnumpy()) |
| 218 | self.assertIs(checkpoint_exclude_wrapper_module._unpack_saved_tensor(packed), packed) | 279 | self.assertIs(checkpoint_exclude_wrapper_module._unpack_saved_tensor(packed), packed) |
| 219 | 280 | ||
| 281 | + def test_saved_tensor_hook_uses_storage_matched_handle(self): | ||
| 282 | + """An exact saved input should be deferred by storage/layout matching.""" | ||
| 283 | + tensor = Tensor(np.arange(4, dtype=np.float32)) | ||
| 284 | + recomputed = Tensor(np.arange(4, dtype=np.float32) * 2) | ||
| 285 | + | ||
| 286 | + packed, bindings = _pack_with_inputs(tensor, (tensor,)) | ||
| 287 | + | ||
| 288 | + self.assertIsInstance(packed, checkpoint_exclude_wrapper_module._RecomputedInputHandle) | ||
| 289 | + self.assertEqual(len(bindings), 1) | ||
| 290 | + self.assertTrue(packed.recipe.exact_input) | ||
| 291 | + with self.assertRaisesRegex(RuntimeError, "before recomputation"): | ||
| 292 | + checkpoint_exclude_wrapper_module._unpack_saved_tensor(packed) | ||
| 293 | + entry = checkpoint_exclude_wrapper_module._ExcludeCacheEntry(None, bindings) | ||
| 294 | + checkpoint_exclude_wrapper_module._materialize_recompute_inputs(entry, (recomputed,), {}) | ||
| 295 | + unpacked = checkpoint_exclude_wrapper_module._unpack_saved_tensor(packed) | ||
| 296 | + np.testing.assert_array_equal(unpacked.asnumpy(), recomputed.asnumpy()) | ||
| 297 | + | ||
| 298 | + def test_saved_input_view_is_rebuilt_from_replay_storage(self): | ||
| 299 | + """A saved slice should be reconstructed with its shape, stride, and offset.""" | ||
| 300 | + base = Tensor(np.arange(24, dtype=np.float32).reshape(2, 3, 4)) | ||
| 301 | + saved_view = base[:, 1:, :] | ||
| 302 | + | ||
| 303 | + packed, bindings = _pack_with_inputs(saved_view, (base,)) | ||
| 304 | + | ||
| 305 | + self.assertIsInstance(packed, checkpoint_exclude_wrapper_module._RecomputedInputHandle) | ||
| 306 | + self.assertFalse(packed.recipe.exact_input) | ||
| 307 | + self.assertEqual(packed.recipe.relative_offset, 4) | ||
| 308 | + | ||
| 309 | + replay = Tensor((np.arange(24, dtype=np.float32) + 100).reshape(2, 3, 4)) | ||
| 310 | + expected = replay[:, 1:, :] | ||
| 311 | + entry = checkpoint_exclude_wrapper_module._ExcludeCacheEntry(None, bindings) | ||
| 312 | + checkpoint_exclude_wrapper_module._materialize_recompute_inputs(entry, (replay,), {}) | ||
| 313 | + restored = checkpoint_exclude_wrapper_module._unpack_saved_tensor(packed) | ||
| 314 | + | ||
| 315 | + np.testing.assert_array_equal(restored.asnumpy(), expected.asnumpy()) | ||
| 316 | + self.assertEqual(tuple(restored.shape), tuple(expected.shape)) | ||
| 317 | + self.assertEqual(tuple(restored.stride()), tuple(expected.stride())) | ||
| 318 | + self.assertEqual(restored.storage_offset(), expected.storage_offset()) | ||
| 319 | + self.assertEqual( | ||
| 320 | + restored.untyped_storage().data_ptr(), | ||
| 321 | + replay.untyped_storage().data_ptr(), | ||
| 322 | + ) | ||
| 323 | + | ||
| 324 | + def test_storage_match_prefers_exact_input_among_shared_storage_views(self): | ||
| 325 | + """Exact layout should disambiguate query/rope-like views sharing storage.""" | ||
| 326 | + packed_q_rope = Tensor(np.arange(48, dtype=np.float32).reshape(2, 3, 8)) | ||
| 327 | + query = packed_q_rope[:, :, :6] | ||
| 328 | + query_rope = packed_q_rope[:, :, 6:] | ||
| 329 | + | ||
| 330 | + packed, bindings = _pack_with_inputs(query_rope, (query, query_rope)) | ||
| 331 | + | ||
| 332 | + self.assertIsInstance(packed, checkpoint_exclude_wrapper_module._RecomputedInputHandle) | ||
| 333 | + self.assertEqual(len(bindings), 1) | ||
| 334 | + self.assertEqual(bindings[0].path, (("arg", 1),)) | ||
| 335 | + self.assertTrue(packed.recipe.exact_input) | ||
| 336 | + | ||
| 337 | + def test_new_storage_saved_tensor_falls_back_to_real_save(self): | ||
| 338 | + """A computed copy cannot be rebuilt from layout metadata alone.""" | ||
| 339 | + tensor = Tensor(np.arange(4, dtype=np.float32)) | ||
| 340 | + copied = tensor + 1 | ||
| 341 | + | ||
| 342 | + packed, bindings = _pack_with_inputs(copied, (tensor,)) | ||
| 343 | + | ||
| 344 | + self.assertNotIsInstance(packed, checkpoint_exclude_wrapper_module._RecomputedInputHandle) | ||
| 345 | + self.assertEqual(bindings, []) | ||
| 346 | + np.testing.assert_array_equal(packed.asnumpy(), copied.asnumpy()) | ||
| 347 | + | ||
| 348 | + def test_recompute_boundary_saves_zero_element_trigger(self): | ||
| 349 | + """The replay boundary should not retain storage from its tensor output.""" | ||
| 350 | + packed_shapes = [] | ||
| 351 | + unpacked_shapes = [] | ||
| 352 | + | ||
| 353 | + def pack_hook(tensor: Tensor) -> Tensor: | ||
| 354 | + """Record the tensor packed by the boundary.""" | ||
| 355 | + packed_shapes.append(tuple(tensor.shape)) | ||
| 356 | + return tensor | ||
| 357 | + | ||
| 358 | + def unpack_hook(tensor: Tensor) -> Tensor: | ||
| 359 | + """Record the tensor unpacked by the boundary.""" | ||
| 360 | + unpacked_shapes.append(tuple(tensor.shape)) | ||
| 361 | + return tensor | ||
| 362 | + | ||
| 363 | + def apply_boundary(tensor: Tensor) -> Tensor: | ||
| 364 | + """Apply the boundary in a differentiable function.""" | ||
| 365 | + return checkpoint_exclude_wrapper_module._RecomputeBoundary.apply(tensor).sum() | ||
| 366 | + | ||
| 367 | + tensor = Tensor(np.arange(4, dtype=np.float32)) | ||
| 368 | + with ms.saved_tensors_hooks(pack_hook, unpack_hook): | ||
| 369 | + _, grad = ms.value_and_grad(apply_boundary)(tensor) | ||
| 370 | + | ||
| 371 | + self.assertEqual(packed_shapes, [(0,)]) | ||
| 372 | + self.assertEqual(unpacked_shapes, [(0,)]) | ||
| 373 | + np.testing.assert_array_equal(grad.asnumpy(), np.ones(4, dtype=np.float32)) | ||
| 374 | + | ||
| 375 | + def test_recompute_boundary_lazily_reuses_trigger(self): | ||
| 376 | + """The boundary trigger should be created lazily and reused globally.""" | ||
| 377 | + get_trigger = checkpoint_exclude_wrapper_module._get_recompute_trigger | ||
| 378 | + get_trigger.cache_clear() | ||
| 379 | + | ||
| 380 | + self.assertEqual(get_trigger.cache_info().currsize, 0) | ||
| 381 | + trigger = get_trigger() | ||
| 382 | + self.assertIs(trigger, get_trigger()) | ||
| 383 | + self.assertEqual(trigger.dtype, ms.float32) | ||
| 384 | + self.assertEqual(tuple(trigger.shape), (0,)) | ||
| 385 | + self.assertEqual(get_trigger.cache_info().currsize, 1) | ||
| 386 | + | ||
| 387 | + def test_parameter_input_is_not_captured(self): | ||
| 388 | + """Parameter inputs should keep the existing saved-tensor behavior.""" | ||
| 389 | + parameter = Parameter(Tensor(np.arange(4, dtype=np.float32)), name="exclude_parameter") | ||
| 390 | + | ||
| 391 | + inputs = checkpoint_exclude_wrapper_module._capture_recompute_inputs((parameter,), {}) | ||
| 392 | + | ||
| 393 | + self.assertEqual(inputs, ()) | ||
| 394 | + | ||
| 395 | + def test_input_capture_does_not_change_active_pack_state_on_failure(self): | ||
| 396 | + """A traversal failure should not leak active pack state.""" | ||
| 397 | + class _BrokenDict(dict): | ||
| 398 | + """Raise while exposing nested items.""" | ||
| 399 | + | ||
| 400 | + def items(self) -> Any: | ||
| 401 | + """Fail before returning nested items.""" | ||
| 402 | + raise RuntimeError("nested traversal failed") | ||
| 403 | + | ||
| 404 | + tensor = Tensor(np.arange(4, dtype=np.float32)) | ||
| 405 | + | ||
| 406 | + with self.assertRaisesRegex(RuntimeError, "nested traversal failed"): | ||
| 407 | + checkpoint_exclude_wrapper_module._capture_recompute_inputs((tensor, _BrokenDict()), {}) | ||
| 408 | + | ||
| 409 | + self.assertIsNone(checkpoint_exclude_wrapper_module._ACTIVE_PACK_STATE.get()) | ||
| 410 | + | ||
| 411 | + def test_nested_input_paths_materialize_matching_replay_tensors(self): | ||
| 412 | + """Nested args and kwargs should bind handles without retaining forward inputs.""" | ||
| 413 | + first = Tensor(np.array([1.0, 2.0], dtype=np.float32)) | ||
| 414 | + second = Tensor(np.array([3.0, 4.0], dtype=np.float32)) | ||
| 415 | + args = ({"items": [first]}, first) | ||
| 416 | + kwargs = {"second": second} | ||
| 417 | + | ||
| 418 | + packed_first, first_bindings = _pack_with_inputs(first, args, kwargs) | ||
| 419 | + packed_second, second_bindings = _pack_with_inputs(second, args, kwargs) | ||
| 420 | + bindings = first_bindings + second_bindings | ||
| 421 | + | ||
| 422 | + replay_first = Tensor(np.array([5.0, 6.0], dtype=np.float32)) | ||
| 423 | + replay_second = Tensor(np.array([7.0, 8.0], dtype=np.float32)) | ||
| 424 | + entry = checkpoint_exclude_wrapper_module._ExcludeCacheEntry(output=None, input_bindings=bindings) | ||
| 425 | + checkpoint_exclude_wrapper_module._materialize_recompute_inputs( | ||
| 426 | + entry, | ||
| 427 | + ({"items": [replay_first]}, replay_first), | ||
| 428 | + {"second": replay_second}, | ||
| 429 | + ) | ||
| 430 | + | ||
| 431 | + self.assertEqual(len(bindings), 2) | ||
| 432 | + self.assertEqual(bindings[0].path, (("arg", 0), ("key", "items"), ("index", 0))) | ||
| 433 | + self.assertEqual(bindings[1].path, (("kwarg", "second"),)) | ||
| 434 | + np.testing.assert_array_equal( | ||
| 435 | + checkpoint_exclude_wrapper_module._unpack_saved_tensor(packed_first).asnumpy(), | ||
| 436 | + replay_first.asnumpy(), | ||
| 437 | + ) | ||
| 438 | + np.testing.assert_array_equal( | ||
| 439 | + checkpoint_exclude_wrapper_module._unpack_saved_tensor(packed_second).asnumpy(), | ||
| 440 | + replay_second.asnumpy(), | ||
| 441 | + ) | ||
| 442 | + | ||
| 443 | + def test_input_traversal_does_not_require_cycle_collection(self): | ||
| 444 | + """Input traversal should release tensor leaves without waiting for cyclic GC.""" | ||
| 445 | + def traverse_tensor() -> weakref.ReferenceType: | ||
| 446 | + """Collect a tensor input and return a non-owning reference to it.""" | ||
| 447 | + tensor = Tensor(np.arange(4, dtype=np.float32)) | ||
| 448 | + inputs = checkpoint_exclude_wrapper_module._capture_recompute_inputs(([tensor],), {}) | ||
| 449 | + self.assertEqual(len(inputs), 1) | ||
| 450 | + return weakref.ref(tensor) | ||
| 451 | + | ||
| 452 | + gc_was_enabled = gc.isenabled() | ||
| 453 | + gc.disable() | ||
| 454 | + try: | ||
| 455 | + tensor_ref = traverse_tensor() | ||
| 456 | + self.assertIsNone(tensor_ref()) | ||
| 457 | + finally: | ||
| 458 | + if gc_was_enabled: | ||
| 459 | + gc.enable() | ||
| 460 | + gc.collect() | ||
| 461 | + | ||
| 462 | + def test_recomputed_input_materializes_before_excluded_backward(self): | ||
| 463 | + """A terminal excluded region should recover its saved input from replay.""" | ||
| 464 | + calls = {"block": 0, "excluded": 0} | ||
| 465 | + _SaveExactInput.calls = 0 | ||
| 466 | + net = checkpoint_wrapper(_SavedInputBlock(calls)) | ||
| 467 | + tensor = Tensor(np.array([1.0, 2.0], dtype=np.float32)) | ||
| 468 | + | ||
| 469 | + value, grad = ms.value_and_grad(net, grad_position=0)(tensor) | ||
| 470 | + | ||
| 471 | + np.testing.assert_allclose(value.asnumpy(), np.array(45.0, np.float32), atol=1e-6, rtol=1e-6) | ||
| 472 | + np.testing.assert_allclose( | ||
| 473 | + grad.asnumpy(), | ||
| 474 | + np.array([18.0, 36.0], np.float32), | ||
| 475 | + atol=1e-6, | ||
| 476 | + rtol=1e-6, | ||
| 477 | + ) | ||
| 478 | + self.assertEqual(calls, {"block": 2, "excluded": 1}) | ||
| 479 | + self.assertEqual(_SaveExactInput.calls, 1) | ||
| 480 | + | ||
| 220 | def test_function_wrapper_skips_middle_recompute(self): | 481 | def test_function_wrapper_skips_middle_recompute(self): |
| 221 | """A wrapped function should run in forward and reuse its output in recompute.""" | 482 | """A wrapped function should run in forward and reuse its output in recompute.""" |
| 222 | calls = {"middle": 0} | 483 | calls = {"middle": 0} |