From 3601a682e95a9616f4c2572df926742a04b27322 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Thu, 6 Aug 2026 19:39:55 +0000 Subject: [PATCH 1/5] [Refactor] Split the selective checkpointing contract below the model layer The contract names callables that live in `xtuner/v1/module`, and the model configs that select units live in `xtuner/v1/model`. Keeping the vocabulary under either package would make the two import each other. Split rather than move: the vocabulary and the unit marker go to `xtuner/v1/utils`, below both packages, while the policy and the checkpoint wrapping stay at the model layer, where knowing `nn.Module` and `CheckpointWrapper` is legitimate. `xtuner.v1.model.utils` keeps exporting the same names, so no import site outside these files changes. The test that guards this imports a decoder layer in a clean interpreter: in process the model layer is already loaded, so a contract that reached back into `model/` would still import fine and the cycle would only appear later. --- tests/model/test_selective_checkpointing.py | 17 ++ xtuner/v1/model/utils/__init__.py | 21 +- .../v1/model/utils/selective_checkpointing.py | 187 ++---------------- xtuner/v1/utils/__init__.py | 12 ++ xtuner/v1/utils/selective_checkpointing.py | 167 ++++++++++++++++ 5 files changed, 223 insertions(+), 181 deletions(-) create mode 100644 xtuner/v1/utils/selective_checkpointing.py diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py index b1738f9366..d7502449bd 100644 --- a/tests/model/test_selective_checkpointing.py +++ b/tests/model/test_selective_checkpointing.py @@ -6,6 +6,8 @@ TestKeptCallables test_kept_callable_reproduces_full_recompute: 按 callable 留驻与全重算逐位相同。 test_unit_marker_outside_a_region_is_noop: 未被包装时不产生任何可观察行为。 +TestContractLayering + test_module_layer_imports_the_contract_without_the_model_layer: 契约模块必须能被 module/ 层单独导入。 TestUnsupportedUnits test_in_place_op_in_kept_unit_is_recomputed_not_refused: 只动自己缓冲区的 in-place 写不该被拒。 test_mutating_a_kept_tensor_is_caught_by_torch: 写到被留驻张量上时由 torch 精确拦下。 @@ -19,6 +21,8 @@ test_kept_unit_matches_full_recompute_under_compile: compile 下同上，且不触发 cached-tensor-mutated。 """ +import subprocess +import sys from collections import Counter import pytest @@ -186,6 +190,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 的 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 8891064953..ce0f5c76e9 100644 --- a/xtuner/v1/model/utils/selective_checkpointing.py +++ b/xtuner/v1/model/utils/selective_checkpointing.py @@ -1,181 +1,34 @@ -"""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. +"""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, 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 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) +__all__ = ["apply_selective_checkpointing", "in_recompute_unit", "resolve_kept_ops"] def apply_selective_checkpointing( @@ -217,8 +70,8 @@ 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. + :class:`~xtuner.v1.utils.selective_checkpointing.KeptCallables` target names, together with + the exclusion from compilation that makes the marker readable at all. Args: unit (RecomputeUnit): The unit this callable implements. 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) From 7d1f68dc79e0fe137f2f9419d29a3175bdc0d1b7 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Thu, 6 Aug 2026 19:40:37 +0000 Subject: [PATCH 2/5] [Fix] Build the DeepEP low-latency buffer when the dispatcher is constructed The buffer is a process-wide singleton created on first dispatch. Its constructor all-gathers device ids and IPC handles, and `all_gather_object` swaps a tensor's storage through `aten.set_`. Left lazy, that one-time setup lands inside whichever forward happens to run first. Under selective checkpointing that is a checkpointed forward, and the storage swap hits a tensor the policy kept resident, so torch fails the step with "Tensor cached during selective activation checkpoint has been mutated" -- on the first step only, from a call site unrelated to the model code. Building it in `__init__` puts the setup where it belongs: once, outside any checkpointed region. `hidden_size` is threaded through `build_dispatcher` because the buffer is sized from it. --- xtuner/v1/module/decoder_layer/moe_decoder_layer.py | 1 + xtuner/v1/module/dispatcher/__init__.py | 2 ++ xtuner/v1/module/dispatcher/deepep.py | 9 +++++++++ 3 files changed, 12 insertions(+) diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 37ec22bb12..baa4f203fb 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -306,6 +306,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", 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( From b628998f6dafe698ecb1a1d3592472192005b418 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Fri, 7 Aug 2026 08:03:07 +0000 Subject: [PATCH 3/5] [Feature] Add `recompute_cfg` and per-model recompute declarations Gives users the selection and models the table it resolves against. A config lists `RecomputeUnit` members; each model declares a `RecomputeTargetMap` binding the units it supports to either op names or callable names, mirroring how `compile_cfg` is declared and resolved. MoE declares three units. Attention resolves by op identity, which costs no compilation. The gate and the dispatch stages resolve to callables, which means withdrawing them from the compiled set -- a contextvar set inside compiled code is neither written nor readable -- so selecting one relaxes `fullgraph` on the entries that would otherwise inline them. `MoEGate.forward` is split so the marker encloses only the routing projection. Top-k selection normalises its weights in place, and an in-place op inside a kept unit is recomputed rather than kept, which would have made the unit keep nothing. --- ci/config/qwen3_moe_30BA3_sac.py | 82 ++++++ ci/config/sac_verify_attn.py | 39 +++ ci/config/sac_verify_none.py | 39 +++ docs/design/selective_checkpointing.md | 240 ++++++++++++++++++ tests/model/test_recompute.py | 201 ++++++++++++++- tests/model/test_selective_checkpointing.py | 14 + tools/verify_sac.sh | 57 +++++ xtuner/v1/model/base.py | 228 +++++++++++++++-- .../compose/intern_s1/modeling_vision.py | 1 + .../model/compose/qwen3_vl/modeling_vision.py | 1 + xtuner/v1/model/dense/dense.py | 32 ++- xtuner/v1/model/moe/moe.py | 85 ++++++- .../module/decoder_layer/moe_decoder_layer.py | 46 +++- xtuner/v1/train/trainer.py | 6 + 14 files changed, 1023 insertions(+), 48 deletions(-) create mode 100644 ci/config/qwen3_moe_30BA3_sac.py create mode 100644 ci/config/sac_verify_attn.py create mode 100644 ci/config/sac_verify_none.py create mode 100644 docs/design/selective_checkpointing.md create mode 100755 tools/verify_sac.sh diff --git a/ci/config/qwen3_moe_30BA3_sac.py b/ci/config/qwen3_moe_30BA3_sac.py new file mode 100644 index 0000000000..b632a87d5f --- /dev/null +++ b/ci/config/qwen3_moe_30BA3_sac.py @@ -0,0 +1,82 @@ +"""Qwen3-MoE-30BA3 with region-level selective checkpointing switched on. + +Copied from `qwen3_moe_30BA3.py`; the only substantive difference is `recompute_cfg` on the model +config. Everything below `recompute_ratio` (1.0 by default) is still wrapped in a checkpoint -- a +unit named here is kept resident *inside* those layers instead of being recomputed. + +To sweep arms, copy this file and change RECOMPUTE_CFG. Do not pass it on the command line. + + RECOMPUTE_CFG = None # recompute everything (baseline) + RECOMPUTE_CFG = True # keep every unit the model supports + RECOMPUTE_CFG = [RecomputeUnit.SAVE_ATTN] # keep exactly these + +Available units for an MoE stack (it supports all of them): + SAVE_ATTN attention call and its projections + SAVE_MOE_GATE router: gating projection, top-k, routing weights + SAVE_MOE_DISPATCH permutation/padding around the dispatch and combine all-to-alls + SAVE_MLP shared-expert MLP, plus the dense MLP of the first_k_dense_replace layers + +Needs torch >= 2.10: regions are carried as fx annotations, which older torch cannot record. With +RECOMPUTE_CFG = None this config runs on 2.9 as well. + +Run (8xH200): + export QWEN3_MOE_PATH=... ALPACA_PATH=... # see ci/scripts/CI_ENV.sh + torchrun --nproc-per-node 8 -m xtuner.v1.train.cli.sft --config ci/config/qwen3_moe_30BA3_sac.py +""" + +import os + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.datasets import FTDPTokenizeFnConfig +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config +from xtuner.v1.train import TrainerConfig +from xtuner.v1.utils.selective_checkpointing import RecomputeUnit + + +QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"] +ALPACA_PATH = os.environ["ALPACA_PATH"] + +RECOMPUTE_CFG = [RecomputeUnit.SAVE_ATTN] + +moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RECOMPUTE_CFG) +# moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RECOMPUTE_CFG) +optim_cfg = AdamWConfig(lr=6e-05) +lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6) +fsdp_cfg = FSDPConfig( + torch_compile=True, + cpu_offload=False, + ep_size=moe_cfg.ep_size, +) + +dataset_config = [ + { + "dataset": DatasetConfig(name="alpaca", anno_path=ALPACA_PATH, sample_ratio=1.0), + "tokenize_fn": FTDPTokenizeFnConfig(max_length=16386), + }, +] + +dataloader_config = DataloaderConfig(pack_max_length=16384) + +loss_cfg = CELossConfig() + + +trainer = TrainerConfig( + load_from=QWEN3_MOE_PATH, + model_cfg=moe_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + dataset_cfg=dataset_config, + dataloader_cfg=dataloader_config, + lr_cfg=lr_cfg, + loss_cfg=loss_cfg, + tokenizer_path=QWEN3_MOE_PATH, + global_batch_size=16, + # total_epoch=1, + total_step=15, + # A 30B HF checkpoint per arm fills the disk fast, and nothing here needs the weights back. + debug_skip_save=True, + work_dir="/tmp/qwen3_moe_30BA3_sac", + seed=0, +) diff --git a/ci/config/sac_verify_attn.py b/ci/config/sac_verify_attn.py new file mode 100644 index 0000000000..889f529e1b --- /dev/null +++ b/ci/config/sac_verify_attn.py @@ -0,0 +1,39 @@ +"""SAC 验证配置（attn）。配套 tools/verify_sac.sh，勿手动单独运行。 + +eager + all2all：关掉 compile 才能把 loss 差异归因到 SAC 本身而不是编译核。 +""" + +import os + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.datasets import FTDPTokenizeFnConfig +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config +from xtuner.v1.train import TrainerConfig + + +moe_cfg = Qwen3MoE30BA3Config(ep_size=4, dispatcher="all2all", recompute_cfg=['save_attn']) +fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=moe_cfg.ep_size) + +trainer = TrainerConfig( + load_from=os.environ["QWEN3_MOE_PATH"], + model_cfg=moe_cfg, + optim_cfg=AdamWConfig(lr=6e-05), + fsdp_cfg=fsdp_cfg, + dataset_cfg=[ + { + "dataset": DatasetConfig(name="alpaca", anno_path=os.environ["ALPACA_PATH"], sample_ratio=1.0), + "tokenize_fn": FTDPTokenizeFnConfig(max_length=8194), + }, + ], + dataloader_cfg=DataloaderConfig(pack_max_length=8192), + lr_cfg=LRConfig(lr_type="cosine", lr_min=1e-6), + loss_cfg=CELossConfig(), + tokenizer_path=os.environ["QWEN3_MOE_PATH"], + global_batch_size=16, + total_step=3, + debug_skip_save=True, + work_dir="/tmp/sac_verify_attn", + seed=0, +) diff --git a/ci/config/sac_verify_none.py b/ci/config/sac_verify_none.py new file mode 100644 index 0000000000..35c7204e42 --- /dev/null +++ b/ci/config/sac_verify_none.py @@ -0,0 +1,39 @@ +"""SAC 验证配置（none）。配套 tools/verify_sac.sh，勿手动单独运行。 + +eager + all2all：关掉 compile 才能把 loss 差异归因到 SAC 本身而不是编译核。 +""" + +import os + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.datasets import FTDPTokenizeFnConfig +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config +from xtuner.v1.train import TrainerConfig + + +moe_cfg = Qwen3MoE30BA3Config(ep_size=4, dispatcher="all2all") +fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=moe_cfg.ep_size) + +trainer = TrainerConfig( + load_from=os.environ["QWEN3_MOE_PATH"], + model_cfg=moe_cfg, + optim_cfg=AdamWConfig(lr=6e-05), + fsdp_cfg=fsdp_cfg, + dataset_cfg=[ + { + "dataset": DatasetConfig(name="alpaca", anno_path=os.environ["ALPACA_PATH"], sample_ratio=1.0), + "tokenize_fn": FTDPTokenizeFnConfig(max_length=8194), + }, + ], + dataloader_cfg=DataloaderConfig(pack_max_length=8192), + lr_cfg=LRConfig(lr_type="cosine", lr_min=1e-6), + loss_cfg=CELossConfig(), + tokenizer_path=os.environ["QWEN3_MOE_PATH"], + global_batch_size=16, + total_step=3, + debug_skip_save=True, + work_dir="/tmp/sac_verify_none", + seed=0, +) diff --git a/docs/design/selective_checkpointing.md b/docs/design/selective_checkpointing.md new file mode 100644 index 0000000000..536f5076f5 --- /dev/null +++ b/docs/design/selective_checkpointing.md @@ -0,0 +1,240 @@ +# Selective Checkpointing & Activation Offload + +## 背景 + +目前 XTuner 的 recompute（gradient checkpointing）和 activation offload 都只能做到 +`DecoderLayer` 级别：在 FSDP sharding 阶段，对整层 `layer` 套 `checkpoint_wrapper` +（见 `xtuner/v1/model/moe/moe.py:1039`），并由 `fsdp_config.recompute_ratio` +（`xtuner/v1/config/fsdp.py:18`）控制重算前若干层。 + +这种粒度太粗：一层里 attention / MoE / MLP 的显存和重算代价差异很大，无法只重算其中一部分。 +本设计的目标是把 recompute 和 activation offload 下沉到 **region（埋点区间）级别**：模型作者在 +forward 里用 `checkpoint_record` 埋点，每个模型在 `model_cfg` 里以 marker 区间声明「哪些段要 +重算 / offload」，再由 `recompute_ratio` 控制作用到哪几层。区间可**跨模块**，不受架构边界约束。 + +## Goals / Non-Goals + +**Goals**（按优先级排序） + +- **[最高优先] 在 torch 2.10 上用非 legacy（`use_reentrant=False`）的 checkpoint 替换现有 legacy 接口。** + 这是本设计的前置项，需落实两件事：①domino EP（`intra_layer_micro_batch > 1`）下非 reentrant 能正确 + 重算；②非 legacy 走 `saved_tensors_hooks`，`DecoderLayer` 可返回 TypedDict、`seq_ctx` 等可 dict 直传， + `_check_signature_of_forward` 随之可删、model / mtp 简化。（后续 region 级 SAC recompute 也依赖此升级， + 见「SAC：region 级 recompute」。）**先做这一步，再据其结论落地后续 schema 与实现。** +- 支持 region（埋点区间）级别的 recompute，用户通过 `recompute_cfg`（语义枚举）声明要留驻哪些子结构，粒度可跨模块。 +- `recompute_cfg` 是每个模型的 per-layer 默认方案，`recompute_ratio` 决定这套方案作用到哪些 layer。 + +**Non-Goals（本 PR 不做，但需保证不冲突）** + +- **activation offload 的实现本身**。本 PR 只做 recompute，`activation_offload_cfg` 的 schema + 与 recompute 对齐、预留位置，保证后续开发不会与本次改动冲突，具体 offload 机制（落在 + `async_save_on_cpu` 上、以及与环境变量 `XTUNER_ACTIVATION_OFFLOAD` 的 BC）留到后续设计。 +- 脱离代码埋点的全局 monkeypatch。recompute 区间只通过模型 forward 里显式的 `checkpoint_record` + 埋点界定，不做对任意自由函数的运行时替换。 + +## 第一优先：调研升级到非 legacy checkpoint wrapper + +> 这是本设计的前置调研项，先于下面的 schema / 实现完成。本章只谈 wrapper 升级本身，就两件事；SAC +> 对非 legacy 的依赖与相关约束放在「SAC：region 级 recompute」里讲。 + +**现状**：`checkpoint_wrapper`（`xtuner/v1/model/utils/checkpointing.py`）封装 legacy 的 +`torch.distributed.algorithms._checkpoint.checkpoint_wrapper` + `CheckpointImpl.REENTRANT`。 +reentrant 版是一个 `autograd.Function`，要求需要求梯度的 tensor **以位置参数直传直返**；`dict` 嵌套 +tensor 会被当成单个非 tensor 参数，wrapper 判定「没有输入需要梯度」→ 断梯度 / 报错。 +`_check_signature_of_forward`（`checkpointing.py:28`）那套严格签名检查正是为在这种情况下早失败而存在。 + +升级到 `use_reentrant=False` 要落实两件事： + +**① domino EP 正确性。** 现在用 `REENTRANT` 的主要原因是 domino EP（`trainer_cfg` 里 +`intra_layer_micro_batch > 1`）下重算要能正常进行。先复现旧接口为何需要 reentrant / 非 reentrant +会失败的现象与根因，再验证 `use_reentrant=False` 在 `intra_layer_micro_batch > 1` 下能否正确重算。 + +**② TypedDict I/O + 删签名检查。** `use_reentrant=False` 用 `saved_tensors_hooks` 实现，对嵌套输入 / +返回值天然友好。升级后 `DecoderLayer` 可以**返回 TypedDict 而不是位置约定的 tuple**、`seq_ctx` / +`position_embeddings` 等可作为 dict 直传，mtp 那类第二路输出不必再挤进 tuple 某个约定位置、由下游按 +index 解包。因此本步的**直接产出**包括：decoder layer 返回值 tuple → TypedDict，删除 +`_check_signature_of_forward`，model / mtp 层代码随之简化。机制已用 POC（玩具模型）验证：非 legacy 下 +「嵌套 dict 输入 + TypedDict 输出」的梯度与无 checkpoint **精确一致**、到达 dict 内每个 tensor；同一个 +dict 进 / 出的 layer 用 legacy REENTRANT 则直接报 `element 0 of tensors does not require grad`。落到 +真实 layer 需覆盖 `seq_ctx` / `position_embeddings` 等实际结构。 + +**验收**：两件事都到位后，domino EP 下做 grad-norm parity + 显存对比（升级前后），确认数值一致、显存 +收益符合预期。 + +## SAC：region 级 recompute + +region 级 recompute 的主体是 SAC（selective activation checkpointing）：整层套一个 checkpoint、由 +policy 逐 op 裁决存或重算。用户侧的 `recompute_cfg`（语义枚举）只是这套机制最终的**表达形式**，先讲机制、再落到配置。 + +### 机制 + +recompute 落在 `torch.utils.checkpoint.checkpoint(..., use_reentrant=False, context_fn=...)` +的 SAC 路径上，而**不是**对 submodule 逐个套 wrapper。被选中的整个 `DecoderLayer` 套**一个** +checkpoint，`context_fn` 由 torch 的 `create_selective_checkpoint_contexts(policy_fn)` 构造， +`policy_fn` 逐 op 裁决：默认 `MUST_RECOMPUTE`（全重算），被点名的区间返回 `MUST_SAVE`（该段激活留驻、 +不重算）。用户看到的是 region 粒度（marker），底层是 op 粒度的 policy 在执行。 + +```python +from torch.utils.checkpoint import checkpoint, create_selective_checkpoint_contexts, CheckpointPolicy + +def policy_fn(ctx, op, *args, **kwargs) -> CheckpointPolicy: + # marker 状态机决定当前是否处于某个 SAVE 区间 + return CheckpointPolicy.MUST_SAVE if _in_save_interval() else CheckpointPolicy.MUST_RECOMPUTE + +ctx_fn = lambda: create_selective_checkpoint_contexts(policy_fn) +out = checkpoint(layer_forward, *args, use_reentrant=False, context_fn=ctx_fn) +``` + +这条路径**必须** `use_reentrant=False`——`create_selective_checkpoint_contexts` / `context_fn` +只有非 legacy 才支持，因此它与「第一优先」的非 legacy wrapper 调研是同一件事：调研结论为真是本 schema +成立的前提。 + +**默认口径的选择**：policy 对每个 op 只回答「SAVE（存下来、不重算）还是 RECOMPUTE（丢掉、backward +重算）」。于是「默认 + 点名例外」有两种反向配法： + +- **A**：默认 `MUST_RECOMPUTE`，点名的区间设成 `MUST_SAVE` —— 默认全重算，点名少数**不重算**。 +- **B**：默认 `MUST_SAVE`，点名的区间设成 `MUST_RECOMPUTE` —— 默认全存，点名少数**重算**。 + +两种 SAC 都能做、也都能跨模块，差别只是 `policy_fn` 默认返回值那一行。**本设计只用 A**：recompute 的 +目的是省显存，常态一定是「重算绝大多数、只把少数重算代价高的段（如 flash-attn）留驻」，所以默认全重算、 +点名 SAVE 才符合直觉，`recompute_cfg` 的语义也单一。不为对称性把 B 也塞进 schema。 + +### 埋点：`checkpoint_record` + +模型作者在 forward 里埋**点位**（不是区间），标注可寻址的语义边界： + +```python +def forward(self, x): + checkpoint_record("attn.core") + ... + checkpoint_record("mlp.up") + ... +``` + +`checkpoint_record(name)` 是命令式点位，不在 SAC 会话内时是 no-op。它翻转一个 context-local 的 +「当前活跃区间」状态机，policy 据此决定当前 op 存还是重算。 + +**为什么是点位 marker 而不是 `with` 上下文**：重算区间可能**跨模块**——例如 `A.a2 + B.b1` +（A 尾部 + B 头部），二者分处 `A.forward` 与 `B.forward`，中间隔着上层的 `A()→B()` 调用。`with` +块受词法作用域限制，无法横跨两个兄弟模块的 forward；命令式 marker 依赖**运行时程序顺序 + 外部状态**， +天然可跨。 + +### 用户配置：语义枚举，不是 marker 字符串 + +用户**不直接填 marker 字符串**（那反人类，且要求用户了解模型内部埋点）。面向用户的是一组**语义枚举**， +每个枚举表示「某类子结构不重算（留驻）」，与 `recompute_ratio`（选层）配合： + +```python +class RecomputeUnit(Enum): + SAVE_ATTN = auto() # attention 段不重算、留驻 + SAVE_MOE_GATE = auto() + ... + +# model config +recompute_cfg: list[RecomputeUnit] | bool | None +``` + +枚举 → 具体 marker 区间的映射，仿照 `default_compile_cfg`，定义在每个 model 的默认 cfg 里（各模型的 +attention / MoE 结构不同，同一个 `SAVE_ATTN` 落到的 marker 区间也不同）： + +```python +@property +def default_recompute_cfg(self) -> dict[RecomputeUnit, list[tuple[str, str]]]: + return { + RecomputeUnit.SAVE_ATTN: [("attn.core", "attn.out")], + RecomputeUnit.SAVE_MOE_GATE: [("moe.gate", "moe.dispatch")], + } +``` + +三层职责因此各自稳定： + +- **用户**：config 里挑语义枚举（`SAVE_ATTN`），不碰字符串、不需了解内部结构。 +- **model 作者**：定义枚举 → marker 区间，每个模型一次、随架构而定（与 `checkpoint_record` 埋点同处一层）。 +- **内部**：把选中的枚举解析成 marker 区间，喂给 policy 的状态机（下节）。 + +> 与 `compile_cfg` 的 `bool | None | dict` 三态一致：`None` / `True` 用 model 默认方案，`False` 关闭 +> （全重算、不留驻）。是否再开放「直传区间」的高级口径（调试用）留待实现时定，默认路径只有枚举。 + +### 内部表示：marker 区间 → policy + +枚举解析出的每个 `(start, end)` 是一段 `[start, end)` 半开区间（`end` 取「下一步起点 marker」）， +区间内的 op 被标 `MUST_SAVE`。区间重叠 / 嵌套用「活跃区间集合非空即 SAVE」定义；policy 只做「命中 +start 置位、命中 end 复位」的状态机，不做配对断言。 + +这带来一个稳健性红利：marker 区间配不平（如 `end` 落在没走到的分支）只会**多留驻一段显存，绝不影响 +梯度**——SAVE 与 RECOMPUTE 两条路数值都精确正确。因此 model 作者写 `default_recompute_cfg` 时无需 +担心边界配对的正确性。 + +### 状态生命周期（实现关键） + +会话状态必须在**被 checkpoint 的 callable 顶部**重新绑定、每趟重置。因为 `use_reentrant=False` +的 backward 会**重跑整个 forward** 来重算，marker 随之重新命中；若状态挂在 module 全局而非随 callable +重置，forward 与 recompute 两趟的 SAVE 集合会错位。用 `contextvars` 在 callable 入口 set、 +`finally` reset。 + +### 待验证的 SAC 风险点 + +玩具模型 POC 已验证：梯度精确、跨模块区间成立、SAVE 段确实不重算。但下面两点玩具模型覆盖不到，需在 +真实 layer 上验证： + +- **op 序确定性**：SAC 的 save-list 靠 forward 与 recompute 两趟 **op 顺序一一对齐**。 + `intra_layer_micro_batch > 1`（domino）下 micro-batch 交织会不会改变 op 顺序，是本方案能否落地的 + 头号拦路虎。 +- **torch.compile 交互**：SAC 的 `context_fn` + `saved_tensors_hooks` 在 compile 下的行为需单独确认。 + +## 寻址语义：与 `recompute_ratio` 的两步分工 + +生效范围两步式，两者正交、不是覆盖关系： + +1. **选层**：`recompute_ratio` 决定哪些 layer 套上 SAC checkpoint（沿用现有 `_should_recompute` 语义）。 +2. **选区间**：在被选中的 layer 内部，`recompute_cfg` 选中的枚举解析成 marker 区间，抠出**留驻不重算**的段，其余全重算。 + +即 `recompute_ratio` 定义「模板铺到前多少层」，`recompute_cfg` 定义「每层内部哪些段不重算」。 + +> 旧设计的「按 module 类匹配 + 最外层去重」不再适用：SAC 下整层只有**一个** checkpoint，不存在 +> 嵌套 wrapper 的双重 checkpoint 问题，SAVE / RECOMPUTE 由 policy 逐 op 裁决，没有「谁套谁」。 + +## recompute 与 activation offload 的关系 + +两者不互斥，但作用对象正好互补：SAC 下 offload 的对象是被**留驻的那部分激活**（`MUST_SAVE` 段） +——被重算的段本就不落地、无需 offload，能省显存的只有留驻下来的 SAVE 张量。 + +本 PR **只实现 recompute**。对 offload 只有一个约束：`activation_offload_cfg` 复用 `recompute_cfg` +同一套语义枚举面向用户、同一套「枚举 → marker 区间」的 model 默认映射机制，且 SAC 的 saved-tensor +注入点要为后续 offload 的 hook 机制（预期落在 `async_save_on_cpu` 上、按留驻 tensor 建 saved-tensor +hook）留出干净入口，不写死实现细节。offload 的具体机制与环境变量 BC 留到后续设计文档。 + +## FSDP 包裹顺序 + +FSDP（`fully_shard`）永远是**最外层**：先对 `DecoderLayer` 套上 SAC checkpoint，再对该 layer +做 `fully_shard`。由于 `fully_shard` 以整层为单位、SAC checkpoint 也以整层为边界，backward 重算 +发生时该层参数已由 FSDP all-gather 就位，重算路径无需额外的 gather 时机处理。 + +## 环境 + +torch 2.10 环境路径：`/mnt/shared-storage-user/yehaochen/miniconda3/envs/py312-pt210/bin/torchrun` + +## 开发计划：分工与 stacked PR + +**硬约束**：非 legacy 升级（第一优先）是**门禁**，不能与后续并行——domino EP 下 +`use_reentrant=False` 能否正确重算的结论，决定整套 SAC schema 是否成立。因此是「门禁串行 → 通过后扇出」， +不是一次并行 N 块。 + +切成 **3 个 owning agent / 3 层 stacked PR**： + +| 层 | agent | 内容 | 主要文件面（尽量不重叠） | +| --- | --- | --- | --- | +| **A（base）** | 非 legacy 基座 | wrapper → 非 legacy + domino EP 验证；`DecoderLayer` 返回 tuple → TypedDict、`seq_ctx` dict 直传、删 `_check_signature_of_forward`、改 model / mtp 调用点。**并定义共享契约**：`checkpoint_record`（先落 no-op 桩）、`RecomputeUnit` 枚举、枚举解析出的区间类型 | `checkpointing.py`、decoder layers、model / mtp、sharding 路径 | +| **B（stack 在 A 上）** | SAC 引擎 | `checkpoint_record` 背后的 `contextvars` session + `policy_fn` + `create_selective_checkpoint_contexts` 接线；FSDP 路径换成整层 SAC；op 序 / `torch.compile` 在真实 layer 上验证 | 新 SAC 引擎模块 + 包裹调用点 | +| **C（stack 在 A 上，与 B 并行）** | 配置 + 逐模型 | `recompute_cfg` 三态解析（`bool \| None \| list[RecomputeUnit]`）+ 与 `recompute_ratio` 打通；各模型 `default_recompute_cfg`；在模型 forward 里埋 `checkpoint_record` | config 模块 + 各 model 文件 | + +**为什么合并原「非 legacy 升级」与「TypedDict 重构」为一个 A**：两者都要改 `checkpointing.py`、decoder +layer、sharding 路径，拆两个 agent 会在同一批文件上冲突；同一 owner 一次到位。 + +**为什么 B、C 能真并行**：B 只碰引擎与包裹点、不碰模型 forward；C 只碰 config 与模型文件、不碰引擎内部。 +C 埋的 `checkpoint_record` 在 B 合入前就是 no-op，无害；两者只在 A 的契约上相遇。 + +**唯一 seam**：FSDP 包裹处「把解析出的区间应用到选中层」这一行归 **B** 拥有（它是机制）；C 只负责把 +config 解析成区间、通过契约函数交给 B，不直接改包裹点，避免 B / C 同时改 sharding 路径。 + +**stack 顺序**：`A → B`、`A → C`，B / C 互为 sibling；A 合入后各自 rebase 到 A，平行推进。 diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py index 1d323cd926..f0140bd529 100644 --- a/tests/model/test_recompute.py +++ b/tests/model/test_recompute.py @@ -5,6 +5,7 @@ test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。 test_checkpointing_keeps_the_module_itself: 换类而非套壳，isinstance/属性/容器协议原生可用。 test_module_without_a_protocol_does_not_gain_one: 被包裹模块没有的协议不会凭空出现。 +TestRecomputeCfgResolution test_unset_cfg_keeps_full_recompute: `None` 不改变显存行为，解析为不留驻。 test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。 test_explicit_units_select_only_themselves: 显式 list 只选中对应 unit。 @@ -13,8 +14,10 @@ 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 撤出编译占比最大的方法。 @@ -23,13 +26,26 @@ """ +import pydoc import pytest import torch from torch import nn from torch.autograd.graph import saved_tensors_hooks -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): @@ -167,3 +183,186 @@ def __init__(self) -> None: make_call(wrapped, x)["out"].square().sum().backward() assert x.data_ptr() in packed + + +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] + + +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 d7502449bd..fdec0e0a18 100644 --- a/tests/model/test_selective_checkpointing.py +++ b/tests/model/test_selective_checkpointing.py @@ -31,6 +31,7 @@ from torch.utils._python_dispatch import TorchDispatchMode from xtuner.v1.model.utils import ( + KeptOps, RecomputeUnit, apply_selective_checkpointing, in_recompute_unit, @@ -219,6 +220,19 @@ 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, + } + + class TestRecomputeIsObservable: """按 reviewer 的三个目标直接观察：精度对齐、重算生效、重算+SAC 生效。""" diff --git a/tools/verify_sac.sh b/tools/verify_sac.sh new file mode 100755 index 0000000000..6a7de75f65 --- /dev/null +++ b/tools/verify_sac.sh @@ -0,0 +1,57 @@ +#!/bin/bash +# 验证 recompute_cfg：跑两次同一份训练，一次不选 unit、一次选 save_attn，比对结果。 +# +# 判据（两条都必须成立）： +# 1. step-1 loss 逐位相同 —— checkpoint 只改 backward，前向不能被动到； +# 2. 峰值显存显著上升 —— unit 真的把激活留驻了，而不是配置被静默忽略。 +# +# 用 eager（torch_compile=False）跑：编译会改变 kernel，从而改变 loss 的低位， +# 那样就无法把差异归因到 SAC 本身。 +set -u +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +OUT="${SAC_VERIFY_OUT:-/tmp/sac_verify}" +PY="${PYTHON:-/mnt/shared-storage-user/yehaochen/miniconda3/envs/py312-pt210/bin/torchrun}" +export QWEN3_MOE_PATH="${QWEN3_MOE_PATH:-/mnt/shared-storage-user/llmrazor-share/model/Qwen3-30B-A3B}" +export ALPACA_PATH="${ALPACA_PATH:-/mnt/shared-storage-user/llmrazor-share/data/alpaca}" +export EP_DISABLE_GIN=1 EP_REUSE_NCCL_COMM=0 +mkdir -p "$OUT" + +for unit in none attn; do + echo "==> 跑 recompute_cfg=$unit ..." + (cd /tmp && PYTHONPATH="$REPO" timeout 3600 "$PY" --master-port "${MASTER_PORT:-21971}" \ + --nproc-per-node "${NPROC:-8}" -m xtuner.v1.train.cli.sft \ + --config "$REPO/ci/config/sac_verify_$unit.py") > "$OUT/$unit.log" 2>&1 + [ $? -eq 0 ] || { echo "失败，日志见 $OUT/$unit.log"; grep -m3 -iE "error|Traceback" "$OUT/$unit.log"; exit 1; } +done + +"${PYTHON_BIN:-python}" - "$OUT" <<'PY' +import re, sys, pathlib + +out = pathlib.Path(sys.argv[1]) +res = {} +for unit in ("none", "attn"): + log = out / f"{unit}.log" + text = log.read_text(errors="ignore") + # 只看 rank 0：各 rank 峰值本就不同，跨 rank 比较得到的是 rank 差异而非配置差异。 + rows = re.findall(r"RANK 0\].*?Step (\d+)/\d+.*?reduced_llm_loss: ([0-9.]+).*?max_memory: ([0-9.]+) GB", text) + by_step = {int(s): (loss, float(mem)) for s, loss, mem in rows} + if 1 not in by_step or len(by_step) < 2: + print("没读到 %s 至少两步的 rank-0 指标，日志：%s" % (unit, log)) + sys.exit(1) + # loss 取第一步（纯前向）；显存取最后一步：第一步的峰值被权重加载/优化器状态初始化这类 + # 一次性分配盖住，激活的稳态峰值要到第二步才显出来。 + res[unit] = (by_step[1][0], by_step[max(by_step)][1]) + +l0, m0 = res["none"] +l1, m1 = res["attn"] +print("") +print(" 无 unit step1_loss=%s 稳态峰值=%.2f GB" % (l0, m0)) +print(" save_attn step1_loss=%s 稳态峰值=%.2f GB (+%.2f GB)" % (l1, m1, m1 - m0)) +print("") + +ok_loss = l0 == l1 +ok_mem = (m1 - m0) > 1.0 +print(" [%s] 前向未受影响：step-1 loss 逐位相同" % ("PASS" if ok_loss else "FAIL")) +print(" [%s] unit 确实生效：稳态峰值显存上升 > 1 GB" % ("PASS" if ok_mem else "FAIL")) +sys.exit(0 if (ok_loss and ok_mem) else 1) +PY diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 697e5f9933..8db1bb21aa 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,22 @@ def compile_cfg(self) -> dict[str, TorchCompileOption]: return _compile_cfg + @property + def default_recompute_cfg(self) -> RecomputeTargetMap: + """What each recompute unit this architecture supports resolves to. + + This is the model author's vocabulary: it declares which :class:`RecomputeUnit` s the architecture supports and + what each one is made of, not which of them are worth enabling. An architecture that supports none 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 +1064,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): @@ -2571,14 +2639,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: @@ -2586,27 +2654,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 b0407e38b4..2be3c7cad7 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: # Compile the class function, then restore descriptor binding on this instance. diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py index f2758b0d5e..a115a88689 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: # Compile the class function, then restore descriptor binding on this instance. diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index 274aebd00b..9eec5da5c9 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 c90033870e..74d77636ae 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, ) @@ -114,6 +118,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 @@ -585,8 +629,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, ): @@ -1243,8 +1288,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 @@ -1294,6 +1339,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 @@ -1456,13 +1526,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/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index baa4f203fb..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) @@ -418,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( @@ -474,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: @@ -601,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/train/trainer.py b/xtuner/v1/train/trainer.py index bd780a7deb..68a26d6a6d 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -1177,6 +1177,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: From cbcc9bac653033366eba9dd0d755f7480af02321 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Thu, 6 Aug 2026 10:49:19 +0000 Subject: [PATCH 4/5] [Docs] Record the routes measured for selective checkpointing under compile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Why a unit needs the compiler's cooperation at all, the two resolutions a unit can have and what each costs, and the four routes measured and not taken -- each with the numbers that decided it, on Qwen3-MoE-30BA3 at `ep_size=4`. The load-bearing results: withdrawing the callable that encloses attention costs −9.9% for a unit that op identity delivers for free, which is why the two resolutions exist; and `fx_traceback.annotate` survives compilation perfectly and still changes peak memory by zero bytes, because the outer checkpoint discards what the partitioner decided. Removing that checkpoint is what would make the tag route work, and it costs 61.8 GiB. --- .../selective_checkpointing_and_compile.md | 266 ++++++++++++++++++ 1 file changed, 266 insertions(+) create mode 100644 docs/design/selective_checkpointing_and_compile.md 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. From e161d2ee267098f50454906646cfc75a8ccc3892 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 11 Aug 2026 13:50:24 +0000 Subject: [PATCH 5/5] [Refactor] Name the save units for what they save, and gather recompute settings Three review points, all about the same confusion: the config said "recompute" while every field in it described what is *not* recomputed. `RecomputeUnit` becomes `SaveUnit`, and its members drop the redundant prefix: `SAVE_ATTN` -> `ATTN`, serialized as `"attn"`. Selecting a member means "save this"; everything unselected is recomputed, which is the default and free. `recompute_cfg` becomes a `RecomputeConfig` holding all of it: recompute_cfg=RecomputeConfig(ratio=0.25, save=["attn"]) `ratio` decides which layers are recomputed, `save` decides what stays resident inside them. They used to live in two different configs -- the ratios on `FSDPConfig`, the unit selection on the model config -- and neither means much without the other. `fsdp_cfg.recompute_ratio` and `vision_recompute_ratio` are kept as deprecated aliases: their default becomes `None` so that "not set" is distinguishable from "set to 1.0", a set value is migrated with a warning, and setting both the old and the new field raises rather than silently picking one. The exit condition is written on the fields. In-repo configs (`autotest/config/*_recompute.py`) are moved to the new field, so CI exercises the new path rather than the shim. Also adds the missing docstring on `_resolve_recompute_cfg`. Landed as one commit rather than folded into the commits it corrects: the rename spans the contract, the engine and the declarations, so splitting it would rewrite three commits that have already been reviewed -- and the `RecomputeConfig` migration is only reviewable as one piece. --- autotest/config/npu_qwen3_recompute.py | 4 +- autotest/config/qwen3_5_recompute.py | 4 +- autotest/config/qwen3_recompute.py | 4 +- .../selective_checkpointing_and_compile.md | 20 ++--- tests/model/test_recompute.py | 75 +++++++++++----- tests/model/test_selective_checkpointing.py | 16 ++-- xtuner/v1/config/fsdp.py | 11 ++- xtuner/v1/model/base.py | 86 +++++++++++++------ .../compose/intern_s1/modeling_vision.py | 2 +- .../model/compose/qwen3_vl/modeling_vision.py | 2 +- xtuner/v1/model/dense/dense.py | 6 +- xtuner/v1/model/moe/moe.py | 12 ++- xtuner/v1/model/utils/__init__.py | 6 +- .../v1/model/utils/selective_checkpointing.py | 6 +- xtuner/v1/utils/__init__.py | 6 +- xtuner/v1/utils/selective_checkpointing.py | 83 +++++++++++++----- 16 files changed, 229 insertions(+), 114 deletions(-) diff --git a/autotest/config/npu_qwen3_recompute.py b/autotest/config/npu_qwen3_recompute.py index 197ea1b0ac..3f8211afc4 100644 --- a/autotest/config/npu_qwen3_recompute.py +++ b/autotest/config/npu_qwen3_recompute.py @@ -10,20 +10,20 @@ from xtuner.v1.loss.ce_loss import CELossConfig from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config from xtuner.v1.train import TrainerConfig +from xtuner.v1.utils import RecomputeConfig QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"] ALPACA_PATH = os.environ["ALPACA_PATH"] -moe_cfg = Qwen3MoE30BA3Config(compile_cfg=False) +moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RecomputeConfig(ratio=0.25), compile_cfg=False) optim_cfg = AdamWConfig(lr=6e-05) lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6) fsdp_cfg = FSDPConfig( cpu_offload=False, ep_size=moe_cfg.ep_size, tp_size=4, - recompute_ratio=0.25, ) dataset_config = [ diff --git a/autotest/config/qwen3_5_recompute.py b/autotest/config/qwen3_5_recompute.py index ada5dfd2e3..2338392275 100644 --- a/autotest/config/qwen3_5_recompute.py +++ b/autotest/config/qwen3_5_recompute.py @@ -10,18 +10,18 @@ from xtuner.v1.loss.ce_loss import CELossConfig from xtuner.v1.model import Qwen3_5_VLMoE35BA3Config from xtuner.v1.train import TrainerConfig +from xtuner.v1.utils import RecomputeConfig QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"] ALPACA_PATH = os.environ["ALPACA_PATH"] -moe_cfg = Qwen3_5_VLMoE35BA3Config(compile_cfg=False) +moe_cfg = Qwen3_5_VLMoE35BA3Config(recompute_cfg=RecomputeConfig(ratio=0.25), compile_cfg=False) optim_cfg = AdamWConfig(lr=6e-05) lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6) fsdp_cfg = FSDPConfig( cpu_offload=False, # qwen3.5 don't support True - recompute_ratio=0.25, ) dataset_config = [ diff --git a/autotest/config/qwen3_recompute.py b/autotest/config/qwen3_recompute.py index 5d7ba68a2c..467bcffd35 100644 --- a/autotest/config/qwen3_recompute.py +++ b/autotest/config/qwen3_recompute.py @@ -10,19 +10,19 @@ from xtuner.v1.loss.ce_loss import CELossConfig from xtuner.v1.model.moe.qwen3 import Qwen3MoE30BA3Config from xtuner.v1.train import TrainerConfig +from xtuner.v1.utils import RecomputeConfig QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"] ALPACA_PATH = os.environ["ALPACA_PATH"] -moe_cfg = Qwen3MoE30BA3Config(compile_cfg=True) +moe_cfg = Qwen3MoE30BA3Config(recompute_cfg=RecomputeConfig(ratio=0.25), compile_cfg=True) optim_cfg = AdamWConfig(lr=6e-05) lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6) fsdp_cfg = FSDPConfig( cpu_offload=True, ep_size=moe_cfg.ep_size, - recompute_ratio=0.25, ) dataset_config = [ diff --git a/docs/design/selective_checkpointing_and_compile.md b/docs/design/selective_checkpointing_and_compile.md index 7a86d79c5c..5e74d0ae67 100644 --- a/docs/design/selective_checkpointing_and_compile.md +++ b/docs/design/selective_checkpointing_and_compile.md @@ -48,7 +48,7 @@ Two mechanics of the withdrawal are easy to get wrong: (`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: +Measured, 8k, `attn`, when it was still resolved by withdrawing the callable that encloses it: | | tgs | peak | compiled graphs / captured calls | kept | |---|---|---|---|---| @@ -71,22 +71,22 @@ 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 | — | — | — | +| `attn` | 10165.6 / 10204.6 / 10071.2 | 10147.2 | 90.4-91.2 GB | 114.4-115.3 GB | +| `moe_gate` | 10095.2 / 10179.5 / 10144.4 | 10139.7 | 93.7-94.5 GB | 117.7-118.7 GB | +| `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, +**Memory is the reproducible effect**: `attn` costs +6.0 GB allocated, `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 +`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 @@ -96,7 +96,7 @@ numerically for one step (see below) and declared for smaller models, not as a d 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. +including the `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: @@ -105,8 +105,8 @@ settings: |---|---|---|---| | 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 | +| `attn` | 2.46262765 | 24.39325905 | 3.1e-5 | +| `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. @@ -177,7 +177,7 @@ Instead of withdrawing the callable, let the marker execute during tracing and a 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 +Measured, 8k, `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". diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py index f0140bd529..07160861ad 100644 --- a/tests/model/test_recompute.py +++ b/tests/model/test_recompute.py @@ -9,11 +9,15 @@ 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_string_units_are_accepted: 配置文件里的字符串能解析成 SaveUnit。 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 序列化成可读字符串并能读回。 +TestRecomputeRatioMigration + test_old_field_still_takes_effect: 旧的 fsdp_cfg.recompute_ratio 仍然生效。 + test_unset_old_field_leaves_the_new_one_alone: 没设旧字段时不覆盖新位置的值。 + test_setting_both_is_an_error: 两处都设时报错而不是静默择一。 TestDeclaredTargets test_declared_targets_resolve: 声明表里的 op 名与 callable 名都能解析到真实对象。 test_no_unit_names_the_method_that_holds_most_compilation: 没有 unit 点名承载最多编译的那个方法。 @@ -33,6 +37,7 @@ from torch import nn from torch.autograd.graph import saved_tensors_hooks +from xtuner.v1.config import FSDPConfig 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 @@ -40,7 +45,8 @@ from xtuner.v1.model.utils import ( KeptCallables, KeptOps, - RecomputeUnit, + RecomputeConfig, + SaveUnit, apply_gradient_checkpointing, resolve_kept_ops, ) @@ -242,34 +248,34 @@ 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() + assert _resolve_units(recompute_cfg=RecomputeConfig(save=None)) == set() def test_true_selects_every_supported_unit(self): - assert _resolve_units(recompute_cfg=True) == set(MOE_RECOMPUTE_CFG) + assert _resolve_units(recompute_cfg=RecomputeConfig(save=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} + assert _resolve_units(recompute_cfg=RecomputeConfig(save=[SaveUnit.MOE_GATE])) == {SaveUnit.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} + assert _resolve_units(recompute_cfg=RecomputeConfig(save=["attn"])) == {SaveUnit.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])) + _ProbeModel(_ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=RecomputeConfig(save=[SaveUnit.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) + config = _ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=RecomputeConfig(save=False)) model = _ProbeModel(config) assert model._selected_recompute_units == set() - assert config.text_config.recompute_cfg is False + assert config.text_config.recompute_cfg.save is False assert config.text_config.compile_cfg is None def test_disable_reaches_every_sub_model_of_a_real_compose_config(self): @@ -277,23 +283,23 @@ def test_disable_reaches_every_sub_model_of_a_real_compose_config(self): # 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) + config = Qwen3VLMoE30BA3Config(recompute_cfg=RecomputeConfig(save=False)) - _disable_nested_switch(config, "recompute_cfg") + _disable_nested_switch(config, "recompute_cfg", subfield="save") for sub_config in (config.vision_config, config.projector_config, config.text_config): - assert sub_config.recompute_cfg is False + assert sub_config.recompute_cfg.save 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]) + config = _build_tiny_moe_config(recompute_cfg=RecomputeConfig(save=[SaveUnit.ATTN, SaveUnit.MOE_GATE])) dumped = config.model_dump(mode="json")["recompute_cfg"] - assert dumped == ["save_attn", "save_moe_gate"] + assert dumped["save"] == ["attn", "moe_gate"] - restored = _build_tiny_moe_config(recompute_cfg=dumped) - assert restored.recompute_cfg == [RecomputeUnit.SAVE_ATTN, RecomputeUnit.SAVE_MOE_GATE] + restored = _build_tiny_moe_config(recompute_cfg=RecomputeConfig(**dumped)) + assert restored.recompute_cfg.save == [SaveUnit.ATTN, SaveUnit.MOE_GATE] class TestDeclaredTargets: @@ -341,28 +347,55 @@ def _compile_cfg(**overrides) -> dict[str, TorchCompileOption]: 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() + assert self._compile_cfg(recompute_cfg=RecomputeConfig(save=[SaveUnit.ATTN])) == self._compile_cfg() def test_a_callable_unit_keeps_its_callers_compiled(self): - unit = RecomputeUnit.SAVE_MOE_GATE + unit = SaveUnit.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]) + relaxed = self._compile_cfg(recompute_cfg=RecomputeConfig(save=[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}" + assert _PRE_MOE_FORWARD in self._compile_cfg(recompute_cfg=RecomputeConfig(save=[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]) + config = _build_tiny_moe_config(recompute_cfg=RecomputeConfig(save=[SaveUnit.ATTN])) with torch.device("meta"): model = MoE(config=config) assert model.kept_ops assert model.keeps_any_recompute_unit + + +class TestRecomputeRatioMigration: + """`fsdp_cfg.recompute_ratio` 已迁到 `recompute_cfg.ratio`，旧字段留作过渡。""" + + def _model(self, fsdp_kwargs, **config_kwargs): + model = _ProbeModel(_ProbeConfig(text_config=_NestedProbeConfig(), **config_kwargs)) + model._migrate_recompute_ratio(FSDPConfig(**fsdp_kwargs)) + return model + + def test_old_field_still_takes_effect(self): + # 过渡期内旧配置不能突然失效，值要被搬到新位置。 + model = self._model({"recompute_ratio": 0.25}) + + assert model.config.recompute_cfg.ratio == 0.25 + + def test_unset_old_field_leaves_the_new_one_alone(self): + # 旧字段的哨兵是 None 而不是 1.0，否则「没设过」会被当成「显式设成 1.0」， + # 把用户在新位置写的值覆盖掉。 + model = self._model({}, recompute_cfg=RecomputeConfig(ratio=0.5)) + + assert model.config.recompute_cfg.ratio == 0.5 + + def test_setting_both_is_an_error(self): + # 同一个设置有两个写法时，静默挑一个是最糟的处理。 + with pytest.raises(ValueError, match="both set"): + self._model({"recompute_ratio": 0.25}, recompute_cfg=RecomputeConfig(ratio=0.5)) diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py index fdec0e0a18..2ac6be43da 100644 --- a/tests/model/test_selective_checkpointing.py +++ b/tests/model/test_selective_checkpointing.py @@ -32,7 +32,7 @@ from xtuner.v1.model.utils import ( KeptOps, - RecomputeUnit, + SaveUnit, apply_selective_checkpointing, in_recompute_unit, resolve_kept_ops, @@ -57,7 +57,7 @@ class _UnitBlock(_Block): def __init__(self) -> None: super().__init__() - self.second_stage = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._second_stage) + self.second_stage = in_recompute_unit(SaveUnit.ATTN, self._second_stage) def _second_stage(self, hidden: torch.Tensor) -> torch.Tensor: return torch.tanh(self.second(hidden)) @@ -72,7 +72,7 @@ class _InPlaceUnitBlock(nn.Module): def __init__(self) -> None: super().__init__() self.linear = nn.Linear(4, 4) - self.unit = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._unit) + self.unit = in_recompute_unit(SaveUnit.ATTN, self._unit) def _unit(self, x: torch.Tensor) -> torch.Tensor: hidden = self.linear(x) @@ -90,7 +90,7 @@ class _MutatesKeptTensorBlock(nn.Module): def __init__(self) -> None: super().__init__() self.linear = nn.Linear(4, 4) - self.unit = in_recompute_unit(RecomputeUnit.SAVE_ATTN, self._unit) + self.unit = in_recompute_unit(SaveUnit.ATTN, self._unit) def _unit(self, x: torch.Tensor) -> torch.Tensor: kept = torch.tanh(self.linear(x)) @@ -225,11 +225,11 @@ 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 isinstance(MOE_RECOMPUTE_CFG[SaveUnit.ATTN], KeptOps) assert set(MOE_RECOMPUTE_CFG) == { - RecomputeUnit.SAVE_ATTN, - RecomputeUnit.SAVE_MOE_GATE, - RecomputeUnit.SAVE_MOE_DISPATCH, + SaveUnit.ATTN, + SaveUnit.MOE_GATE, + SaveUnit.MOE_DISPATCH, } diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py index 7335d6d7a5..0d0b5f070b 100644 --- a/xtuner/v1/config/fsdp.py +++ b/xtuner/v1/config/fsdp.py @@ -15,8 +15,15 @@ class FSDPConfig(BaseModel): tp_size: Annotated[int, Parameter(help="Tensor parallel size")] = 1 ep_size: Annotated[int, Parameter(help="Expert parallel size")] = 1 reshard_after_forward: Annotated[bool, Parameter(help="Reshard model parameters after forward pass")] = True - recompute_ratio: Annotated[float, Parameter(help="Gradient checkpointing ratio for memory optimization")] = 1.0 - vision_recompute_ratio: Annotated[float, Parameter(help="Recompute ratio for vision modules")] = 1.0 + # Deprecated: moved to `model_cfg.recompute_cfg.ratio` / `.vision_ratio`, so that "which layers + # are recomputed" and "what stays saved inside them" live together -- neither means much without + # the other. `None` marks "not set by the user", which is what lets the migration tell a real + # request from a default. Remove both fields, and `BaseModel._migrate_recompute_ratio`, once + # downstream configs have moved. + recompute_ratio: Annotated[float | None, Parameter(help="Deprecated, use `model_cfg.recompute_cfg.ratio`")] = None + vision_recompute_ratio: Annotated[ + float | None, Parameter(help="Deprecated, use `model_cfg.recompute_cfg.vision_ratio`") + ] = None checkpoint_preserve_rng_state: Annotated[bool, Parameter(help="Preserve RNG state during checkpointing")] = True # Training-time FSDP CPU offload is version-sensitive for XTuner model configs # that keep selected fp32 trainable parameters outside FSDP via diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 8db1bb21aa..655f140af9 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -65,8 +65,9 @@ KeptCallables, KeptOps, ModelForwardExtraLogInfo, + RecomputeConfig, RecomputeTargetMap, - RecomputeUnit, + SaveUnit, in_recompute_unit, resolve_kept_ops, ) @@ -150,21 +151,14 @@ class XTunerBaseModelConfig(PydanticBaseModel): "`dict[str, TorchCompileOption]`: Customize the compile option", ), ] = None - # `activation_offload_cfg` belongs here, next to `recompute_cfg` and sharing its `RecomputeUnit` + # `activation_offload_cfg` belongs here, next to `recompute_cfg` and sharing its `SaveUnit` # 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 + RecomputeConfig, + Parameter(group="model", help="Which layers are recomputed, and what stays saved inside them"), + ] = RecomputeConfig() 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() @@ -560,7 +554,7 @@ def _save_file( save_file(tensors, filename, metadata=metadata) -def _disable_nested_switch(obj: Any, field_name: str) -> None: +def _disable_nested_switch(obj: Any, field_name: str, *, subfield: str | None = None) -> None: """Turn a tri-state feature switch off on every config reachable from ``obj``. @@ -573,21 +567,27 @@ def _disable_nested_switch(obj: Any, field_name: str) -> None: Args: obj (Any): Config, container, or leaf value to walk. - field_name (str): Name of the switch field to set to ``False``. + field_name (str): Name of the field holding the switch. + subfield (str | None): When the switch is one field *inside* ``field_name`` -- as + ``recompute_cfg.save`` is inside ``RecomputeConfig`` -- the name of that inner field. + Defaults to None, i.e. ``field_name`` is itself the switch. """ if isinstance(obj, PydanticBaseModel): if not hasattr(obj, field_name): return - setattr(obj, field_name, False) + if subfield is None: + setattr(obj, field_name, False) + else: + setattr(getattr(obj, field_name), subfield, False) for nested_field in type(obj).model_fields: - _disable_nested_switch(getattr(obj, nested_field), field_name) + _disable_nested_switch(getattr(obj, nested_field), field_name, subfield=subfield) elif isinstance(obj, Mapping): for value in obj.values(): - _disable_nested_switch(value, field_name) + _disable_nested_switch(value, field_name, subfield=subfield) # 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) + _disable_nested_switch(value, field_name, subfield=subfield) class BaseModel(nn.Module): @@ -610,7 +610,7 @@ def __init__(self, config: XTunerBaseModelConfig): # 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._selected_recompute_units: set[SaveUnit] = 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 @@ -648,6 +648,7 @@ def fully_shard( ) -> Self: """Fully shard the model parameters.""" self.fsdp_config = fsdp_config + self._migrate_recompute_ratio(fsdp_config) self.fsdp_mesh = self._init_world_mesh() self._world_mesh = self.fsdp_mesh @@ -1044,7 +1045,7 @@ def compile_cfg(self) -> dict[str, TorchCompileOption]: def default_recompute_cfg(self) -> RecomputeTargetMap: """What each recompute unit this architecture supports resolves to. - This is the model author's vocabulary: it declares which :class:`RecomputeUnit` s the architecture supports and + This is the model author's vocabulary: it declares which :class:`SaveUnit` s the architecture supports and what each one is made of, not which of them are worth enabling. An architecture that supports none declares nothing. @@ -2656,13 +2657,44 @@ def _resolve_compile_cfg( return self._without_compiled_selected_regions(compile_cfg) + def _migrate_recompute_ratio(self, fsdp_config: FSDPConfig) -> None: + # Compatibility shim for the two ratios that used to live on `FSDPConfig`. Remove together + # with those fields once downstream configs have moved to `recompute_cfg`. + moves = ( + ("recompute_ratio", "ratio", "recompute_cfg.ratio"), + ("vision_recompute_ratio", "vision_ratio", "recompute_cfg.vision_ratio"), + ) + for old_name, new_name, new_path in moves: + old_value = getattr(fsdp_config, old_name, None) + if old_value is None: + continue + if new_name in self.config.recompute_cfg.model_fields_set: + raise ValueError( + f"`fsdp_cfg.{old_name}` and `model_cfg.{new_path}` are both set. They are the same setting; " + f"keep only `model_cfg.{new_path}`." + ) + log_rank0.warning( + f"`fsdp_cfg.{old_name}` is deprecated and will be removed; use `model_cfg.{new_path}` instead." + ) + setattr(self.config.recompute_cfg, new_name, old_value) + def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> tuple[frozenset, set[str]]: - selected = config.recompute_cfg + """Turn the user's ``recompute_cfg.save`` selection into what the + policy needs. + + Args: + config (XTunerBaseModelConfig): The model config carrying the selection. + + Returns: + tuple[frozenset, set[str]]: The op overloads to keep by identity, and the qualified + names of the callables to keep by marker. + """ + selected = config.recompute_cfg.save # `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") + _disable_nested_switch(self.config, "recompute_cfg", subfield="save") return frozenset(), set() # `None` means "keep the memory profile of plain full recompute", not "use the model default" as it does @@ -2680,7 +2712,7 @@ def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> tuple[frozens # 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 " + f"`recompute_cfg.save=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 @@ -2691,9 +2723,9 @@ def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> tuple[frozens 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"`recompute_cfg.save` 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." + "of its own -- set `recompute_cfg.save` on the sub-model config that owns the layers instead." ) target = supported[unit] if isinstance(target, KeptOps): @@ -2753,12 +2785,12 @@ def _install_recompute_units(self) -> None: for name in sorted(target.names): self._recompute_unit_overwrite(unit, name) - def _recompute_unit_overwrite(self, unit: RecomputeUnit, func_name: str) -> None: + def _recompute_unit_overwrite(self, unit: SaveUnit, 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. + unit (SaveUnit): 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)) diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py index 2be3c7cad7..0d4a0e486d 100644 --- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py +++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py @@ -396,7 +396,7 @@ def fully_shard( for param in self.parameters(): param.requires_grad = False - recompute_ratio = fsdp_config.vision_recompute_ratio + recompute_ratio = self.config.recompute_cfg.vision_ratio num_recompute_layers = int(len(self.encoder.layer) * recompute_ratio) generator = torch.Generator() diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py index a115a88689..85e8a5d49c 100644 --- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py +++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py @@ -333,7 +333,7 @@ def fully_shard( param.requires_grad = False checkpoint_preserve_rng_state = fsdp_config.checkpoint_preserve_rng_state - num_recompute_layers = int(len(self.blocks) * fsdp_config.vision_recompute_ratio) + num_recompute_layers = int(len(self.blocks) * self.config.recompute_cfg.vision_ratio) for layer_idx in tqdm(list(range(len(self.blocks))), desc="[Vision Fully Shard]"): layer = self.blocks[layer_idx] diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index 9eec5da5c9..bf4457d852 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 KeptOps, RecomputeTargetMap, RecomputeUnit, apply_selective_checkpointing +from xtuner.v1.model.utils import KeptOps, RecomputeTargetMap, SaveUnit, apply_selective_checkpointing from xtuner.v1.module import ( GatedDeltaNetConfig, LMHead, @@ -54,7 +54,7 @@ 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( + SaveUnit.ATTN: KeptOps( "flash_attn::_flash_attn_varlen_forward_v3", "flash_attn::_flash_attn_varlen_forward_v2", ), @@ -254,7 +254,7 @@ def fully_shard( lm_head_mp_policy = MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32) else: lm_head_mp_policy = mp_policy - num_recompute_layers = int(self.config.num_hidden_layers * self.fsdp_config.recompute_ratio) + num_recompute_layers = int(self.config.num_hidden_layers * self.config.recompute_cfg.ratio) generator = torch.Generator() generator.manual_seed(dist.get_rank()) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 74d77636ae..e11937b4d7 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -49,7 +49,7 @@ KeptOps, ModelForwardExtraLogInfo, RecomputeTargetMap, - RecomputeUnit, + SaveUnit, apply_selective_checkpointing, module_dict_repr, ) @@ -138,7 +138,7 @@ # 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( + SaveUnit.ATTN: KeptOps( "flash_attn::_flash_attn_varlen_forward_v3", "flash_attn::_flash_attn_varlen_forward_v2", ), @@ -147,15 +147,13 @@ # 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"), + SaveUnit.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) - ), + SaveUnit.MOE_DISPATCH: KeptCallables(*(f"{cls}.{stage}" for cls in _DISPATCHERS for stage in _DISPATCH_STAGES)), } @@ -1564,7 +1562,7 @@ def _should_recompute( mtp_layers = 1 if self.config.mtp_config.share_weights else self.config.mtp_config.num_layers else: mtp_layers = 0 - recompute_ratio = self.fsdp_config.recompute_ratio if self.fsdp_config is not None else 0.0 + recompute_ratio = self.config.recompute_cfg.ratio total_layers = num_layers + mtp_layers num_recompute_layers = int(total_layers * recompute_ratio) diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py index 7f35465937..948e45653d 100644 --- a/xtuner/v1/model/utils/__init__.py +++ b/xtuner/v1/model/utils/__init__.py @@ -1,8 +1,9 @@ from xtuner.v1.utils.selective_checkpointing import ( KeptCallables, KeptOps, + RecomputeConfig, RecomputeTargetMap, - RecomputeUnit, + SaveUnit, ) from .checkpointing import apply_gradient_checkpointing @@ -18,7 +19,8 @@ "KeptCallables", "KeptOps", "RecomputeTargetMap", - "RecomputeUnit", + "RecomputeConfig", + "SaveUnit", "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 ce0f5c76e9..3b39a00915 100644 --- a/xtuner/v1/model/utils/selective_checkpointing.py +++ b/xtuner/v1/model/utils/selective_checkpointing.py @@ -23,7 +23,7 @@ from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts from xtuner.v1.utils import log_rank0 -from xtuner.v1.utils.selective_checkpointing import RecomputeUnit, active_recompute_unit, recompute_unit +from xtuner.v1.utils.selective_checkpointing import SaveUnit, active_recompute_unit, recompute_unit from .checkpointing import apply_gradient_checkpointing, checkpoint_flattened, install_checkpointing @@ -66,7 +66,7 @@ def apply_selective_checkpointing( return install_checkpointing(module, partial(_run_checkpointed_region, kept_ops, preserve_rng_state)) -def in_recompute_unit(unit: RecomputeUnit, func: Any) -> Any: +def in_recompute_unit(unit: SaveUnit, func: Any) -> Any: """Wrap ``func`` so everything it dispatches belongs to ``unit``. Installed by the model layer on each callable a @@ -74,7 +74,7 @@ def in_recompute_unit(unit: RecomputeUnit, func: Any) -> Any: the exclusion from compilation that makes the marker readable at all. Args: - unit (RecomputeUnit): The unit this callable implements. + unit (SaveUnit): The unit this callable implements. func (Any): The callable to wrap. Returns: diff --git a/xtuner/v1/utils/__init__.py b/xtuner/v1/utils/__init__.py index 5c4c85618b..638d67a640 100644 --- a/xtuner/v1/utils/__init__.py +++ b/xtuner/v1/utils/__init__.py @@ -25,9 +25,10 @@ from .selective_checkpointing import ( KeptCallables, KeptOps, + RecomputeConfig, RecomputeTarget, RecomputeTargetMap, - RecomputeUnit, + SaveUnit, ) from .state import ForwardState from .type_helper import copy_method_signature, copy_signature, ray_method @@ -79,5 +80,6 @@ "KeptOps", "RecomputeTarget", "RecomputeTargetMap", - "RecomputeUnit", + "RecomputeConfig", + "SaveUnit", ] diff --git a/xtuner/v1/utils/selective_checkpointing.py b/xtuner/v1/utils/selective_checkpointing.py index 50f0f2e6da..50d7a41e49 100644 --- a/xtuner/v1/utils/selective_checkpointing.py +++ b/xtuner/v1/utils/selective_checkpointing.py @@ -2,7 +2,7 @@ This module holds the vocabulary that the SAC layers agree on and nothing else: -- Users select :class:`RecomputeUnit` members in the model config. +- Users select :class:`SaveUnit` 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. @@ -35,11 +35,17 @@ from dataclasses import dataclass from typing import TypeAlias +from cyclopts import Parameter +from pydantic import BaseModel as PydanticBaseModel +from pydantic import ConfigDict +from typing_extensions import Annotated + from .enum_helper import StrEnum __all__ = [ - "RecomputeUnit", + "SaveUnit", + "RecomputeConfig", "KeptOps", "KeptCallables", "RecomputeTarget", @@ -49,19 +55,22 @@ ] -class RecomputeUnit(StrEnum): - """Semantic units of activation that may be kept resident instead of - recomputed. +class SaveUnit(StrEnum): + """Semantic units of activation that are *saved* -- kept resident instead + of recomputed. + + Naming is deliberate: selecting a member means "save this", not "recompute this". Everything not + selected is recomputed, which is the default and costs no memory. - 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. + Each member names a *class of sub-structure* whose activations are worth keeping because + recomputing it is expensive relative to what it costs to store. 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" + ATTN = "attn" """Keep the attention kernel's output -- the flash-attention call itself, not the projections around it. @@ -70,10 +79,10 @@ class RecomputeUnit(StrEnum): is always visible to the checkpoint policy, compiled or not. """ - SAVE_MOE_GATE = "save_moe_gate" + MOE_GATE = "moe_gate" """Keep the MoE router: gating projection, top-k selection, and routing weights.""" - SAVE_MOE_DISPATCH = "save_moe_dispatch" + MOE_DISPATCH = "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. @@ -123,31 +132,63 @@ def __init__(self, *names: str) -> None: RecomputeTarget: TypeAlias = KeptOps | KeptCallables -"""What a single :class:`RecomputeUnit` resolves to for one architecture.""" +"""What a single :class:`SaveUnit` resolves to for one architecture.""" + + +class RecomputeConfig(PydanticBaseModel): + """All recompute settings in one place: which layers, and what stays saved + inside them. + + ``ratio`` and ``vision_ratio`` decide *which layers* are recomputed at all; ``save`` decides + *what inside a recomputed layer* is kept resident anyway. They used to live in two different + configs -- the ratios on ``FSDPConfig``, the unit selection on the model config -- which made + the pair hard to reason about, since neither means much without the other. + + Args: + ratio (float): Fraction of decoder layers to recompute, counted from the end. Defaults to + 1.0, i.e. recompute every layer. + vision_ratio (float): The same, for a compose model's vision tower. Defaults to 1.0. + save (list[SaveUnit] | bool | None): What stays resident inside the recomputed layers. + ``None`` and ``False`` save nothing, ``True`` saves every unit the model declares, and a + list saves exactly those. Defaults to None. + + ``None`` deliberately does not mean "model default", unlike ``compile_cfg``: saving a + unit trades memory for speed, so an unset config must not silently change an existing + run's peak memory. + """ + + model_config = ConfigDict(extra="forbid") + + ratio: Annotated[float, Parameter(help="Fraction of decoder layers to recompute, from the end")] = 1.0 + vision_ratio: Annotated[float, Parameter(help="The same, for a compose model's vision tower")] = 1.0 + save: Annotated[ + list[SaveUnit] | bool | None, + Parameter(help="What stays resident inside recomputed layers: None/False none, True all declared, or a list"), + ] = None + -RecomputeTargetMap: TypeAlias = dict[RecomputeUnit, RecomputeTarget] -"""Per-model declaration of what each supported :class:`RecomputeUnit` resolves -to. +RecomputeTargetMap: TypeAlias = dict[SaveUnit, RecomputeTarget] +"""Per-model declaration of what each supported :class:`SaveUnit` 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: +def active_recompute_unit() -> SaveUnit | 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. + SaveUnit | None: The innermost open unit, or None outside one. """ return _ACTIVE_UNIT.get() @contextmanager -def recompute_unit(unit: RecomputeUnit) -> Iterator[None]: +def recompute_unit(unit: SaveUnit) -> 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 @@ -155,7 +196,7 @@ def recompute_unit(unit: RecomputeUnit) -> Iterator[None]: 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. + unit (SaveUnit): The unit the enclosed call belongs to. """ token = _ACTIVE_UNIT.set(unit) try: @@ -164,4 +205,4 @@ def recompute_unit(unit: RecomputeUnit) -> Iterator[None]: _ACTIVE_UNIT.reset(token) -_ACTIVE_UNIT: ContextVar[RecomputeUnit | None] = ContextVar("xtuner_recompute_unit", default=None) +_ACTIVE_UNIT: ContextVar[SaveUnit | None] = ContextVar("xtuner_recompute_unit", default=None)