diff --git a/docs/design/selective_checkpointing_and_compile.md b/docs/design/selective_checkpointing_and_compile.md new file mode 100644 index 0000000000..7a86d79c5c --- /dev/null +++ b/docs/design/selective_checkpointing_and_compile.md @@ -0,0 +1,266 @@ +# Selective checkpointing under `torch.compile` + +Why unit-level SAC needs the compiler's cooperation, which routes exist, and what each one +measured. All numbers are Qwen3-MoE-30BA3, `ep_size=4` on 8 GPUs, `dispatcher="deepep"`, +torch 2.10, reported as the mean of steps 5-8 of an 8-step run. Two shapes recur: + +- **8k** — `pack_max_length=8192`, no domino. Baseline 8904.4 tgs, 84.32 GB peak. +- **4k** — `pack_max_length=4096` with domino (`intra_layer_micro_batch=2`). Baseline 7812.4 tgs. + +The 4k shape is CPU-bound (GPU busy 55-60%), so every CPU-side cost lands fully in wall time there +and is largely hidden at 8k. The same change reads as −0.9% at 8k and −13.9% at 4k; neither number +is wrong, they measure different regimes. + +## The problem + +The checkpoint policy is asked about every op that reaches the dispatcher, and it has to answer one +question: does this op belong to a unit the user asked to keep. There are exactly two ways to know, +and which one applies is a property of the unit, not a style choice. + +**Naming the op.** Works when the op is specific enough to identify the unit on its own — +`flash_attn::_flash_attn_varlen_forward` appears nowhere else in the model. Costs nothing: the +compile set is untouched, no graph breaks, and it works identically inside and outside compiled +code, because the ops worth keeping are exactly the ones inductor cannot fuse and those are the ones +still reaching the dispatcher. + +**Naming the callable.** Works for anything, at a price. The marker is ordinary python state — a +`ContextVar` set around the callable — and compiled code neither writes nor reads it: reading one +inside a `fullgraph=True` region is a hard compile error, and a value written during tracing is +traced away. So the callable has to leave the compiled set, via `torch._dynamo.disable`. + +This is not a visibility problem. With the whole layer compiled the policy is still consulted +**117597 times per run**, because ops that inductor cannot fuse still go through the dispatcher. +Measured on this model, 26 distinct ops reach it, among them `flash_attn::_flash_attn_varlen_forward_v2`, +`moe::m_grouped_gemm`, `moe::permute`/`unpermute` and `aten::mm.out`; 61% of the calls are +`aten::record_stream`. What is missing for a callable-scoped unit is the region identity, not the op. + +## The route taken: op identity where possible, withdrawal where not + +`RecomputeTargetMap` binds each unit a model supports to one of the two resolutions. MoE declares +`SAVE_ATTN` as `KeptOps` and the gate and dispatch stages as `KeptCallables`. + +Two mechanics of the withdrawal are easy to get wrong: + +- Removing an entry from `compile_cfg` does **not** keep it out of the compiled set. Dynamo inlines + it into whichever compiled caller reaches it, and the marker is traced away exactly as before. + `torch._dynamo.disable` is what makes it run in python. +- A disabled callee inside a `fullgraph=True` region is a hard error + (`Skip inlining torch.compiler.disable()d function`), not a split. The surviving entries are + therefore relaxed to `fullgraph=False`. + +Measured, 8k, `save_attn`, when it was still resolved by withdrawing the callable that encloses it: + +| | tgs | peak | compiled graphs / captured calls | kept | +|---|---|---|---|---| +| no unit | 8904.4 | 84.32 GB | 6 / 154 | — | +| withdrawal, marker never fired | 8883.1 | 84.32 GB | 6 / 154 | nothing | +| withdrawal, marker firing | 8021.7 | 112.14 GB | 4 / 26 | 42112 tensors | + +−9.9% for the same unit that op identity delivers for free. The cost is concentrated, not diffuse: +`_pre_moe_forward` accounts for 128 of the 154 captured calls, because it exists to give +`torch.compile` the ops on either side of attention. Withdrawing it is most of what an MoE layer +compiles. That is why attention resolves by op identity and why `KeptCallables` should name the +*smallest* callable that covers the unit: the compilation given up is the whole callable's. + +## What the shipped units cost + +Measured on the shipped resolutions, 16k domino (`pack_max_length=8192`, +`intra_layer_micro_batch=2`), `ep_size=4`, deepep. **Three independent runs per setting**, because +one run per setting is not enough to say anything about throughput here: + +| unit | tgs (3 runs) | mean | peak allocated | peak reserved | +|---|---|---|---|---| +| none | 10206.0 / 9990.6 / 10205.9 | 10134.2 | 84.4-85.3 GB | 108.5-109.3 GB | +| `save_attn` | 10165.6 / 10204.6 / 10071.2 | 10147.2 | 90.4-91.2 GB | 114.4-115.3 GB | +| `save_moe_gate` | 10095.2 / 10179.5 / 10144.4 | 10139.7 | 93.7-94.5 GB | 117.7-118.7 GB | +| `save_moe_dispatch` | OOM | — | — | — | + +**Throughput is unchanged.** The three means are within 0.1% of each other, while the baseline's own +three runs span 2.1%. Any single-run comparison of these units reads as ±2% in whichever direction +the run-to-run variance happened to fall, and means nothing. + +**Memory is the reproducible effect**: `save_attn` costs +6.0 GB allocated, `save_moe_gate` +9.2 GB, +each within ±0.4 GB across runs. + +So a unit here buys nothing on this model and costs memory. That is not a statement about the +mechanism -- it is a statement about which activations an MoE layer's recompute is actually spent +on, and the census below says why. + +`save_moe_dispatch` is not usable at this shape at all. It keeps the permutation and padding buffers +on both sides of the all-to-all, the widest tensors in the layer, and domino already runs at 109 GB +reserved of 140 GB. It OOMs during the first step -- and also OOMs without domino, and at +`pack_max_length=4096`, where it peaks at 111.9 GB against the baseline's 85.9 GB. It is verified +numerically for one step (see below) and declared for smaller models, not as a default here. + +## Numerical correctness + +Checkpointing changes only the backward pass, so **step-1 loss must be bit-identical** whatever is +kept. It is, across every setting and both shapes -- `2.46262765` at 8k, `2.38410378` at 4k, +including the `save_moe_dispatch` run that OOMs immediately afterwards. + +Gradients need a noise floor to interpret, which is why the baseline was run twice under identical +settings: + +| | step-1 loss | step-1 grad_norm | vs baseline | +|---|---|---|---| +| baseline | 2.46262765 | 24.39401245 | — | +| baseline, second run | 2.46262765 | 24.39329147 | 3.0e-5 | +| `save_attn` | 2.46262765 | 24.39325905 | 3.1e-5 | +| `save_moe_gate` | 2.46262765 | 24.39400864 | 1.6e-8 | + +Every unit sits at or below the floor two identical runs produce. The floor itself comes from +reduction ordering in the grouped GEMM and the all-to-all, and exists with no unit selected. + +Measured eager (`torch_compile=False`, all2all, no domino) so that compilation cannot be the source +of a difference; the compiled path is covered by the throughput runs above. + +## Whether a unit repays its cost + +On this model: no. Every unit raises peak reserved -- keeping activations is the trade -- and none +of them buys back measurable throughput. The census below says why. + +A resident-activation census under the shipped configuration (domino, 8k, whole layer checkpointed) +found **245 tensors totalling 15.27 GiB**: + +| producer | count | bytes | +|---|---|---| +| `LogSoftmaxBackward0 (16384, 151936) float32` | 1 | 9.27 GiB | +| layer-boundary `(1, 8192, 2048)` bf16 | 94 | 2.94 GiB | +| `TBackward0 (2048, 151936)` bf16 | 1 | 0.58 GiB | +| everything else (60 producers) | 149 | ~2.5 GiB | + +The loss logits alone are 61% of the resident set, and no recompute unit addresses them. Nothing +inside the MoE layers survives — the whole-layer checkpoint is doing its job. That bounds how much +any unit-level policy can win here. + +## Routes measured + +### Keeping the expert GEMM by op identity + +The same resolution `SAVE_ATTN` uses, pointed at the ops that dominate an MoE layer's recompute: + +| kept | tgs (8k) | peak | captured | +|---|---|---|---| +| baseline | 8904.4 | 84.32 GB | 154 | +| `flash_attn::_flash_attn_varlen_forward_v2` | 8828.2 | 87.31 GB | 154 | +| `moe::m_grouped_gemm` | **9477.3 (+6.4%)** | 105.25 GB | 154 | +| `moe::m_grouped_gemm` + `permute` + `unpermute` | **9616.7 (+8.0%)** | 122.61 GB | 154 | + +Not declared as a unit here because +20 to +38 GB is not a trade a user can currently tune — the +selection is per-model, not per-layer. It is the obvious next unit once the selection can say +"in the first N layers". The limit of op identity is expressiveness: it cannot reach anything fused +into a generated kernel, and it cannot scope a unit to part of the model. It is what torchtitan does +(`torchtitan/distributed/activation_checkpoint.py`), with a preset op list and no region concept. + +### Withdrawing a smaller callable + +Making the excluded callable smaller does not make the exclusion cheaper in proportion, because +`torch._dynamo.disable` breaks the graph once **per level of the inline stack it has to unwind**, +not once per call: + +| cut point | inline depth | breaks | captured | tgs (8k) | +|---|---|---|---|---| +| `MultiHeadAttention.forward` | 1 | 2 | 60 | 8101.7 | +| `attn_imp.flash_attention` | 2 | — | 96 | 8664.5 | +| `flash_attn_varlen_func` | 3 | 4 | 104 | 8660.3 | + +Cutting deeper leaves more code compiled but unwinds more levels; the two effects cancel. A +withdrawal is cheapest when its boundary is *shallow*, near the compile entry point — the opposite of intuition. + +The two resolutions also cannot be made equivalent. Even at the tightest cut, the withdrawal keeps +`aten::alias` alongside the attention kernel, because a `dynamo.disable`d segment runs eagerly and +eager execution dispatches bookkeeping ops that do not exist in the compiled path. + +### Letting the marker run while compiling, and taking the graph break + +Instead of withdrawing the callable, let the marker execute during tracing and accept the graph +break Dynamo takes at it. The marker becomes live and the policy sees the unit — but only for ops +that reach the dispatcher, so it keeps 752 tensors where withdrawing the callable keeps 42112. + +Measured, 8k, `save_attn`: 8578.4 tgs, 86.03 GB, 11 graphs / 90 captured, 5 breaks. Cheaper than +withdrawing the callable and more expensive than op identity, with a third semantics again. Not +adopted because "keep this unit" should not silently mean "keep the handful of ops in it that +inductor happened not to fuse". + +### `fx_traceback.annotate` plus a joint-graph pass + +The only marker that survives compilation intact. Dynamo executes it during tracing +(`_dynamo/variables/ctx_manager.py`), stamps `node.meta["custom"]`, and AOT carries it through +decomposition and into the backward nodes. Verified on GPU under `fullgraph=True`: 6 graphs, 154 +captured calls, **0 graph breaks** — identical to the baseline — with the joint pass seeing 410 +annotated nodes and pinning 386. + +And it changes nothing. Peak memory stayed at 84.32 GB, byte for byte. + +The reason is a level mismatch. `torch.utils.checkpoint(use_reentrant=False)` is implemented with +saved-tensor hooks: its `pack_hook` replaces every tensor autograd saves inside the region with a +holder and frees it. A compiled region is an ordinary `autograd.Function`, so the tensors the +partitioner chose to save are handed to `ctx.save_for_backward` and discarded like everything else. +The partitioner's decision is made and honoured — its *product* is thrown away. + +So the tag route requires the partitioner to be the only decider, which means removing the outer +checkpoint. + +### Removing the outer checkpoint + +Then the partitioner governs, and the tags work. It also costs 61.8 GiB. + +| | resident tensors | resident bytes | +|---|---|---| +| whole layer checkpointed | 245 | 15.27 GiB | +| no checkpoint, everything forced to `MUST_RECOMPUTE` | 1706 | 77.04 GiB | + +`MUST_RECOMPUTE` applies *within* one joint graph. It cannot recompute across a graph boundary, +because the backward of each graph needs that graph's inputs. Domino + EP splits a layer into many +graphs, and the boundaries land on the widest tensors in the model: one `(49152, 2048)` per layer +(9.00 GiB total), one `(65536, 768)` per layer (4.50 GiB), 96 layer-boundary `(1, 8192, 2048)` +(3.00 GiB). Of the 77.04 GiB, 18.33 GiB has no `grad_fn` at all — pure graph-boundary input. + +Forcing every node to recompute changes nothing versus a `budget=0.0`: 131.14 GB against 131.34 GB. +The memory that cannot be reclaimed is not the partitioner's to reclaim. + +`torch._functorch.config.activation_memory_budget` on the same shape, with no checkpoint: + +| budget | tgs | peak | +|---|---|---| +| baseline (whole layer checkpointed) | 8904.4 | 84.32 GB | +| 0.0 | 10032.5 (+12.7%) | 107.71 GB | +| 0.25 | **10282.0 (+15.5%)** | 111.68 GB | +| 0.5 | 9750.1 (+9.5%) | 113.25 GB | + +The fastest result measured anywhere, for +23 GB and no way back down: `budget=0.0` is the floor and +it is 23.4 GB above the baseline. The curve is not monotone (0.25 beats both neighbours) and that is +unexplained; it is a single point per budget, not a characterised curve. + +**Under domino this route is unusable.** Domino already needs ~19 GB more reserved, and removing the +checkpoint adds ~24 GB more, which crosses the 140 GB device limit: `budget=0.0` runs at 1925 tgs +(5.4× slower than the checkpointed baseline) with 134.60 GB reserved, and 0.25 and 0.5 both OOM. + +## The route that would be complete, and what blocks it + +If the checkpoint became a `tag_activation_checkpoint` HOP — that is, if `torch.compile` wrapped the +checkpoint call rather than sitting inside it — the two mechanisms collapse into one. The policy +result is written to `node.meta["recompute"]` during tracing (`torch/utils/checkpoint.py`) and the +partitioner consumes it; there is no second decider to overrule it, no dispatch mode at runtime, and +no callable to withdraw from the compile set. + +The HOP only forms if the checkpointed region speculates into a single subgraph. On the domino path +it never does: `torch_all2all.py` reads `permuted_hidden_states.grad_fn` to register a backward +pre-hook for the CUDA-event choreography, and Dynamo cannot trace a tensor's `grad_fn`. Speculation +fails, everything traced is discarded, and the graph before the checkpoint call is empty — measured +at 30B as `unique_graphs=0, calls_captured=0`, a silent degradation to eager rather than an error. + +torchtitan reaches this topology because its MoE dispatch is a custom op (`deepep.dispatch.default`), +which enters the graph whole. Making xtuner's dispatcher expressible the same way — replacing the +`grad_fn.register_prehook` choreography — is the one change that unlocks it. + +## Memory wall, for context + +Independent of checkpointing: domino at 16384 is 4.44× *slower* than not using it (14.20 s vs +3.199 s per step) purely because reserved memory reaches 134.43 GB of 140.06 GB and the caching +allocator thrashes. At 8192 the same code is **16.1% faster** with domino than without (1.583 s vs +1.838 s). `expandable_segments` does not help, so this is total demand rather than fragmentation, and +switching between the all2all and deepep backends does not help either (deepep is 7.1% faster only +when domino is off). + +Any activation-memory work on the domino path should check headroom before it profiles kernels. diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py index 7e55511d72..a06e7c6541 100644 --- a/tests/model/test_recompute.py +++ b/tests/model/test_recompute.py @@ -1,18 +1,52 @@ -"""Gradient checkpointing regression tests. +"""Gradient checkpointing and recompute-unit regression tests. TestCheckpointWrapper test_wrapper_is_transparent_to_state_dict_and_attributes: 包裹后参数名/state_dict/属性访问不变。 test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。 test_wrapper_forwards_container_protocols: 被包裹模块的 len/iter/in/索引协议在包裹后仍可用。 test_wrapper_does_not_claim_protocols_the_module_lacks: 被包裹模块没有的协议不会出现在包裹层上。 +TestRecomputeCfgResolution + test_unset_cfg_keeps_full_recompute: `None` 不改变显存行为，解析为不留驻。 + test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。 + test_explicit_units_select_only_themselves: 显式 list 只选中对应 unit。 + test_string_units_are_accepted: 配置文件里的字符串能解析成 RecomputeUnit。 + test_unsupported_unit_is_rejected: 模型不支持的 unit 在构造时报错并列出支持项。 + test_disable_propagates_into_nested_configs: `False` 递归关闭嵌套子模型配置。 + test_disable_reaches_every_sub_model_of_a_real_compose_config: 真实 compose 配置的三个子配置都被关闭。 + test_units_round_trip_through_json: enum 序列化成可读字符串并能读回。 +TestDeclaredTargets + test_declared_targets_resolve: 声明表里的 op 名与 callable 名都能解析到真实对象。 + test_no_unit_names_the_method_that_holds_most_compilation: 没有 unit 点名承载最多编译的那个方法。 +TestUnitCostIsProportionate + test_an_op_identity_unit_costs_no_compilation: KeptOps 不改动编译集合。 + test_a_callable_unit_keeps_its_callers_compiled: KeptCallables 只退出自身，调用者仍编译。 + test_no_unit_withdraws_the_method_that_holds_most_compilation: 没有 unit 撤出编译占比最大的方法。 + test_attention_is_kept_by_op_identity: attention 走 op identity 而非撤出 callable。 """ +import ast +import inspect +import pydoc +import textwrap + import pytest import torch from torch import nn -from xtuner.v1.model.utils import apply_gradient_checkpointing +from xtuner.v1.model.base import BaseModel, TorchCompileOption, XTunerBaseModelConfig, _disable_nested_switch +from xtuner.v1.model.compose.qwen3_vl import Qwen3VLMoE30BA3Config +from xtuner.v1.model.dense.dense import DENSE_RECOMPUTE_CFG +from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG, MoE, MoEConfig +from xtuner.v1.model.utils import ( + KeptCallables, + KeptOps, + RecomputeUnit, + apply_gradient_checkpointing, + resolve_kept_ops, +) +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.router import NoAuxRouterConfig class _KeywordOnlyBlock(nn.Module): @@ -107,3 +141,244 @@ def test_wrapper_does_not_claim_protocols_the_module_lacks(self): assert not hasattr(type(wrapped), "__len__") with pytest.raises(TypeError, match="CheckpointWrapper"): len(wrapped) + + +def _build_tiny_moe_config(**overrides) -> MoEConfig: + """A MoE small enough to instantiate on the meta device in a config-only test.""" + router_config = NoAuxRouterConfig( + scoring_func="sigmoid", + router_scaling_factor=1.0, + n_group=1, + topk_group=1, + norm_topk_prob=True, + ) + return MoEConfig( + vocab_size=256, + max_position_embeddings=128, + pad_token_id=0, + eos_token_id=0, + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=16), + tie_word_embeddings=False, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + first_k_dense_replace=1, + hidden_factor=1.0, + moe_intermediate_size=64, + router=router_config, + compile_cfg=False, + **overrides, + ) + + +def _resolve_units(**overrides) -> set: + with torch.device("meta"): + return set(MoE(config=_build_tiny_moe_config(**overrides))._selected_recompute_units) + + +class _NestedProbeConfig(XTunerBaseModelConfig): + """Stand-in for a sub-model config, as a compose model nests one.""" + + +class _ProbeConfig(XTunerBaseModelConfig): + text_config: XTunerBaseModelConfig + + +class _ProbeModel(BaseModel): + """A model that contributes nothing but ``BaseModel.__init__``'s config resolution.""" + + config: _ProbeConfig + + +class TestRecomputeCfgResolution: + def test_unset_cfg_keeps_full_recompute(self): + # `None` must not change the memory profile of an existing training run: unlike `compile_cfg`, + # it resolves to "retain nothing" rather than to the model's declared units. + assert _resolve_units(recompute_cfg=None) == set() + + def test_true_selects_every_supported_unit(self): + assert _resolve_units(recompute_cfg=True) == set(MOE_RECOMPUTE_CFG) + + def test_explicit_units_select_only_themselves(self): + assert _resolve_units(recompute_cfg=[RecomputeUnit.SAVE_MOE_GATE]) == {RecomputeUnit.SAVE_MOE_GATE} + + def test_string_units_are_accepted(self): + # Configs arrive as JSON/py files where units are written as plain strings. + assert _resolve_units(recompute_cfg=["save_attn"]) == {RecomputeUnit.SAVE_ATTN} + + def test_unsupported_unit_is_rejected(self): + # A model declaring no units cannot honour any selection, so this is a user configuration + # error rather than something to silently drop. It surfaces at construction, before the run + # spends anything on materializing and sharding weights. + with pytest.raises(ValueError, match="does not support"): + _ProbeModel(_ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=[RecomputeUnit.SAVE_ATTN])) + + def test_disable_propagates_into_nested_configs(self): + # A sub-model resolves its own switch, so `False` on the outer config only means something + # if it reaches the nested ones. `compile_cfg` must stay untouched: the walk is per switch. + config = _ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=False) + + model = _ProbeModel(config) + + assert model._selected_recompute_units == set() + assert config.text_config.recompute_cfg is False + assert config.text_config.compile_cfg is None + + def test_disable_reaches_every_sub_model_of_a_real_compose_config(self): + # The probe above has one nested config; a shipped compose config has three, one of them a + # further-derived MoE config. Exercised on the config walk rather than through the model, + # because constructing a 30B compose model is the expensive part and contributes nothing: + # what can regress here is which nested configs the walk reaches. + config = Qwen3VLMoE30BA3Config(recompute_cfg=False) + + _disable_nested_switch(config, "recompute_cfg") + + for sub_config in (config.vision_config, config.projector_config, config.text_config): + assert sub_config.recompute_cfg is False + + def test_units_round_trip_through_json(self): + # Trainer resume reads the config back, and serialized runs are read by humans, so units + # must survive as their readable names. + config = _build_tiny_moe_config(recompute_cfg=[RecomputeUnit.SAVE_ATTN, RecomputeUnit.SAVE_MOE_GATE]) + + dumped = config.model_dump(mode="json")["recompute_cfg"] + assert dumped == ["save_attn", "save_moe_gate"] + + restored = _build_tiny_moe_config(recompute_cfg=dumped) + assert restored.recompute_cfg == [RecomputeUnit.SAVE_ATTN, RecomputeUnit.SAVE_MOE_GATE] + + +def _recorded_markers(func) -> set[str]: + """Marker names a function passes to ``checkpoint_record``.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + return { + node.args[0].value + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "checkpoint_record" + and node.args + and isinstance(node.args[0], ast.Constant) + } + + +def _region_coverage(func) -> dict[str, set[str]]: + """Which of the layer's own operations each marker region encloses. + + Regions are keyed by the shared prefix of a ``.begin`` / ``.end`` pair, and an operation is a call on + ``self`` -- ``self.experts(...)``, ``self.dispatcher.dispatch(...)``. Calls on tensors are ignored: a region is + defined by the sub-modules it covers, not by the reshapes threaded between them. + """ + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + # `ast.walk` is breadth-first; the marker state machine needs source order. + nodes = sorted( + (node for node in ast.walk(tree) if isinstance(node, ast.Call)), + key=lambda node: (node.lineno, node.col_offset), + ) + + coverage: dict[str, set[str]] = {} + active: set[str] = set() + for node in nodes: + name = _called_name(node) + if name == "checkpoint_record" and node.args and isinstance(node.args[0], ast.Constant): + region, _, edge = node.args[0].value.rpartition(".") + if edge == "begin": + active.add(region) + coverage.setdefault(region, set()) + elif edge == "end": + active.discard(region) + elif name is not None and name.startswith("self."): + for region in active: + coverage[region].add(name.removeprefix("self.")) + return coverage + + +def _called_name(node: ast.Call) -> str | None: + """Dotted name of a call target, e.g. ``self.dispatcher.dispatch``.""" + parts: list[str] = [] + target: ast.expr = node.func + while isinstance(target, ast.Attribute): + parts.append(target.attr) + target = target.value + if not isinstance(target, ast.Name): + return None + parts.append(target.id) + return ".".join(reversed(parts)) + + +class TestDeclaredTargets: + @pytest.mark.parametrize("target_map", [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG], ids=["moe", "dense"]) + def test_declared_targets_resolve(self, target_map): + # A renamed method or op would not fail anywhere at runtime on its own: the unit would + # simply keep nothing and the region would stay recomputed, silently costing the memory the + # user asked to keep. Resolve every name a model declares so a rename fails here instead. + for unit, target in target_map.items(): + if isinstance(target, KeptOps): + # A build registers only one flash-attention version, so *some* name must resolve + # rather than all of them. + assert resolve_kept_ops(target.names), f"{unit} names no op that resolves: {target.names}" + else: + for name in target.names: + assert pydoc.locate(name) is not None, f"{unit} names {name}, which does not resolve" + + def test_no_unit_names_the_method_that_holds_most_compilation(self): + # `_pre_moe_forward` exists to give `torch.compile` the ops on either side of attention. + # Naming it would withdraw most of what an MoE layer compiles, which is the cost this + # design exists to avoid -- attention is kept by op identity and the gate by its own + # callable instead. + for unit, target in MOE_RECOMPUTE_CFG.items(): + if isinstance(target, KeptCallables): + assert _PRE_MOE_FORWARD not in target.names, f"{unit} withdraws {_PRE_MOE_FORWARD}" + + +_PRE_MOE_FORWARD = "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer._pre_moe_forward" + + +class TestUnitCostIsProportionate: + """What a unit costs follows from how it is resolved, and nothing costs a withdrawn method. + + An op-identity unit changes nothing about compilation. A callable unit is excluded from the + compiled set, which a compiled caller sees as a graph break -- so the callers are relaxed to + `fullgraph=False`, but they stay compiled. No unit withdraws `_pre_moe_forward`, where most of + an MoE layer's compilation lives. + """ + + @staticmethod + def _compile_cfg(**overrides) -> dict[str, TorchCompileOption]: + # The tiny config disables compilation; this asks what would be compiled if it did not. + config = _build_tiny_moe_config(**overrides).model_copy(update={"compile_cfg": None}) + with torch.device("meta"): + return MoE(config=config).compile_cfg + + def test_an_op_identity_unit_costs_no_compilation(self): + assert self._compile_cfg(recompute_cfg=[RecomputeUnit.SAVE_ATTN]) == self._compile_cfg() + + def test_a_callable_unit_keeps_its_callers_compiled(self): + unit = RecomputeUnit.SAVE_MOE_GATE + # Relaxed, not removed: the caller still compiles, it just splits at the excluded callee. + # Being absent from `compile_cfg` is not enough to be outside the compiled set, because + # Dynamo inlines a callee into whichever compiled caller reaches it. + baseline = self._compile_cfg() + relaxed = self._compile_cfg(recompute_cfg=[unit]) + + assert set(relaxed) == set(baseline), "a unit must not remove a caller from the compiled set" + assert not any(option.get("fullgraph") for option in relaxed.values()) + + def test_no_unit_withdraws_the_method_that_holds_most_compilation(self): + for unit in MOE_RECOMPUTE_CFG: + assert _PRE_MOE_FORWARD in self._compile_cfg(recompute_cfg=[unit]), f"{unit} withdrew {_PRE_MOE_FORWARD}" + + def test_attention_is_kept_by_op_identity(self): + # The attention kernel is a custom op, so it reaches the policy from inside a compiled + # region -- which is why this unit costs nothing. + config = _build_tiny_moe_config(recompute_cfg=[RecomputeUnit.SAVE_ATTN]) + with torch.device("meta"): + model = MoE(config=config) + assert model.kept_ops + assert model.keeps_any_recompute_unit diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py index 39e8a1d6d8..945e6895a5 100644 --- a/tests/model/test_selective_checkpointing.py +++ b/tests/model/test_selective_checkpointing.py @@ -17,11 +17,15 @@ test_kept_unit_matches_full_recompute_under_compile: compile 下同上，且不触发 cached-tensor-mutated。 """ +import subprocess +import sys + import pytest import torch from torch import nn from xtuner.v1.model.utils import ( + KeptOps, RecomputeUnit, apply_selective_checkpointing, in_recompute_unit, @@ -149,6 +153,19 @@ def test_unit_marker_outside_a_region_is_noop(self): assert inputs.grad is not None +class TestContractLayering: + def test_module_layer_imports_the_contract_without_the_model_layer(self): + # 契约之所以在 xtuner/v1/utils 而不是挨着 engine，是因为它命名的 callable 在 + # xtuner/v1/module 里：一旦契约里出现指向 model/ 的 import，这条独立导入就会变成 + # 循环导入而失败。必须用干净的解释器：同进程里 xtuner.v1.model 早就被导入了。 + result = subprocess.run( + [sys.executable, "-c", "import xtuner.v1.module.decoder_layer.moe_decoder_layer"], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + class TestUnsupportedUnits: def test_in_place_op_in_kept_unit_is_recomputed_not_refused(self): # in-place op 本身不危险，危险的是它写到了被 unit 留驻的张量上——那由 torch 的 @@ -163,3 +180,16 @@ def test_mutating_a_kept_tensor_is_caught_by_torch(self): def test_in_place_op_outside_kept_unit_is_fine(self): _run(_InPlaceOutsideBlock(), keeps_any_unit=True) + + +class TestDeclarations: + def test_kept_ops_and_callables_are_distinct_targets(self): + # 两种解析方式的代价完全不同（一个零编译代价，一个要退出编译集合），所以类型必须能区分。 + from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG + + assert isinstance(MOE_RECOMPUTE_CFG[RecomputeUnit.SAVE_ATTN], KeptOps) + assert set(MOE_RECOMPUTE_CFG) == { + RecomputeUnit.SAVE_ATTN, + RecomputeUnit.SAVE_MOE_GATE, + RecomputeUnit.SAVE_MOE_DISPATCH, + } diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 9e811610ab..c580163544 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -19,7 +19,6 @@ import torch.nn as nn import torch.nn.functional as F from cyclopts import Parameter -from more_itertools import consume from pydantic import BaseModel as PydanticBaseModel from pydantic import ConfigDict, Field, computed_field, model_validator from pydantic.fields import FieldInfo @@ -62,7 +61,15 @@ set_async_save_process_qos, ) -from .utils import ModelForwardExtraLogInfo +from .utils import ( + KeptCallables, + KeptOps, + ModelForwardExtraLogInfo, + RecomputeTargetMap, + RecomputeUnit, + in_recompute_unit, + resolve_kept_ops, +) logger = get_logger() @@ -143,6 +150,21 @@ class XTunerBaseModelConfig(PydanticBaseModel): "`dict[str, TorchCompileOption]`: Customize the compile option", ), ] = None + # `activation_offload_cfg` belongs here, next to `recompute_cfg` and sharing its `RecomputeUnit` + # vocabulary and per-model interval declarations: offloading applies to the regions SAC keeps + # resident, since recomputed regions never land in memory to begin with. Not declared until it is + # implemented -- a config field nothing reads is a switch that silently does nothing. + recompute_cfg: Annotated[ + list[RecomputeUnit] | bool | None, + Parameter( + group="model", + help="Which activation regions stay resident instead of being recomputed, inside the layers " + "selected by `fsdp_cfg.recompute_ratio`. " + "`None` | `False`: Recompute everything, " + "`True`: Keep every region the model declares in `default_recompute_cfg`, " + "`list[RecomputeUnit]`: Keep exactly the listed regions", + ), + ] = None hf_key_mapping: Annotated[dict[str, str] | None, "Remapping hf key based on the `to_hf_key_list`"] = None dcp_ignore_frozen_params: bool = True lm_loss_cfg: BaseLossConfig = CELossConfig() @@ -538,6 +560,36 @@ def _save_file( save_file(tensors, filename, metadata=metadata) +def _disable_nested_switch(obj: Any, field_name: str) -> None: + """Turn a tri-state feature switch off on every config reachable from + ``obj``. + + Sub-model configs carry their own copy of switches such as ``compile_cfg`` and ``recompute_cfg``, and each + sub-model resolves its own. Turning a feature off on the outer config therefore only means something if the + ``False`` reaches them, which is what this walk does. + + Traversal stops at configs that do not declare ``field_name``: a config that does not take part in the feature + cannot hide a participant behind it, and stopping keeps unrelated sub-configs out of the walk. + + Args: + obj (Any): Config, container, or leaf value to walk. + field_name (str): Name of the switch field to set to ``False``. + """ + if isinstance(obj, PydanticBaseModel): + if not hasattr(obj, field_name): + return + setattr(obj, field_name, False) + for nested_field in type(obj).model_fields: + _disable_nested_switch(getattr(obj, nested_field), field_name) + elif isinstance(obj, Mapping): + for value in obj.values(): + _disable_nested_switch(value, field_name) + # str & bytes are Iterables of themselves and would recurse forever. + elif isinstance(obj, Iterable) and not isinstance(obj, (str, bytes)): + for value in obj: + _disable_nested_switch(value, field_name) + + class BaseModel(nn.Module): load_spec_mapping: dict[str, LoadSpec] = {} fsdp_mesh: DeviceMesh | None = None @@ -556,6 +608,10 @@ def __init__(self, config: XTunerBaseModelConfig): self._pending_async_hf: AsyncHFSaveHandle | None = None self._async_hf_resources: AsyncHFResources | None = None + # Recompute resolves first: selecting a unit withdraws the callables enclosing its region + # from the compiled set, so the compile config depends on the answer. + self._selected_recompute_units: set[RecomputeUnit] = set() + self._kept_ops, self._kept_callables = self._resolve_recompute_cfg(self.config) self._compile_cfg = self._resolve_compile_cfg(self.config) self._float8_handler: Float8Handler | None = None @@ -984,6 +1040,23 @@ def compile_cfg(self) -> dict[str, TorchCompileOption]: return _compile_cfg + @property + def default_recompute_cfg(self) -> RecomputeTargetMap: + """Marker intervals this architecture can keep resident, keyed by + semantic unit. + + This is the model author's vocabulary: it declares which :class:`RecomputeUnit` s the architecture supports and + where each one lives, not which of them are worth enabling. A model that has no ``checkpoint_record`` markers, + or whose markers are all inert in the setup it ships with, declares nothing. + + Like ``default_compile_cfg``, an override must be answerable from ``self.config`` alone: it is read while + ``BaseModel.__init__`` resolves the user's selection, which is before the subclass has built its layers. + + Returns: + RecomputeTargetMap: Supported units mapped to what implements them for this architecture. + """ + return {} + @property def kept_ops(self) -> frozenset: """Op overloads kept resident wherever they run, as selected by @@ -992,20 +1065,16 @@ def kept_ops(self) -> frozenset: Returns: frozenset: Overloads to keep. Empty means no op-identity unit was selected. """ - return frozenset() + return self._kept_ops @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 + return bool(self._selected_recompute_units) @property def float8_handler(self): @@ -2573,14 +2642,14 @@ def _resolve_compile_cfg( custom_cfg = config.compile_cfg if custom_cfg is False: - self._disable_compile_cfg(self.config) + _disable_nested_switch(self.config, "compile_cfg") return {} # torch.compile is not supported on NPU if DEVICE == "npu": if custom_cfg is not False: log_rank0.warning("torch.compile is not supported on NPU, disabling torch.compile.") - self._disable_compile_cfg(self.config) + _disable_nested_switch(self.config, "compile_cfg") return {} if custom_cfg is True or custom_cfg is None: @@ -2588,27 +2657,141 @@ def _resolve_compile_cfg( else: compile_cfg = custom_cfg - return compile_cfg - - def _disable_compile_cfg(self, obj): - if isinstance(obj, PydanticBaseModel) and hasattr(obj, "compile_cfg"): - obj.compile_cfg = False - consume(self._disable_compile_cfg(getattr(obj, x)) for x in obj.__class__.model_fields) - elif isinstance(obj, Mapping): - consume(map(self._disable_compile_cfg, obj.values())) - # str&bytes are special Iterable, need to exclude it, otherwise it will infinite loop - elif isinstance(obj, Iterable) and not isinstance(obj, (str, bytes)): - consume(map(self._disable_compile_cfg, obj)) - else: - return + return self._without_compiled_selected_regions(compile_cfg) + + def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> tuple[frozenset, set[str]]: + selected = config.recompute_cfg + + # `False` has to reach the nested sub-model configs, which a compose model builds after this + # returns and which resolve their own switch against them. + if selected is False: + _disable_nested_switch(self.config, "recompute_cfg") + return frozenset(), set() + + # `None` means "keep the memory profile of plain full recompute", not "use the model default" as it does + # for `compile_cfg`. `default_recompute_cfg` is the vocabulary of what an architecture *can* keep resident, + # not a recommendation: keeping a unit trades memory for speed, and retaining everything can even exceed + # the peak of not checkpointing at all, because SAC storage duplicates what autograd already holds. Only the + # user knows which side of that trade they want, so an unset config changes nothing. + if selected is None: + return frozenset(), set() + + supported = self.default_recompute_cfg + + # `True` asks for whatever the model offers, so an empty vocabulary is not a user error -- + # but it does mean the request is a no-op, which is worth saying out loud rather than + # letting the run look configured when it is not. + if selected is True and not supported: + log_rank0.warning( + f"`recompute_cfg=True` has no effect: {type(self).__name__} declares no recompute units, so every " + "region is recomputed." + ) + units = list(supported) if selected is True else selected + + op_names: list[str] = [] + callables: set[str] = set() + for unit in units: + if unit not in supported: + supported_desc = ", ".join(sorted(supported)) if supported else "none" + raise ValueError( + f"`recompute_cfg` selects {unit!r}, which {type(self).__name__} does not support. " + f"Units supported by this model: {supported_desc}. Note that a compose model declares no units " + "of its own -- set `recompute_cfg` on the sub-model config that owns the layers instead." + ) + target = supported[unit] + if isinstance(target, KeptOps): + op_names.extend(target.names) + else: + callables.update(target.names) + self._selected_recompute_units.add(unit) + return resolve_kept_ops(tuple(op_names)), callables + + def _without_compiled_selected_regions( + self, compile_cfg: dict[str, TorchCompileOption] + ) -> dict[str, TorchCompileOption]: + """Let the compiled callables tolerate the gap a selected unit leaves + in them. + + A unit resolved by callable is excluded from compilation, and a compiled caller that reaches it sees that as + a graph break -- which ``fullgraph=True`` reports as an error rather than a split. Which callers those are is + not knowable from the config: Dynamo inlines across call sites, so the break can surface in any compiled + method upstream of the unit. Relaxing all of them is the only rule that is correct without tracing. + + Units resolved by op identity cost nothing here; they never leave a gap. + """ + if not self._kept_callables: + return compile_cfg + + log_rank0.info( + f"Keeping {sorted(self._selected_recompute_units)} resident excludes {sorted(self._kept_callables)} " + f"from torch.compile, so the callers that reach them are compiled with `fullgraph=False`." + ) + return { + target: TorchCompileOption(**{**option, "fullgraph": False}) + for target, option in compile_cfg.items() + if target not in self._kept_callables + } def _maybe_enable_compile(self, compile_cfg: dict[str, TorchCompileOption]): if compile_cfg: torch._dynamo.config.cache_size_limit = 256 + # Before compiling anything: install the unit wrappers, each excluded from compilation. + # Being absent from `compile_cfg` is not enough to be outside the compiled set -- Dynamo + # inlines a callee into whichever compiled caller reaches it, and the wrapper's ContextVar + # is then a hard error rather than a graph break. On a callable no compiled caller reaches, + # the exclusion costs nothing. + self._install_recompute_units() + for target, option in compile_cfg.items(): self._compile_overwrite(target, option) + def _install_recompute_units(self) -> None: + """Wrap each selected unit's callables so the policy can recognise + their ops.""" + for unit in sorted(self._selected_recompute_units): + target = self.default_recompute_cfg.get(unit) + if not isinstance(target, KeptCallables): + continue + for name in sorted(target.names): + self._recompute_unit_overwrite(unit, name) + + def _recompute_unit_overwrite(self, unit: RecomputeUnit, func_name: str) -> None: + """Install the unit wrapper on one callable, and keep it out of the + compiled set. + + Args: + unit (RecomputeUnit): The unit this callable implements. + func_name (str): Qualified name of the callable, as written in ``compile_cfg``. + """ + function = cast(FunctionType | MaybeCompile, pydoc.locate(func_name)) + if function is None: + raise AttributeError(f"Recompute config error! Cannot locate the callable: {func_name}") + + def install(fn: Any) -> Any: + return torch._dynamo.disable(in_recompute_unit(unit, fn)) + + if isinstance(function, MaybeCompile): + function.disable_compile() + function.func = install(function.func) + return + + function = cast(FunctionType, function) + if get_function_type(function) is not FunctionEnum.CLASS_LEVEL_FUNCTION: + raise ValueError( + f"Recompute config error! {func_name} must be a method or a `@maybe_compile` function to implement " + f"a recompute unit." + ) + qualname_split = function.__qualname__.split(".") + assert len(qualname_split) == 2, ( + f"XTuner Internal Error! the name of {function} should be recognized as " + f"., but got {qualname_split}" + ) + class_name, method_name = qualname_split + cls = getattr(import_module(function.__module__), class_name) + setattr(cls, method_name, install(function)) + logger.debug(f"{func_name} now implements {unit} and is excluded from compilation.") + def _mark_dynamic(self, seq_ctx: SequenceContext, dim=0): """`cu_seq_lens_q` and `cu_seq_lens_k` are dynamic shapes in each fwd/bwd pass. diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py index 7ec701056a..e6cd88742b 100644 --- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py +++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py @@ -412,6 +412,7 @@ def fully_shard( self.kept_ops, keeps_any_unit=self.keeps_any_recompute_unit, preserve_rng_state=checkpoint_preserve_rng_state, + # The layer's own forward is compiled just below, making it one opaque region. ) 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 efacd5f9e4..a4f04b1cba 100644 --- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py +++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py @@ -343,6 +343,7 @@ def fully_shard( self.kept_ops, keeps_any_unit=self.keeps_any_recompute_unit, preserve_rng_state=checkpoint_preserve_rng_state, + # The layer's own forward is compiled just below, making it one opaque region. ) 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 fbabdaa743..61ae0c48a3 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_selective_checkpointing +from xtuner.v1.model.utils import KeptOps, RecomputeTargetMap, RecomputeUnit, apply_selective_checkpointing from xtuner.v1.module import ( GatedDeltaNetConfig, LMHead, @@ -51,6 +51,15 @@ **DEFAULT_FLOAT8_CFG, } +DENSE_RECOMPUTE_CFG: RecomputeTargetMap = { + # A dense stack has no router and no expert dispatch, so attention is the only unit it can keep + # -- and it keeps it by op identity, which costs no compilation. + RecomputeUnit.SAVE_ATTN: KeptOps( + "flash_attn::_flash_attn_varlen_forward_v3", + "flash_attn::_flash_attn_varlen_forward_v2", + ), +} + class Dense(BaseModel): config: TransformerConfig @@ -164,6 +173,26 @@ def build_layers(self, config: TransformerConfig) -> nn.ModuleDict: def default_compile_cfg(self) -> dict[str, TorchCompileOption]: return DENSE_COMPILE_CFG + @property + @override + def default_recompute_cfg(self) -> RecomputeTargetMap: + """Marker intervals this architecture can keep resident, keyed by + semantic unit. + + Both units are **eager-only**: ``DenseDecoderLayer.forward`` is compiled as one fullgraph region, so the + markers that delimit them are folded away and every region falls back to being recomputed. They are declared + because eager training addresses them normally, and because the regions become effective as soon as the layer + is compiled at a finer granularity. + + Returns: + RecomputeTargetMap: Supported units mapped to the marker intervals that implement them. + """ + # Linear-attention layers carry in-place convolution state across the forward, so replaying them under a + # selective checkpoint is untested. Hybrid models declare no units until it is. + if "linear_attention" in self.config.layers_type: + return {} + return DENSE_RECOMPUTE_CFG + # NOTE: Add this overload for inferring the return type for easier type checking and using @overload # type: ignore def __call__( # type: ignore @@ -240,6 +269,7 @@ def fully_shard( self.kept_ops, keeps_any_unit=self.keeps_any_recompute_unit, preserve_rng_state=checkpoint_preserve_rng_state, + # The layer's own forward is compiled just below, making it one opaque region. ) # Linear-attention (GatedDeltaNet) layers write ``seq_ctx.seq_idx`` inside the diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 04e0f9b3fb..3613409514 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -45,7 +45,11 @@ TransformerConfig, ) from xtuner.v1.model.utils import ( + KeptCallables, + KeptOps, ModelForwardExtraLogInfo, + RecomputeTargetMap, + RecomputeUnit, apply_selective_checkpointing, module_dict_repr, ) @@ -113,6 +117,46 @@ MOE_EP_COMPILE_CFG = MOE_NON_EP_COMPILE_CFG.copy() MOE_EP_COMPILE_CFG.pop(MOE_DECODER_LAYER_FORWARD) +_MOE_GATE = "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEGate" +_DISPATCHERS = ( + "xtuner.v1.module.dispatcher.torch_all2all.TorchAll2AllDispatcher", + "xtuner.v1.module.dispatcher.deepep.DeepEPDispatcher", + "xtuner.v1.module.dispatcher.agrs.MoEAGRSDispatcher", + "xtuner.v1.module.dispatcher.base.NaiveDispatcher", +) +_DISPATCH_STAGES = ( + "dispatch_preprocess", + "dispatch", + "dispatch_postprocess", + "combine_preprocess", + "combine", + "combine_postprocess", +) + +MOE_RECOMPUTE_CFG: RecomputeTargetMap = { + # The attention kernel identifies itself, so this unit needs nothing taken out of the compiled + # set: a custom op is never fused, so it reaches the checkpoint policy compiled or not. Both + # flash-attention spellings are listed because only one is registered in a given build. + RecomputeUnit.SAVE_ATTN: KeptOps( + "flash_attn::_flash_attn_varlen_forward_v3", + "flash_attn::_flash_attn_varlen_forward_v2", + ), + # The router's ops are ordinary `addmm`/`topk` that appear all over the layer, so it has to be + # named by callable -- and the callable is the projection alone, not `MoEGate.forward`. The + # projection holds all the activation memory the router has (the logits are + # `tokens x n_routed_experts` while the router's own tensors are top-k sized), and unlike the + # router it contains no in-place write, which the policy refuses to keep. + RecomputeUnit.SAVE_MOE_GATE: KeptCallables(f"{_MOE_GATE}.project"), + # The dispatcher stages are already the finest callables there are, and none of them is + # compiled, so this unit costs no compilation either. Its stages call into the backend library, + # so ops the library performs on its own buffers -- deep_ep swaps storage with `aten.set_` -- + # fall inside the unit and are reported once; they are recomputed rather than kept, and torch + # catches the case where such a write lands on a tensor the unit did keep. + RecomputeUnit.SAVE_MOE_DISPATCH: KeptCallables( + *(f"{cls}.{stage}" for cls in _DISPATCHERS for stage in _DISPATCH_STAGES) + ), +} + class MoEModelOutputs(ModelOutputs): router_logits: dict[str, torch.Tensor] | None = None @@ -575,8 +619,9 @@ def _micro_batch_forward( d2h_stream=self.offload_stream, block_idx=layer_idx - self.config.first_k_dense_replace, group="text", - custom_check_fn=lambda x: x.data_ptr() - in [hidden_states.data_ptr() for hidden_states in hidden_states_list], + custom_check_fn=lambda x: ( + x.data_ptr() in [hidden_states.data_ptr() for hidden_states in hidden_states_list] + ), prefetch=True, reserve_pin_memory=True, ): @@ -1228,8 +1273,8 @@ 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 - # Marker intervals are declared against the decoder layer, so an MTP layer - # gets the same mechanism with nothing kept: recomputed whole, as before. + # Recompute units 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 @@ -1274,6 +1319,31 @@ def default_compile_cfg(self) -> dict[str, TorchCompileOption]: else: return MOE_NON_EP_COMPILE_CFG + @property + @override + def default_recompute_cfg(self) -> RecomputeTargetMap: + """What each supported unit resolves to for this architecture. + + Both units are fine-grained: neither names ``_pre_moe_forward``, the method that holds attention, the + layernorms and the gate together so that ``torch.compile`` can cover the ops on either side of attention. + Withdrawing that would cost most of what an MoE layer compiles. + + - ``SAVE_ATTN`` is resolved by op identity and costs no compilation at all. + - ``SAVE_MOE_GATE`` names ``MoEGate.project``, which holds the router's activation memory. Its callers are + compiled with ``fullgraph=False`` so they can split at it. + + ``SAVE_MOE_DISPATCH`` is not declared here; see the note next to :data:`MOE_RECOMPUTE_CFG`. + + Returns: + RecomputeTargetMap: Supported units mapped to the marker intervals that implement them. + """ + # Linear-attention layers carry in-place convolution state across the forward, so replaying them under a + # selective checkpoint is untested. Hybrid models declare no units until it is. + if "linear_attention" in self.config.layers_type: + return {} + + return MOE_RECOMPUTE_CFG + @property def need_update_bias(self) -> bool: router_config = self.config.router @@ -1436,13 +1506,6 @@ 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 09cd02bff7..7f35465937 100644 --- a/xtuner/v1/model/utils/__init__.py +++ b/xtuner/v1/model/utils/__init__.py @@ -1,31 +1,24 @@ -from .checkpointing import apply_gradient_checkpointing -from .misc import ModelForwardExtraLogInfo, module_dict_repr -from .selective_checkpointing import ( +from xtuner.v1.utils.selective_checkpointing import ( KeptCallables, KeptOps, - RecomputeTarget, RecomputeTargetMap, RecomputeUnit, - active_recompute_unit, - apply_selective_checkpointing, - in_recompute_unit, - recompute_unit, - resolve_kept_ops, ) +from .checkpointing import apply_gradient_checkpointing +from .misc import ModelForwardExtraLogInfo, module_dict_repr +from .selective_checkpointing import apply_selective_checkpointing, in_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", + "in_recompute_unit", + "resolve_kept_ops", ] diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py index 93950780fa..b21f27bb6b 100644 --- a/xtuner/v1/model/utils/selective_checkpointing.py +++ b/xtuner/v1/model/utils/selective_checkpointing.py @@ -1,177 +1,34 @@ -"""Region-level selective activation checkpointing (SAC): contract and engine. - -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. - -**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. +"""The selective activation checkpointing engine. + +Drives a per-op checkpoint policy for the units a user selected. The vocabulary those units are +written in lives in :mod:`xtuner.v1.utils.selective_checkpointing`; the config resolution that turns +a user's ``recompute_cfg`` into resolved targets lives with the model configs. + +A unit reaches the policy by one of two routes, and the policy answers the same question either way +-- "does this op belong to a kept unit": + +- :class:`~xtuner.v1.utils.selective_checkpointing.KeptOps` targets are recognised from the op + itself, so they need nothing installed and work identically inside and outside compiled code. +- :class:`~xtuner.v1.utils.selective_checkpointing.KeptCallables` targets are recognised from a + ``ContextVar`` the engine's wrapper sets around the callable. That only works because the model + layer has also taken those callables out of the compiled set: a ``ContextVar`` set inside compiled + code is neither written nor readable. """ -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 +from typing import Any import torch import torch.nn as nn from torch.utils.checkpoint import CheckpointPolicy, checkpoint, create_selective_checkpoint_contexts from xtuner.v1.utils import log_rank0 -from xtuner.v1.utils.enum_helper import StrEnum +from xtuner.v1.utils.selective_checkpointing import RecomputeUnit, active_recompute_unit, recompute_unit from .checkpointing import CheckpointWrapper, apply_gradient_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) +__all__ = ["apply_selective_checkpointing", "in_recompute_unit", "resolve_kept_ops"] def apply_selective_checkpointing( @@ -213,7 +70,7 @@ 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 + :class:`~xtuner.v1.utils.selective_checkpointing.KeptCallables` target names, together with the exclusion from compilation that makes the marker readable at all. Args: diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 37ec22bb12..edc8af0fe6 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -139,11 +139,21 @@ def __init__( if self.gate_bias: self.bias = nn.Parameter(torch.zeros(self.n_routed_experts)) - def forward( - self, hidden_states: torch.Tensor, rollout_routed_experts: torch.Tensor | None = None - ) -> RouterResults: + def project(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute the routing logits. + + Split from :meth:`forward` so that the projection can be addressed on its own: it holds all the activation + memory the router has -- the logits are ``tokens x n_routed_experts`` while everything the router itself + produces is top-k sized -- and unlike the router it contains no in-place write, which selective checkpointing + cannot keep. + + Args: + hidden_states (torch.Tensor): Layer input, of any leading shape. + + Returns: + torch.Tensor: Routing logits, flattened to ``(tokens, n_routed_experts)``. + """ _, _, h = hidden_states.shape - ### compute gating score hidden_states = hidden_states.view(-1, h) if isinstance(self.weight, DTensor): @@ -156,11 +166,14 @@ def forward( bias = self.bias.to_local() if isinstance(self.bias, DTensor) else self.bias if self.router_compute_dtype == "native": - logits = F.linear(hidden_states, weight, bias) - else: - bias = bias.float() if bias is not None else None - logits = F.linear(hidden_states.float(), weight.float(), bias) - return self.router(logits, rollout_routed_experts) + return F.linear(hidden_states, weight, bias) + bias = bias.float() if bias is not None else None + return F.linear(hidden_states.float(), weight.float(), bias) + + def forward( + self, hidden_states: torch.Tensor, rollout_routed_experts: torch.Tensor | None = None + ) -> RouterResults: + return self.router(self.project(hidden_states), rollout_routed_experts) # Debug for aligning with hf implementation. # logits = F.linear(hidden_states, weight, bias) @@ -306,6 +319,7 @@ def __init__( self.dispatcher = build_dispatcher( dispatcher=dispatcher, n_routed_experts=n_routed_experts, + hidden_size=hidden_size, ep_group=process_group, training_dtype="fp8" if float8_cfg is not None else "bf16", generate_dtype=generate_config.dtype if generate_config is not None else "bf16", @@ -417,11 +431,14 @@ def _forward( origin_shape = hidden_states.shape # reshape hidden_states to (batch_size * seq_len, hidden_size) + # Flattened before the marker so that the dispatch region covers the same operations here as + # it does in `_micro_batch_forward`, which reshapes outside the region too. + flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) # ProberList.before_dispatch( # self.layer_idx, hidden_states, router_results["topk_ids"], router_results["topk_weights"] # ) pre_dispatched = self.dispatcher.dispatch_preprocess( - hidden_states=hidden_states.view(-1, hidden_states.shape[-1]), + hidden_states=flat_hidden_states, topk_ids=router_results["topk_ids"], ) dispatched = self.dispatcher.dispatch( @@ -473,14 +490,15 @@ def _forward( pre_combined=pre_combined, combined=combined, ) - combined_hidden_states = post_combined["hidden_states"] - combined_hidden_states = combined_hidden_states.view(*origin_shape) + combined_hidden_states = post_combined["hidden_states"].view(*origin_shape) # debug for aligning with hf implementation. # combined_hidden_states = self._hf_expert_forward_for_debug(hidden_states, router_results, origin_shape) # ProberList.after_combine(self.layer_idx, combined_hidden_states) + # Recorded outside the branch so the region never depends on a configuration-dependent path: + # without shared experts it is simply empty rather than left open until the next marker. if self.n_shared_experts > 0: shared_experts_out = self._shared_experts_forward(hidden_states=hidden_states) else: @@ -600,6 +618,11 @@ def _micro_batch_forward( shared_experts_out_list: list[torch.Tensor | None] + # Recorded outside the branch so the region never depends on a configuration-dependent path: + # without shared experts it is simply empty rather than left open until the next marker. It + # also sits outside the loop, unlike the single-batch path's per-call region: this stage runs + # every micro-batch back to back with nothing else between them, so one region spanning the + # whole stage covers exactly the same operations as one region per micro-batch would. if self.n_shared_experts > 0: shared_experts_out_list = [] for pre_moe_forward_out in pre_moe_forward_out_list: diff --git a/xtuner/v1/module/dispatcher/__init__.py b/xtuner/v1/module/dispatcher/__init__.py index e4981392d0..4017e340d8 100644 --- a/xtuner/v1/module/dispatcher/__init__.py +++ b/xtuner/v1/module/dispatcher/__init__.py @@ -30,6 +30,7 @@ def build_dispatcher( dispatcher: Literal["deepep", "all2all", "agrs"] | None, n_routed_experts: int, + hidden_size: int, ep_group: dist.ProcessGroup | None = None, training_dtype: Literal["bf16", "fp8"] = "bf16", generate_dtype: Literal["bf16", "fp8"] = "bf16", @@ -55,6 +56,7 @@ def build_dispatcher( # TODO: remove type ignore here return DeepEPDispatcher( n_routed_experts=n_routed_experts, + hidden_size=hidden_size, process_group=ep_group, training_dtype=training_dtype, generate_dtype=generate_dtype, diff --git a/xtuner/v1/module/dispatcher/deepep.py b/xtuner/v1/module/dispatcher/deepep.py index 679253e2ba..10d5f38f20 100644 --- a/xtuner/v1/module/dispatcher/deepep.py +++ b/xtuner/v1/module/dispatcher/deepep.py @@ -13,6 +13,7 @@ combine_forward, dispatch_backward, dispatch_forward, + get_low_latency_buffer, ) from xtuner.v1.utils import copy_method_signature, get_device, get_logger @@ -257,6 +258,7 @@ def __init__( self, *, n_routed_experts: int, + hidden_size: int, process_group: torch.distributed.ProcessGroup, training_dtype: Literal["fp8", "bf16"] = "bf16", generate_dtype: Literal["fp8", "bf16"] = "bf16", @@ -273,6 +275,13 @@ def __init__( "Process group must be provided for `DeepEPDispatcher`. " "If you are training a MoE model, it means that `expert parallel` is not enabled in the config." ) + # Built here rather than on first dispatch. The buffer is a process-wide singleton whose + # constructor all-gathers device ids and IPC handles, and `all_gather_object` swaps a + # tensor's storage with `aten.set_`. Left lazy, that lands inside whichever forward happens + # to run first -- which under selective checkpointing is a checkpointed one, where the write + # hits a tensor the policy kept and torch rejects the step with "Tensor cached during + # selective activation checkpoint has been mutated". + get_low_latency_buffer(self._process_group, hidden=hidden_size, num_experts=n_routed_experts) @override def dispatch_preprocess( diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index 1d279156d9..524ba79351 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -1152,6 +1152,12 @@ def build_engine( if engine.model.compile_cfg is not None: log_rank0.info(f"The `compile_cfg` of model is {json.dumps(engine.model.compile_cfg, indent=4)}") + # Only reported when something is actually kept resident: this is the memory knob to look at + # first when a run's peak does not match the one it was tuned for. + if engine.model.keeps_any_recompute_unit: + log_rank0.info( + f"The `recompute_cfg` of model keeps these units resident: {sorted(engine.model._selected_recompute_units)}" + ) return engine def build_lr_scheduler(self, lr_cfg: LRConfig, scheduler_step: int) -> torch.optim.lr_scheduler.LRScheduler: diff --git a/xtuner/v1/utils/__init__.py b/xtuner/v1/utils/__init__.py index 915e2a0c79..5c4c85618b 100644 --- a/xtuner/v1/utils/__init__.py +++ b/xtuner/v1/utils/__init__.py @@ -22,6 +22,13 @@ ) from .pad import pad_to_max_length, pad_to_multiple_of from .profile import profile_time, profile_time_and_memory, timer, timer_logger +from .selective_checkpointing import ( + KeptCallables, + KeptOps, + RecomputeTarget, + RecomputeTargetMap, + RecomputeUnit, +) from .state import ForwardState from .type_helper import copy_method_signature, copy_signature, ray_method from .update_weights_utils import monkey_unpatch_torch_reductions @@ -68,4 +75,9 @@ "trim_memory", "group_tensors_by_device_mesh_and_placements", "cal_total_norm", + "KeptCallables", + "KeptOps", + "RecomputeTarget", + "RecomputeTargetMap", + "RecomputeUnit", ] diff --git a/xtuner/v1/utils/selective_checkpointing.py b/xtuner/v1/utils/selective_checkpointing.py new file mode 100644 index 0000000000..50f0f2e6da --- /dev/null +++ b/xtuner/v1/utils/selective_checkpointing.py @@ -0,0 +1,167 @@ +"""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 typing import TypeAlias + +from .enum_helper import StrEnum + + +__all__ = [ + "RecomputeUnit", + "KeptOps", + "KeptCallables", + "RecomputeTarget", + "RecomputeTargetMap", + "active_recompute_unit", + "recompute_unit", +] + + +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)