diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py new file mode 100644 index 000000000..b1738f936 --- /dev/null +++ b/tests/model/test_selective_checkpointing.py @@ -0,0 +1,238 @@ +"""选择性重算（SAC）的回归测试。 + +TestKeptOps + test_kept_op_reproduces_full_recompute: 按 op 留驻与全重算的输出/梯度逐位相同且梯度存活。 + test_unknown_op_name_is_skipped: 本次构建没注册的 op 名字被跳过，不报错。 +TestKeptCallables + test_kept_callable_reproduces_full_recompute: 按 callable 留驻与全重算逐位相同。 + test_unit_marker_outside_a_region_is_noop: 未被包装时不产生任何可观察行为。 +TestUnsupportedUnits + test_in_place_op_in_kept_unit_is_recomputed_not_refused: 只动自己缓冲区的 in-place 写不该被拒。 + test_mutating_a_kept_tensor_is_caught_by_torch: 写到被留驻张量上时由 torch 精确拦下。 + test_in_place_op_outside_kept_unit_is_fine: 单元之外的 in-place 写不受影响。 +TestRecomputeIsObservable + test_recompute_reproduces_the_plain_gradients: 重算与不重算的梯度逐位相同。 + test_checkpointing_actually_recomputes: 重算确实发生——op 执行次数翻倍。 + test_a_kept_op_is_not_recomputed: 被留驻的 op 不参与重算，计数不翻倍。 +TestRegionRecomputeUnderDominoEP + test_kept_unit_matches_full_recompute_under_domino_ep: domino EP 下留驻与全重算数值一致。 + test_kept_unit_matches_full_recompute_under_compile: compile 下同上，且不触发 cached-tensor-mutated。 +""" + +from collections import Counter + +import pytest +import torch +from torch import nn +from torch.utils._python_dispatch import TorchDispatchMode + +from xtuner.v1.model.utils import ( + RecomputeUnit, + apply_selective_checkpointing, + in_recompute_unit, + resolve_kept_ops, +) + + +class _Block(nn.Module): + """两级 linear，方便留驻其中一级、重算另一级。""" + + def __init__(self) -> None: + super().__init__() + self.first = nn.Linear(4, 4) + self.second = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + hidden = torch.tanh(self.first(x)) + return torch.tanh(self.second(hidden)) + + +class _UnitBlock(_Block): + """第二级被包成一个 unit，形态与引擎实际装上去的一致。""" + + def __init__(self) -> None: + super().__init__() + self.second_stage = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._second_stage) + + def _second_stage(self, hidden: torch.Tensor) -> torch.Tensor: + return torch.tanh(self.second(hidden)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.second_stage(torch.tanh(self.first(x))) + + +class _InPlaceUnitBlock(nn.Module): + """unit 内部做原地累加——这种写法不能被留驻。""" + + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(4, 4) + self.unit = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._unit) + + def _unit(self, x: torch.Tensor) -> torch.Tensor: + hidden = self.linear(x) + accumulator = torch.zeros_like(hidden) + accumulator.add_(hidden) + return accumulator + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.unit(x) + + +class _MutatesKeptTensorBlock(nn.Module): + """unit 原地改写了一个已被留驻的张量，这才是真正不安全的情形。""" + + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(4, 4) + self.unit = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._unit) + + def _unit(self, x: torch.Tensor) -> torch.Tensor: + kept = torch.tanh(self.linear(x)) + kept.mul_(2.0) + return kept + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.unit(x) + + +class _InPlaceOutsideBlock(_InPlaceUnitBlock): + """同样的原地写，但没有进入 unit。""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self._unit(x) + + +class _OpCounter(TorchDispatchMode): + """直接数每个 op 实际执行了多少次。 + + 这是判断「重算/SAC 到底有没有生效」最直接的观察点：不看 policy 返回了什么，只看 op 跑了几遍。 + 前向跑一遍、重算再跑一遍，所以被重算的 op 计数翻倍；被留驻的 op 不参与重算，计数保持一遍。 + """ + + def __init__(self) -> None: + super().__init__() + self.counts: Counter = Counter() + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + self.counts[func] += 1 + return func(*args, **(kwargs or {})) + + +def _count_ops(make_module, kept_ops=frozenset(), *, checkpointed: bool, keeps_any_unit: bool = False): + """跑一次前反向，返回 (op 计数, 输入梯度, 权重梯度)。 + + 收工厂而不是收实例：权重必须在同一个 seed 下构造，否则比较的是两组不同的随机权重。 + """ + torch.manual_seed(0) + module = make_module() + target = apply_selective_checkpointing(module, kept_ops, keeps_any_unit=keeps_any_unit) if checkpointed else module + x = torch.randn(2, 4, requires_grad=True) + + counter = _OpCounter() + with counter: + target(x).square().sum().backward() + return counter.counts, x.grad.clone(), module.first.weight.grad.clone() + + +def _run(module: nn.Module, kept_ops=frozenset(), *, keeps_any_unit: bool = False): + torch.manual_seed(0) + for parameter in module.parameters(): + nn.init.normal_(parameter, std=0.5) + wrapped = apply_selective_checkpointing(module, kept_ops, keeps_any_unit=keeps_any_unit) + + inputs = torch.arange(12, dtype=torch.float32).reshape(3, 4) / 12 + inputs.requires_grad_(True) + output = wrapped(inputs) + output.sum().backward() + return output.detach().clone(), [parameter.grad.clone() for parameter in module.parameters()] + + +def _assert_matches(kept, recomputed): + kept_out, kept_grads = kept + recomputed_out, recomputed_grads = recomputed + torch.testing.assert_close(kept_out, recomputed_out, atol=0.0, rtol=0.0) + for one, other in zip(kept_grads, recomputed_grads): + torch.testing.assert_close(one, other, atol=0.0, rtol=0.0) + # 梯度断掉时 loss 依然有限、上面的比较也依然成立（两边都是 None/零），所以单独断言存活。 + for grad in kept_grads: + assert grad is not None + assert torch.count_nonzero(grad) > 0 + + +class TestKeptOps: + def test_kept_op_reproduces_full_recompute(self): + # 留驻与重算两条路都必须是精确值，不是近似：任何差异都说明 save-list 与重算对不上。 + kept_ops = resolve_kept_ops(("aten::tanh",)) + assert kept_ops, "aten::tanh should resolve in every build" + _assert_matches(_run(_Block(), kept_ops, keeps_any_unit=True), _run(_Block())) + + def test_unknown_op_name_is_skipped(self): + # 模型会同时列出同一个 kernel 的多种拼写（flash-attn v2/v3），只有一种在本次构建里注册。 + assert resolve_kept_ops(("nonexistent_namespace::nonexistent_op",)) == frozenset() + assert len(resolve_kept_ops(("aten::tanh", "nonexistent_namespace::nope"))) == 1 + + +class TestKeptCallables: + def test_kept_callable_reproduces_full_recompute(self): + _assert_matches(_run(_UnitBlock(), keeps_any_unit=True), _run(_UnitBlock())) + + def test_unit_marker_outside_a_region_is_noop(self): + # 模型可以先被包装、后开启 recompute；包装本身不能有任何可观察行为。 + module = _UnitBlock() + inputs = torch.zeros(3, 4, requires_grad=True) + module(inputs).sum().backward() + + assert inputs.grad is not None + + +class TestUnsupportedUnits: + def test_in_place_op_in_kept_unit_is_recomputed_not_refused(self): + # in-place op 本身不危险，危险的是它写到了被 unit 留驻的张量上——那由 torch 的 + # version 检查精确判定。schema 是 mutable 只说明"可疑"，据此拒绝会误杀那些 + # 只动自己缓冲区的库调用（deep_ep 的 set_）。所以这里只重算、不拒绝。 + _run(_InPlaceUnitBlock(), keeps_any_unit=True) + + def test_mutating_a_kept_tensor_is_caught_by_torch(self): + # 真正不安全的那种：被留驻的张量随后被原地改写，重算取回的就是改写后的值。 + with pytest.raises(RuntimeError, match="has been mutated"): + _run(_MutatesKeptTensorBlock(), resolve_kept_ops(("aten::tanh",)), keeps_any_unit=True) + + def test_in_place_op_outside_kept_unit_is_fine(self): + _run(_InPlaceOutsideBlock(), keeps_any_unit=True) + + +class TestRecomputeIsObservable: + """按 reviewer 的三个目标直接观察：精度对齐、重算生效、重算+SAC 生效。""" + + def test_recompute_reproduces_the_plain_gradients(self): + # 目标 1：精度与不重算完全一致（逐位，不给容差）。 + _, plain_x, plain_w = _count_ops(_Block, checkpointed=False) + _, ckpt_x, ckpt_w = _count_ops(_Block, checkpointed=True) + + assert torch.equal(plain_x, ckpt_x) + assert torch.equal(plain_w, ckpt_w) + + def test_checkpointing_actually_recomputes(self): + # 目标 2：重算真的发生了——前向的 op 在 backward 里又跑了一遍，计数翻倍。 + plain, _, _ = _count_ops(_Block, checkpointed=False) + ckpt, _, _ = _count_ops(_Block, checkpointed=True) + + tanh = torch.ops.aten.tanh.default + assert plain[tanh] == 2, plain[tanh] + assert ckpt[tanh] == 2 * plain[tanh] + + def test_a_kept_op_is_not_recomputed(self): + # 目标 3：SAC 生效——被留驻的那个 op 不参与重算，计数回到只跑一遍。 + kept = resolve_kept_ops(("aten::tanh",)) + ckpt, _, _ = _count_ops(_Block, checkpointed=True) + sac, sac_x, sac_w = _count_ops(_Block, kept, checkpointed=True, keeps_any_unit=True) + + tanh = torch.ops.aten.tanh.default + assert ckpt[tanh] == 4 + assert sac[tanh] == 2, "留驻的 op 不该在重算里再跑一遍" + + # 少跑不等于算对，所以顺带把梯度也对一遍。 + _, plain_x, plain_w = _count_ops(_Block, checkpointed=False) + assert torch.equal(plain_x, sac_x) + assert torch.equal(plain_w, sac_w) diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 13166e218..697e5f993 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -984,6 +984,29 @@ def compile_cfg(self) -> dict[str, TorchCompileOption]: return _compile_cfg + @property + def kept_ops(self) -> frozenset: + """Op overloads kept resident wherever they run, as selected by + ``config.recompute_cfg``. + + Returns: + frozenset: Overloads to keep. Empty means no op-identity unit was selected. + """ + return frozenset() + + @property + def keeps_any_recompute_unit(self) -> bool: + """Whether any unit at all was selected. + + This is what tells the sharding paths whether to install the per-op policy at all: with + nothing selected, torch's own default already recomputes everything, so a policy that + reaches the same answer would only put a dispatch mode in the way of every op. + + Returns: + bool: True when ``config.recompute_cfg`` selected something, by either resolution. + """ + return False + @property def float8_handler(self): if ( diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py index a971efcc6..7ec701056 100644 --- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py +++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py @@ -34,7 +34,7 @@ fully_shard, ) from xtuner.v1.ops.attn_imp import attn_impl_mapping, AttnOpOutputs -from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing +from xtuner.v1.model.utils import apply_selective_checkpointing from xtuner.v1.module import RMSNorm from xtuner.v1.ops.others import Dropout from xtuner.v1.ops.act_fn import get_act_fn @@ -407,7 +407,12 @@ def fully_shard( layer = self.encoder.layer[layer_idx] if layer_idx < num_recompute_layers: - layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state) + layer = apply_selective_checkpointing( + layer, + self.kept_ops, + keeps_any_unit=self.keeps_any_recompute_unit, + preserve_rng_state=checkpoint_preserve_rng_state, + ) if self.config.drop_path_rate == 0.0 and self.compile_cfg: layer.forward = torch.compile(layer.forward, fullgraph=True) diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py index aa1dd2a83..efacd5f9e 100644 --- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py +++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py @@ -24,7 +24,7 @@ from torch.distributed.device_mesh import init_device_mesh import torch.distributed as dist from xtuner.v1.utils.compile import maybe_compile -from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing +from xtuner.v1.model.utils import apply_selective_checkpointing from xtuner.v1.module import AttnOutputs from torch.distributed.device_mesh import DeviceMesh from tqdm import tqdm @@ -338,7 +338,12 @@ def fully_shard( layer = self.blocks[layer_idx] if layer_idx < num_recompute_layers: - layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state) + layer = apply_selective_checkpointing( + layer, + self.kept_ops, + keeps_any_unit=self.keeps_any_recompute_unit, + preserve_rng_state=checkpoint_preserve_rng_state, + ) if self.compile_cfg: layer.forward = torch.compile(layer.forward, fullgraph=True) diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index 59a9a7df0..fbabdaa74 100644 --- a/xtuner/v1/model/dense/dense.py +++ b/xtuner/v1/model/dense/dense.py @@ -26,7 +26,7 @@ TorchCompileOption, TransformerConfig, ) -from xtuner.v1.model.utils import apply_gradient_checkpointing +from xtuner.v1.model.utils import apply_selective_checkpointing from xtuner.v1.module import ( GatedDeltaNetConfig, LMHead, @@ -235,7 +235,12 @@ def fully_shard( layer = self.layers[str(int(layer_idx))] layer_idx = int(layer_idx) if layer_idx < num_recompute_layers: - layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state) + layer = apply_selective_checkpointing( + layer, + self.kept_ops, + keeps_any_unit=self.keeps_any_recompute_unit, + preserve_rng_state=checkpoint_preserve_rng_state, + ) # Linear-attention (GatedDeltaNet) layers write ``seq_ctx.seq_idx`` inside the # checkpoint region; compiling the checkpointed layer with ``fullgraph=True`` turns diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 316bdfca0..04e0f9b3f 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -46,7 +46,7 @@ ) from xtuner.v1.model.utils import ( ModelForwardExtraLogInfo, - apply_gradient_checkpointing, + apply_selective_checkpointing, module_dict_repr, ) from xtuner.v1.module import ( @@ -90,9 +90,13 @@ logger = get_logger() +# Compiling the decoder layer as a whole is what makes marker intervals inert, so the key is named +# rather than spelled out at each use site. +MOE_DECODER_LAYER_FORWARD = "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward" + MOE_NON_EP_COMPILE_CFG: dict[str, TorchCompileOption] = { "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEBlock.forward": TorchCompileOption(fullgraph=True), - "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward": TorchCompileOption(fullgraph=True), + MOE_DECODER_LAYER_FORWARD: TorchCompileOption(fullgraph=True), "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer._pre_moe_forward": TorchCompileOption( fullgraph=True ), @@ -107,7 +111,7 @@ } MOE_EP_COMPILE_CFG = MOE_NON_EP_COMPILE_CFG.copy() -MOE_EP_COMPILE_CFG.pop("xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward") +MOE_EP_COMPILE_CFG.pop(MOE_DECODER_LAYER_FORWARD) class MoEModelOutputs(ModelOutputs): @@ -1168,7 +1172,11 @@ def fully_shard( layer_idx=layer_idx, mtp_idx=None, ): - layer = apply_gradient_checkpointing(layer) + layer = apply_selective_checkpointing( + layer, + self.kept_ops, + keeps_any_unit=self.keeps_any_recompute_unit, + ) self.layers[str(layer_idx)] = layer if layer_idx >= len(self.layers) - 1 and self.mtp_block is None: @@ -1220,7 +1228,9 @@ def fully_shard( if self._should_recompute(None, mtp_idx=mtp_idx) or ( self.config.mtp_config is not None and self.config.mtp_config.share_weights ): # share mtp head must recompute - mtp_layer = apply_gradient_checkpointing(mtp_layer) + # Marker intervals are declared against the decoder layer, so an MTP layer + # gets the same mechanism with nothing kept: recomputed whole, as before. + mtp_layer = apply_selective_checkpointing(mtp_layer) self.mtp_block.layers[mtp_idx] = mtp_layer reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1 @@ -1426,6 +1436,13 @@ def patched_emb_forward(self, input): self.sparse, ) + def _compiles_whole_decoder_layer(self) -> bool: + # Without EP the decoder layer's own forward is compiled, so the layer is one opaque region + # and no marker inside it can delimit anything. With EP that entry is dropped and only the + # methods below it are compiled, which leaves the dispatcher-call boundaries in eager python + # for markers to land on. + return MOE_DECODER_LAYER_FORWARD in self.compile_cfg + def _should_recompute( self, layer_idx: int | None, diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py index a398224af..09cd02bff 100644 --- a/xtuner/v1/model/utils/__init__.py +++ b/xtuner/v1/model/utils/__init__.py @@ -1,9 +1,31 @@ from .checkpointing import apply_gradient_checkpointing from .misc import ModelForwardExtraLogInfo, module_dict_repr +from .selective_checkpointing import ( + KeptCallables, + KeptOps, + RecomputeTarget, + RecomputeTargetMap, + RecomputeUnit, + active_recompute_unit, + apply_selective_checkpointing, + in_recompute_unit, + recompute_unit, + resolve_kept_ops, +) __all__ = [ "apply_gradient_checkpointing", + "apply_selective_checkpointing", + "in_recompute_unit", + "resolve_kept_ops", "module_dict_repr", "ModelForwardExtraLogInfo", + "KeptCallables", + "KeptOps", + "RecomputeTarget", + "RecomputeTargetMap", + "RecomputeUnit", + "active_recompute_unit", + "recompute_unit", ] diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py new file mode 100644 index 000000000..889106495 --- /dev/null +++ b/xtuner/v1/model/utils/selective_checkpointing.py @@ -0,0 +1,359 @@ +"""Shared contract for region-level selective activation checkpointing (SAC). + +This module holds the vocabulary that the SAC layers agree on and nothing else: + +- Users select :class:`RecomputeUnit` members in the model config. +- Model authors declare a :data:`RecomputeTargetMap` saying what each unit they support resolves to + for their architecture. +- The SAC engine turns the selection into a per-op checkpoint policy. + +It lives under ``xtuner.v1.utils`` rather than next to the models because the targets name callables +in ``xtuner.v1.module`` while the unit vocabulary is consumed by ``xtuner.v1.model`` configs; a home +inside either package would make the two import each other. + +**A unit resolves to one of two things, and which one is not a style choice.** The checkpoint policy +is asked about every op that reaches the dispatcher, and it has to answer "is this op part of a kept +unit". There are exactly two ways to know: + +- :class:`KeptOps` names the ops directly. This is the cheapest possible form -- it changes nothing + about compilation -- but it only works when the op is *specific enough to identify the unit on its + own*. Attention qualifies: `flash_attn::_flash_attn_varlen_forward` appears nowhere else. A gate's + ``addmm`` does not. +- :class:`KeptCallables` names callables whose whole body belongs to the unit. Everything they + dispatch is kept. This can express any region, at the price of taking those callables out of the + compiled set -- a marker is ordinary python state and the ops it would cover run as fused kernels, + so a region inside compiled code is invisible either way. + +Prefer :class:`KeptOps`. Reach for :class:`KeptCallables` when no op identifies the unit, and then +name the *smallest* callable that covers it: the compilation given up is the whole callable's, not +the region's. +""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from functools import partial +from typing import Any, TypeAlias + +import torch +import torch.nn as nn +from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts + +from xtuner.v1.utils import log_rank0 +from xtuner.v1.utils.enum_helper import StrEnum + +from .checkpointing import apply_gradient_checkpointing, checkpoint_flattened, install_checkpointing + + +__all__ = [ + "RecomputeUnit", + "KeptOps", + "KeptCallables", + "RecomputeTarget", + "RecomputeTargetMap", + "active_recompute_unit", + "recompute_unit", + "apply_selective_checkpointing", + "in_recompute_unit", + "resolve_kept_ops", +] + + +class RecomputeUnit(StrEnum): + """Semantic units of activation that may be kept resident instead of + recomputed. + + Each member names a *class of sub-structure* whose activations are worth keeping in memory because recomputing it + is expensive relative to what it costs to store. A unit selected by the user means "do not recompute this part"; + everything not selected is recomputed. What a unit resolves to is architecture-specific and is declared per model, + so the same unit can cover different code in different models. + + Members are strings so pydantic round-trips them to readable names in serialized configs. + """ + + SAVE_ATTN = "save_attn" + """Keep the attention kernel's output -- the flash-attention call itself, not the projections + around it. + + This is the narrowest unit and the only one that costs no compilation, because the attention + kernel is a custom op: inductor cannot fuse it, so it is always called as a fallback kernel and + is always visible to the checkpoint policy, compiled or not. + """ + + SAVE_MOE_GATE = "save_moe_gate" + """Keep the MoE router: gating projection, top-k selection, and routing weights.""" + + SAVE_MOE_DISPATCH = "save_moe_dispatch" + """Keep the tensors produced around expert dispatch and combine: the permutation, padding and + unpermutation buffers on either side of the all-to-all. + + The collective itself is always recomputed, never kept. Keeping a collective would elide it from + the recompute pass, which is only sound if nothing it communicates with is replayed; the safe + rule is to replay it. So this unit trades memory for the surrounding tensor work, not for the + communication. + """ + + +@dataclass(frozen=True) +class KeptOps: + """Resolve a unit by op identity: keep these ops wherever they run. + + Costs nothing -- compilation is untouched and no graph breaks are introduced -- because the ops + worth naming this way are exactly the ones inductor cannot fuse, which are the ones that still + reach the dispatcher from inside a compiled region. + + Args: + names (tuple[str, ...]): Qualified op names, e.g. ``"flash_attn::_flash_attn_varlen_forward_v3"``. + A name that is not registered in this build is skipped, so a model may list several + backends' spellings of the same kernel. + """ + + names: tuple[str, ...] + + def __init__(self, *names: str) -> None: + object.__setattr__(self, "names", tuple(names)) + + +@dataclass(frozen=True) +class KeptCallables: + """Resolve a unit by callable: keep everything these callables dispatch. + + The callables are taken out of the compiled set, because a region inside compiled code cannot be + addressed at all -- so the compilation given up is the whole callable's. Name the smallest + callable that covers the unit. + + Args: + names (tuple[str, ...]): Qualified callable names, in the same form ``compile_cfg`` uses. + """ + + names: tuple[str, ...] + + def __init__(self, *names: str) -> None: + object.__setattr__(self, "names", tuple(names)) + + +RecomputeTarget: TypeAlias = KeptOps | KeptCallables +"""What a single :class:`RecomputeUnit` resolves to for one architecture.""" + +RecomputeTargetMap: TypeAlias = dict[RecomputeUnit, RecomputeTarget] +"""Per-model declaration of what each supported :class:`RecomputeUnit` resolves +to. + +Models expose this as a ``default_recompute_cfg`` property, mirroring ``default_compile_cfg``. A unit absent from the +mapping is not supported by that architecture. +""" + + +def active_recompute_unit() -> RecomputeUnit | None: + """Return the unit whose callable is currently executing, if any. + + Meaningful only while a checkpoint policy is running, which is the only caller. Units resolved by + :class:`KeptOps` never set this -- they are recognised from the op itself. + + Returns: + RecomputeUnit | None: The innermost open unit, or None outside one. + """ + return _ACTIVE_UNIT.get() + + +@contextmanager +def recompute_unit(unit: RecomputeUnit) -> Iterator[None]: + """Mark the enclosed call as belonging to ``unit``. + + Entered by the wrapper the engine installs on a :class:`KeptCallables` target, never by model + code. It is a plain ``ContextVar``, which is only readable because the wrapped callable has been + taken out of the compiled set; inside compiled code it would neither be set nor read. + + Args: + unit (RecomputeUnit): The unit the enclosed call belongs to. + """ + token = _ACTIVE_UNIT.set(unit) + try: + yield + finally: + _ACTIVE_UNIT.reset(token) + + +_ACTIVE_UNIT: ContextVar[RecomputeUnit | None] = ContextVar("xtuner_recompute_unit", default=None) + + +def apply_selective_checkpointing( + module: nn.Module, + kept_ops: frozenset = frozenset(), + *, + keeps_any_unit: bool = False, + preserve_rng_state: bool = True, +) -> nn.Module: + """Wrap ``module`` so its forward is recomputed during backward, keeping + the selected units. + + An op is kept when it is one of ``kept_ops``, or when it runs inside a callable the model layer + wrapped with :func:`in_recompute_unit`. Everything else in the layer is recomputed. With nothing + selected this is plain full recompute -- the degenerate case of the same mechanism -- so a + sharding path can call this for every layer ``recompute_ratio`` selects. + + Args: + module (nn.Module): The layer to checkpoint. + kept_ops (frozenset): Op overloads to keep wherever they run, as resolved from the user's + ``recompute_cfg``. Defaults to keeping none. + keeps_any_unit (bool): Whether any unit at all was selected, including ones installed as + callable wrappers rather than op names. Defaults to False, which skips the per-op policy + entirely -- torch's own default already recomputes everything, so running a policy to + reach the same answer would only put a dispatch mode in the way of every op. + preserve_rng_state (bool): Restore the RNG state before recomputing, so dropout and other + stochastic ops replay identically. Defaults to True. + + Returns: + nn.Module: The checkpoint-wrapped layer, transparent to parameter names and ``state_dict``. + """ + if not keeps_any_unit: + return apply_gradient_checkpointing(module, preserve_rng_state=preserve_rng_state) + + return install_checkpointing(module, partial(_run_checkpointed_region, kept_ops, preserve_rng_state)) + + +def in_recompute_unit(unit: RecomputeUnit, func: Any) -> Any: + """Wrap ``func`` so everything it dispatches belongs to ``unit``. + + Installed by the model layer on each callable a + :class:`KeptCallables` target names, together with the + exclusion from compilation that makes the marker readable at all. + + Args: + unit (RecomputeUnit): The unit this callable implements. + func (Any): The callable to wrap. + + Returns: + Any: A callable with the same signature that marks its own execution. + """ + + # No `functools.wraps`: it sets `__wrapped__`, which Dynamo resolves through to the unbound + # function and then loses `self` ("Missing required positional argument: self"). + def in_unit(*args: Any, **kwargs: Any) -> Any: + with recompute_unit(unit): + return func(*args, **kwargs) + + in_unit.__name__ = getattr(func, "__name__", "in_unit") + in_unit.__qualname__ = getattr(func, "__qualname__", "in_unit") + in_unit.__module__ = getattr(func, "__module__", __name__) + return in_unit + + +def resolve_kept_ops(names: tuple[str, ...]) -> frozenset: + """Resolve qualified op names to overloads, skipping ones this build does + not register. + + A model may list several backends' spellings of the same kernel -- flash-attention v2 and v3, a + vendor kernel and its fallback -- and only one of them exists at runtime. + + Args: + names (tuple[str, ...]): Qualified op names such as ``"flash_attn::_flash_attn_varlen_forward_v3"``. + + Returns: + frozenset: The overloads that resolve in this build. + """ + resolved = set() + for name in names: + namespace, _, opname = name.partition("::") + packet = getattr(getattr(torch.ops, namespace, None), opname, None) + if packet is None: + log_rank0.debug(f"Selective checkpointing: {name} is not registered in this build, skipping it.") + continue + resolved.add(packet.default) + return frozenset(resolved) + + +# Ops from these namespaces are never kept, whatever unit they fall in. See `_checkpoint_policy`. +_NEVER_KEPT_NAMESPACES = ("c10d", "_c10d_functional") + +# Mutating ops that leave tensor *values* alone, so a kept unit may contain them. Anything else with +# a mutable schema goes through `_reject_non_replayable_op_in_kept_unit`. +_VALUE_PRESERVING_MUTATING_OPS = frozenset({torch.ops.aten.record_stream.default}) + + +def _run_checkpointed_region( + kept_ops: frozenset, + preserve_rng_state: bool, + original_call: Any, + *args: Any, + **kwargs: Any, +) -> Any: + # `context_fn` must be a module-level function or a `functools.partial` of one: Dynamo's + # checkpoint higher-order op rejects anything else (lambdas, closures, bound methods) with + # `NotImplementedError: ... LazyVariableTracker context_fn`. Keep it that way. + return checkpoint_flattened( + original_call, + *args, + preserve_rng_state=preserve_rng_state, + context_fn=partial(_selective_checkpoint_contexts, kept_ops), + **kwargs, + ) + + +def _selective_checkpoint_contexts(kept_ops: frozenset) -> tuple[Any, Any]: + return create_selective_checkpoint_contexts(partial(_checkpoint_policy, kept_ops)) + + +def _checkpoint_policy(kept_ops: frozenset, ctx: Any, op: Any, *args: Any, **kwargs: Any) -> CheckpointPolicy: + if op not in kept_ops and active_recompute_unit() is None: + return CheckpointPolicy.MUST_RECOMPUTE + + # Keeping a collective would elide it from the recompute pass. That is only correct while the op + # that allocated its destination buffer is kept too, and a unit boundary falling between the + # allocation and the collective would leave the recompute reading an uninitialised buffer -- + # silently, and differently on each rank. + if op.namespace in _NEVER_KEPT_NAMESPACES: + return CheckpointPolicy.MUST_RECOMPUTE + + if op._schema.is_mutable: + _warn_non_replayable_op_in_kept_unit(op) + # Never kept: what the recompute pass would get back from the cache is whatever the last + # writer left behind. + return CheckpointPolicy.MUST_RECOMPUTE + + return CheckpointPolicy.MUST_SAVE + + +def _warn_non_replayable_op_in_kept_unit(op: Any) -> None: + """Name an in-place op inside a kept unit, once, as a lead for a later + failure. + + A kept unit has to survive being replayed on top of its own results: the recompute pass gets the *forward's* + tensors back from the cache, so a read-modify-write op such as ``add_`` would apply its update a second time to a + value that already includes it -- finite loss, no error, wrong gradients. + + Whether that actually happens is not decidable from the op's schema, only from whether the tensor it writes to + was cached. Torch decides exactly that, by version counter, and raises "Tensor cached during selective activation + checkpoint has been mutated" when it does (``torch/utils/checkpoint.py``, ``_VersionWrapper``). Refusing here on + the schema alone would reject units that are perfectly safe -- a library mutating its own buffers, which nothing + in the unit ever cached. + + So this only reports. It runs in the forward pass and names the op, which torch's check -- raised in backward, + naming only the tensor -- cannot. + + Writing through ``out`` is not reported at all: those overwrite the destination with a value that does not depend + on what was there, which is what inductor's extern kernels do, and replaying them is idempotent. + """ + if op in _VALUE_PRESERVING_MUTATING_OPS or _writes_only_through_out(op): + return + if op in _REPORTED_MUTATING_OPS: + return + _REPORTED_MUTATING_OPS.add(op) + log_rank0.warning( + f"Selective checkpointing: the in-place op {op} runs inside a kept unit. It is recomputed rather than kept, " + f"which is safe unless it writes to a tensor the unit did keep -- torch reports that case as " + f'"Tensor cached during selective activation checkpoint has been mutated".' + ) + + +# Reported once per op, not once per layer or per step. +_REPORTED_MUTATING_OPS: set = set() + + +def _writes_only_through_out(op: Any) -> bool: + return all( + argument.name == "out" + for argument in op._schema.arguments + if argument.alias_info is not None and argument.alias_info.is_write + )