From d6b3091fcea7007befaa41157cb743ab2e29b8d5 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 28 Jul 2026 04:26:55 +0000 Subject: [PATCH 1/5] [Refactor] Split the selective checkpointing contract below the model layer `checkpoint_record` is called from `xtuner/v1/module/decoder_layer/*.py`, and importing it from there pulls in `xtuner/v1/model/__init__.py`, which imports the decoder layers back. A module the `module/` layer depends on cannot live under `model/`. Split rather than move: the vocabulary and the marker session go to `xtuner/v1/utils`, below both packages, while the policy and the wrapping stay at the model layer, where knowing `nn.Module`, `CheckpointWrapper` and the DSA lifecycle predicate is legitimate. The session's API becomes public because it is now driven across a package boundary. `xtuner.v1.model.utils` keeps exporting the same names, so no import site outside these files changes. --- tests/model/test_selective_checkpointing.py | 36 +- xtuner/v1/model/utils/__init__.py | 9 +- .../v1/model/utils/selective_checkpointing.py | 249 ++------------ xtuner/v1/utils/__init__.py | 5 + xtuner/v1/utils/selective_checkpointing.py | 309 ++++++++++++++++++ 5 files changed, 363 insertions(+), 245 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 e3b8577f42..93f46aa1e7 100644 --- a/tests/model/test_selective_checkpointing.py +++ b/tests/model/test_selective_checkpointing.py @@ -5,6 +5,8 @@ test_unbalanced_interval_is_safe: end marker 不执行时只多留驻，不影响梯度。 test_overlapping_intervals_keep_the_union: 区间重叠时按并集留驻。 test_marker_outside_session_is_noop: 不在会话内时埋点不做任何事。 +TestContractLayering + test_module_layer_imports_the_contract_without_the_model_layer: 契约模块必须能被 module/ 层单独导入。 TestUnsupportedRegions test_a_second_model_still_gets_its_own_diagnosis: 诊断去重不跨模型，第二个模型仍会告警。 test_in_place_op_in_kept_region_is_rejected: 留驻区间内的 in-place 写会明确报错而不是静默改梯度。 @@ -14,6 +16,8 @@ """ import os +import subprocess +import sys import pytest import torch @@ -28,7 +32,7 @@ from xtuner.v1.model.utils import apply_selective_checkpointing, checkpoint_record from xtuner.v1.module.attention import MHAConfig from xtuner.v1.module.router import NoAuxRouterConfig -from xtuner.v1.model.utils import selective_checkpointing as contract +from xtuner.v1.utils import selective_checkpointing as contract class _MarkedBlock(nn.Module): @@ -52,22 +56,6 @@ class _OtherMarkedBlock(_MarkedBlock): """A second layer class, standing in for another model living in the same process.""" -class _InPlaceBlock(nn.Module): - """A region whose body accumulates in place, which selective checkpointing cannot keep.""" - - def __init__(self) -> None: - super().__init__() - self.linear = nn.Linear(4, 4) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - checkpoint_record("region.start") - hidden = self.linear(x) - accumulator = torch.zeros_like(hidden) - accumulator.add_(hidden) - checkpoint_record("region.end") - return accumulator - - class _EmptyRegionBlock(nn.Module): """A layer whose declared region encloses no op the policy can keep. @@ -163,6 +151,20 @@ def test_marker_outside_session_is_noop(self): assert inputs.grad is not None +class TestContractLayering: + def test_module_layer_imports_the_contract_without_the_model_layer(self): + # 契约（含 marker session）之所以在 xtuner/v1/utils 而不是挨着 engine，就是因为 + # `checkpoint_record` 的调用点在 xtuner/v1/module 的 forward 里：一旦契约里出现 + # 指向 model/ 或 module/ 的 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 TestUnsupportedRegions: def test_in_place_op_in_kept_region_is_rejected(self): # 留驻的张量会原样交给重算那趟，read-modify-write 会在重算时二次累加：loss 有限、 diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py index 273ef0668b..e3b710af5d 100644 --- a/xtuner/v1/model/utils/__init__.py +++ b/xtuner/v1/model/utils/__init__.py @@ -1,13 +1,14 @@ -from .checkpointing import apply_gradient_checkpointing -from .misc import ModelForwardExtraLogInfo, module_dict_repr -from .selective_checkpointing import ( +from xtuner.v1.utils.selective_checkpointing import ( MarkerInterval, RecomputeIntervalMap, RecomputeUnit, - apply_selective_checkpointing, checkpoint_record, ) +from .checkpointing import apply_gradient_checkpointing +from .misc import ModelForwardExtraLogInfo, module_dict_repr +from .selective_checkpointing import apply_selective_checkpointing + __all__ = [ "apply_gradient_checkpointing", diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py index fd27415784..f5a4eb6ca8 100644 --- a/xtuner/v1/model/utils/selective_checkpointing.py +++ b/xtuner/v1/model/utils/selective_checkpointing.py @@ -1,13 +1,10 @@ -"""Region-level selective activation checkpointing (SAC). +"""The selective activation checkpointing engine. -The three SAC layers meet here: - -- Model authors call :func:`checkpoint_record` inside ``forward`` to name addressable semantic - boundaries, and declare a :data:`RecomputeIntervalMap` mapping each :class:`RecomputeUnit` they - support to the marker intervals that implement it for their architecture. -- Users select :class:`RecomputeUnit` members in the model config; they never see marker strings. -- The sharding paths hand the resolved intervals to :func:`apply_selective_checkpointing`, which - wraps one layer in a single checkpoint whose per-op policy keeps those intervals resident. +Applies a recompute strategy to one decoder layer: the whole layer goes under a single +``torch.utils.checkpoint`` region whose per-op policy keeps the marker intervals the config layer +selected and recomputes everything else. The vocabulary and the marker session it drives live in +:mod:`xtuner.v1.utils.selective_checkpointing`, below the model layer, because the marker calls +themselves sit in ``xtuner.v1.module`` forwards. Granularity note, because it decides which units are worth declaring for a given model: a region is addressable only where its **contents** run in eager python, not merely its endpoints. A @@ -17,80 +14,26 @@ compiled ``_shared_experts_forward``. In eager, every region is addressable. """ -import contextvars -import weakref from collections.abc import Sequence from functools import partial -from typing import Any, TypeAlias +from typing import Any import torch import torch.nn as nn from torch.utils.checkpoint import CheckpointPolicy, checkpoint, create_selective_checkpoint_contexts from xtuner.v1.utils import log_rank0 -from xtuner.v1.utils.enum_helper import StrEnum +from xtuner.v1.utils.selective_checkpointing import ( + MarkerInterval, + active_marker_session, + declare_selective_regions, + marker_session, +) from .checkpointing import CheckpointWrapper, apply_gradient_checkpointing -__all__ = [ - "RecomputeUnit", - "MarkerInterval", - "RecomputeIntervalMap", - "apply_selective_checkpointing", - "checkpoint_record", -] - - -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. Which marker intervals a unit resolves to is architecture-specific and is - declared per model, so the same unit can cover different regions in different models. - - Members are strings so pydantic round-trips them to readable names in serialized configs. - """ - - SAVE_ATTN = "save_attn" - """Keep the attention core: the flash-attention / SDPA call and its immediate surroundings.""" - - 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 expert dispatch / combine communication region.""" - - SAVE_MLP = "save_mlp" - """Keep the dense MLP / shared-expert region.""" - - -MarkerInterval: TypeAlias = tuple[str, str] -"""A half-open ``[start, end)`` interval between two :func:`checkpoint_record` -marker names. - -Ops executed at or after the ``start`` marker and strictly before the ``end`` marker belong to the -interval. ``end`` is the name of the marker that begins the *next* region, which is why the -interval is half-open: regions are delimited by points, not by explicit closing markers. - -Intervals are resolved by program order at runtime, so an interval may span module boundaries (its -``start`` in one module's ``forward`` and its ``end`` in a sibling module's). Unbalanced intervals --- for instance an ``end`` marker that sits on a branch the forward did not take -- are not an -error: they only widen the region that stays resident. Both the kept and the recomputed paths are -numerically exact, so marker bookkeeping can never corrupt gradients. -""" - - -RecomputeIntervalMap: TypeAlias = dict[RecomputeUnit, list[MarkerInterval]] -"""Per-model declaration of how each supported :class:`RecomputeUnit` maps to -marker intervals. - -Models expose this as a ``default_recompute_cfg`` property, mirroring ``default_compile_cfg``. A -unit absent from the mapping is simply not supported by that architecture and contributes no -intervals. -""" +__all__ = ["apply_selective_checkpointing"] def apply_selective_checkpointing( @@ -131,170 +74,31 @@ def apply_selective_checkpointing( Returns: nn.Module: The checkpoint-wrapped layer, transparent to parameter names and ``state_dict``. """ - intervals = tuple(intervals) + kept_intervals = tuple(intervals) - if not intervals: + if not kept_intervals: return apply_gradient_checkpointing(module, preserve_rng_state=preserve_rng_state) if layer_compiled_as_one_region: - _warn_intervals_unsupported(module, intervals, "it is compiled as a single region") + _warn_intervals_unsupported(module, kept_intervals, "it is compiled as a single region") # Not routed through `apply_gradient_checkpointing` because the marker session has to open # *inside* the checkpointed region: `use_reentrant=False` re-runs this forward to recompute, and # a session opened around the checkpoint call would be long gone by then. - _declare_selective_regions(owner, intervals) - checkpointed_call = partial(_run_checkpointed_region, intervals, preserve_rng_state, owner) + declare_selective_regions(owner, kept_intervals) + checkpointed_call = partial(_run_checkpointed_region, kept_intervals, preserve_rng_state, owner) return CheckpointWrapper(module, checkpointed_call) -def checkpoint_record(name: str) -> None: - """Mark an addressable semantic boundary in the current forward pass. - - This is a point marker, not a region: it records "execution reached ``name`` here". Regions are - formed by pairing markers into :data:`MarkerInterval` s, which is why a marker can open a region - in one module and have it closed by a marker in another -- ordering is by runtime program order, - not lexical scope. - - Outside an active SAC session the call is a no-op, so models may be instrumented independently - of whether selective checkpointing is enabled. It carries no autograd semantics and is safe to - call during recomputation. - - Args: - name (str): Marker name, unique within a layer's forward. Use a dotted, structural name - such as ``"attn.core"`` or ``"moe.dispatch"`` so that ``default_recompute_cfg`` entries - read as coordinates into the architecture. - """ - # The marker session is backed by contextvars, which Dynamo cannot trace: reading a ContextVar - # inside a `fullgraph=True` region is a hard compile error rather than a graph break, and xtuner - # compiles several MoE forward methods that way. Dynamo constant-folds this check, so the body - # never enters the graph and instrumented models stay compilable. - # - # This branch must stay free of side effects for the same reason -- a global mutation or a log - # call here would break `fullgraph=True`. Markers that turn out to be inert are reported from - # `_report_pass`, which runs in eager python once the owner's layers have all been through. - if torch.compiler.is_compiling(): - return - - session = _MARKER_SESSION.get() - if session is not None: - session.record(name) - - -_MARKER_SESSION: contextvars.ContextVar["_MarkerSession | None"] = contextvars.ContextVar( - "xtuner_sac_marker_session", default=None -) - # Ops from these namespaces are never kept, whatever the markers say. See `_checkpoint_policy`. _NEVER_KEPT_NAMESPACES = ("c10d", "_c10d_functional") # Mutating ops that leave tensor *values* alone, so a kept region may contain them. Anything else -# with a mutable schema is rejected by `_reject_in_place_op_in_kept_region`. +# with a mutable schema goes through `_reject_non_replayable_op_in_kept_region`. _VALUE_PRESERVING_MUTATING_OPS = frozenset({torch.ops.aten.record_stream.default}) -# Diagnostics accumulate per model and are held weakly, so they neither keep models alive nor let -# one model's conclusions silence another's. -_DIAGNOSTICS: "weakref.WeakKeyDictionary[nn.Module, _OwnerDiagnostics]" = weakref.WeakKeyDictionary() -_REPORTED_UNSUPPORTED_DIAGNOSES: set[tuple[type, tuple["MarkerInterval", ...], str]] = set() - - -class _MarkerSession: - """Tracks which marker intervals are open at the current point of one pass - through a checkpointed layer.""" - - def __init__(self, intervals: tuple[MarkerInterval, ...], owner: nn.Module | None = None) -> None: - self._intervals = intervals - self._owner = owner - self._open: set[MarkerInterval] = set() - self._recorded: set[str] = set() - self._kept: set[MarkerInterval] = set() - - @property - def keeping(self) -> bool: - # Overlapping and nested intervals are defined as "any interval open => keep", which is why - # this is set emptiness rather than a paired counter: no pairing is asserted anywhere. - return bool(self._open) - - def record(self, name: str) -> None: - self._recorded.add(name) - for interval in self._intervals: - start, end = interval - if name == start: - self._open.add(interval) - elif name == end: - self._open.discard(interval) - - def note_kept(self) -> None: - # The only evidence that a region is addressable at all. Markers merely delimit; an interval - # whose contents run inside a compiled region has both endpoints fire and still keeps - # nothing, because those ops execute as fused kernels and never reach the policy. - self._kept |= self._open - - def finish(self) -> None: - _report_pass(self._owner, self._intervals, self._recorded, self._kept) - - -class _OwnerDiagnostics: - def __init__(self) -> None: - self.expected_layers = 0 - self.completed_layers = 0 - self.recorded_markers: set[str] = set() - self.kept_intervals: set[MarkerInterval] = set() - self.declared_intervals: set[MarkerInterval] = set() - self.reported: set[MarkerInterval] = set() - - -def _declare_selective_regions(owner: nn.Module | None, intervals: tuple[MarkerInterval, ...]) -> None: - # Diagnostics are only trustworthy once every layer has run: a marker missing from one layer - # type is normal when the intervals span several -- `mlp.begin` never runs in a MoE layer and - # `moe.gate.begin` never runs in a dense one -- and warning per layer fires on models that are - # configured correctly. This is how the diagnostics know how many layers to wait for. - if owner is None or not intervals: - return - state = _DIAGNOSTICS.get(owner) - if state is None: - state = _OwnerDiagnostics() - _DIAGNOSTICS[owner] = state - state.expected_layers += 1 - - -def _report_pass( - owner: nn.Module | None, - intervals: tuple[MarkerInterval, ...], - recorded: set[str], - kept: set[MarkerInterval], -) -> None: - state = _DIAGNOSTICS.get(owner) if owner is not None else None - if state is None: - # No owner declared these regions -- a direct caller, or a test. Diagnose the region alone. - state = _OwnerDiagnostics() - state.expected_layers = 1 - - state.completed_layers += 1 - state.recorded_markers |= recorded - state.kept_intervals |= kept - state.declared_intervals |= set(intervals) - - # Wait for a full pass over the owner's layers before concluding anything: only then has every - # layer type had its chance to record a marker and keep an op. - if state.completed_layers < state.expected_layers: - return - - for interval in sorted(state.declared_intervals): - if interval in state.kept_intervals or interval in state.reported: - continue - state.reported.add(interval) - if interval[0] not in state.recorded_markers: - log_rank0.warning( - f"Selective checkpointing: marker {interval[0]!r} of interval {interval} never ran, so this " - f"interval keeps nothing resident and its region is recomputed. Either the model does not record " - f"that marker, or it records it inside a torch.compile'd region, where markers have no effect." - ) - else: - log_rank0.warning( - f"Selective checkpointing: interval {interval} opened but kept nothing resident, so its region is " - f"recomputed. Its contents run inside a torch.compile'd region or consist only of ops that are " - f"never kept, such as collectives; markers delimiting a compiled region do not make it addressable." - ) +# Reported once per distinct diagnosis, not once per layer. See `_warn_intervals_unsupported`. +_REPORTED_UNSUPPORTED_DIAGNOSES: set[tuple[type, tuple[MarkerInterval, ...], str]] = set() def _run_checkpointed_region( @@ -334,12 +138,9 @@ def _forward_with_marker_session( if torch.compiler.is_compiling(): return module(*args, **kwargs) - session = _MarkerSession(intervals, owner) - token = _MARKER_SESSION.set(session) - try: + with marker_session(intervals, owner=owner) as session: output = module(*args, **kwargs) - finally: - _MARKER_SESSION.reset(token) + # Only a pass that ran to the end counts: with `set_checkpoint_early_stop` on, the recompute # pass is cut short by an exception once the last needed tensor is repacked, and folding a # truncated pass into the diagnostics would blame regions that simply had not come up yet. @@ -352,7 +153,7 @@ def _selective_checkpoint_contexts() -> tuple[Any, Any]: def _checkpoint_policy(ctx: Any, op: Any, *args: Any, **kwargs: Any) -> CheckpointPolicy: - session = _MARKER_SESSION.get() + session = active_marker_session() if session is None or not session.keeping: return CheckpointPolicy.MUST_RECOMPUTE diff --git a/xtuner/v1/utils/__init__.py b/xtuner/v1/utils/__init__.py index 915e2a0c79..7fc8819089 100644 --- a/xtuner/v1/utils/__init__.py +++ b/xtuner/v1/utils/__init__.py @@ -22,6 +22,7 @@ ) 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 MarkerInterval, RecomputeIntervalMap, RecomputeUnit, checkpoint_record 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 +69,8 @@ "trim_memory", "group_tensors_by_device_mesh_and_placements", "cal_total_norm", + "MarkerInterval", + "RecomputeIntervalMap", + "RecomputeUnit", + "checkpoint_record", ] diff --git a/xtuner/v1/utils/selective_checkpointing.py b/xtuner/v1/utils/selective_checkpointing.py new file mode 100644 index 0000000000..4381916861 --- /dev/null +++ b/xtuner/v1/utils/selective_checkpointing.py @@ -0,0 +1,309 @@ +"""Shared contract for region-level selective activation checkpointing (SAC). + +This module holds the vocabulary that the three SAC layers agree on and nothing else: + +- Model authors call :func:`checkpoint_record` inside ``forward`` to name addressable semantic + boundaries, and declare a :data:`RecomputeIntervalMap` mapping each :class:`RecomputeUnit` they + support to the marker intervals that implement it for their architecture. +- Users select :class:`RecomputeUnit` members in the model config; they never see marker strings. +- The SAC engine turns the selected units into intervals and drives a per-op checkpoint policy. + +It lives under ``xtuner.v1.utils`` rather than next to the models because the marker calls sit in +``xtuner.v1.module`` forwards while the unit vocabulary is consumed by ``xtuner.v1.model`` configs; +a home inside either package would make the two import each other. + +The vocabulary and the marker session live here -- the session because :func:`checkpoint_record` +is what reads it, and that call sits in ``xtuner.v1.module``. The checkpoint policy and the wrapping +that drives a session live in ``xtuner.v1.model.utils.selective_checkpointing``, which is free to +import from the model and module layers; the config resolution that turns user selections into +intervals lives with the model configs. + +Splitting the session down here means the engine has to drive it across a package boundary, which +is the *only* reason the session's API is public: + +- :func:`marker_session` opens a session for one pass through a checkpointed region, and must be + entered inside the checkpointed callable so that the recompute pass rebuilds its own state. +- :func:`active_marker_session` is what the per-op policy asks whether a region is currently open. +- :class:`MarkerSession` carries that state: ``keeping`` (is any interval open), ``record`` (driven + by :func:`checkpoint_record`), ``note_kept`` (the policy reporting that it kept an op) and + ``finish`` (fold this pass into the owner's diagnostics). +- :func:`declare_selective_regions` tells the diagnostics how many layers to wait for before + concluding that an interval kept nothing. + +Treat that surface as internal to the engine rather than as an extension point: it is public +because of where the package boundary fell, and it carries no stability promise. +""" + +import contextvars +import weakref +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from typing import Any, TypeAlias + +import torch + +from .enum_helper import StrEnum +from .logger import log_rank0 + + +__all__ = [ + "RecomputeUnit", + "MarkerInterval", + "RecomputeIntervalMap", + "MarkerSession", + "active_marker_session", + "declare_selective_regions", + "marker_session", + "checkpoint_record", +] + + +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. Which marker intervals a unit resolves to is architecture-specific and is + declared per model, so the same unit can cover different regions in different models. + + Members are strings so pydantic round-trips them to readable names in serialized configs. + """ + + SAVE_ATTN = "save_attn" + """Keep the attention core: the flash-attention / SDPA call and its immediate surroundings.""" + + 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 expert dispatch / combine communication region.""" + + SAVE_MLP = "save_mlp" + """Keep the dense MLP / shared-expert region.""" + + +MarkerInterval: TypeAlias = tuple[str, str] +"""A half-open ``[start, end)`` interval between two :func:`checkpoint_record` +marker names. + +Ops executed at or after the ``start`` marker and strictly before the ``end`` marker belong to the +interval. + +Intervals are resolved by program order at runtime, so an interval may span module boundaries (its +``start`` in one module's ``forward`` and its ``end`` in a sibling module's). Unbalanced intervals +-- for instance an ``end`` marker that sits on a branch the forward did not take -- are not an +error: they only widen the region that stays resident. Both the kept and the recomputed paths are +numerically exact, so marker bookkeeping can never corrupt gradients. +""" + + +RecomputeIntervalMap: TypeAlias = dict[RecomputeUnit, list[MarkerInterval]] +"""Per-model declaration of how each supported :class:`RecomputeUnit` maps to +marker intervals. + +Models expose this as a ``default_recompute_cfg`` property, mirroring ``default_compile_cfg``. A +unit absent from the mapping is simply not supported by that architecture and contributes no +intervals. +""" + + +def checkpoint_record(name: str) -> None: + """Mark an addressable semantic boundary in the current forward pass. + + This is a point marker, not a region: it records "execution reached ``name`` here". Regions are + formed by pairing markers into :data:`MarkerInterval` s, which is why a marker can open a region + in one module and have it closed by a marker in another -- ordering is by runtime program order, + not lexical scope. + + Outside an active SAC session the call is a no-op, so models may be instrumented independently + of whether selective checkpointing is enabled. It carries no autograd semantics and is safe to + call during recomputation. + + Args: + name (str): Marker name, unique within a layer's forward. Use a dotted, structural name + such as ``"attn.begin"`` or ``"moe.dispatch.end"`` so that ``default_recompute_cfg`` + entries read as coordinates into the architecture. + """ + # The marker session is backed by contextvars, which Dynamo cannot trace: reading a ContextVar + # inside a `fullgraph=True` region is a hard compile error rather than a graph break, and xtuner + # compiles several MoE forward methods that way. Dynamo constant-folds this check, so the body + # never enters the graph and instrumented models stay compilable. + # + # This branch must stay free of side effects for the same reason -- a global mutation or a log + # call here would break `fullgraph=True`. Markers that turn out to be inert are reported from + # `_report_pass`, which runs in eager python once the owner's layers have all been through. + if torch.compiler.is_compiling(): + return + + session = _MARKER_SESSION.get() + if session is not None: + session.record(name) + + +def active_marker_session() -> "MarkerSession | None": + """Return the marker session of the region currently executing, if any. + + Returns: + MarkerSession | None: The active session, or None outside a selectively checkpointed region. + """ + return _MARKER_SESSION.get() + + +@contextmanager +def marker_session(intervals: Sequence[MarkerInterval], owner: Any = None) -> Iterator["MarkerSession"]: + """Open a marker session for one pass through a checkpointed region. + + The session must be opened *inside* the checkpointed callable: ``use_reentrant=False`` re-runs + that callable to recompute, and both passes have to build their marker state from scratch or the + two sets of kept ops drift apart. + + Args: + intervals (Sequence[MarkerInterval]): Intervals this region keeps resident. + owner (Any): The model these layers belong to, used only to aggregate diagnostics across + its layers. Pass the same object that was given to :func:`declare_selective_regions`. + Defaults to None, which diagnoses this region on its own. + + Returns: + Iterator[MarkerSession]: The session that :func:`checkpoint_record` will drive. + """ + session = MarkerSession(tuple(intervals), owner) + token = _MARKER_SESSION.set(session) + try: + yield session + finally: + _MARKER_SESSION.reset(token) + + +def declare_selective_regions(owner: Any, intervals: Sequence[MarkerInterval]) -> None: + """Declare that one more layer of ``owner`` will run with ``intervals``. + + Diagnostics are only trustworthy once every such layer has run: a marker missing from one layer + type is normal when the intervals span several -- ``mlp.begin`` never runs in a MoE layer and + ``moe.gate.begin`` never runs in a dense one -- and warning per layer would fire on models that + are configured correctly. This tells the diagnostics how many layers to wait for. + + Args: + owner (Any): The model that owns the layer, held weakly. + intervals (Sequence[MarkerInterval]): The intervals that layer will keep resident. + """ + if owner is None or not intervals: + return + state = _DIAGNOSTICS.get(owner) + if state is None: + state = _OwnerDiagnostics() + _DIAGNOSTICS[owner] = state + state.expected_layers += 1 + + +class MarkerSession: + """Tracks which marker intervals are open at the current point of one pass + through a region. + + Driven by :func:`checkpoint_record` and read by the checkpoint policy; the SAC engine owns its + lifetime through :func:`marker_session`. + """ + + def __init__(self, intervals: tuple[MarkerInterval, ...], owner: Any = None) -> None: + self._intervals = intervals + self._owner = owner + self._open: set[MarkerInterval] = set() + self._recorded: set[str] = set() + self._kept: set[MarkerInterval] = set() + + @property + def keeping(self) -> bool: + """Whether execution is currently inside at least one interval. + + Returns: + bool: True while any interval is open. + """ + # Overlapping and nested intervals are defined as "any interval open => keep", which is why + # this is set emptiness rather than a paired counter: no pairing is asserted anywhere. + return bool(self._open) + + def record(self, name: str) -> None: + """Register that execution reached the marker ``name``. + + Args: + name (str): The marker name passed to :func:`checkpoint_record`. + """ + self._recorded.add(name) + for interval in self._intervals: + start, end = interval + if name == start: + self._open.add(interval) + elif name == end: + self._open.discard(interval) + + def note_kept(self) -> None: + """Register that an op was kept resident while the open intervals were + open.""" + # This is the only evidence that a region is addressable at all. Markers merely delimit; an + # interval whose contents run inside a compiled region has both endpoints fire and still + # keeps nothing, because those ops execute as fused kernels and never reach the policy. + self._kept |= self._open + + def finish(self) -> None: + """Fold this pass into the owner's diagnostics, warning once the model + has run in full.""" + _report_pass(self._owner, self._intervals, self._recorded, self._kept) + + +class _OwnerDiagnostics: + def __init__(self) -> None: + self.expected_layers = 0 + self.completed_layers = 0 + self.recorded_markers: set[str] = set() + self.kept_intervals: set[MarkerInterval] = set() + self.declared_intervals: set[MarkerInterval] = set() + self.reported: set[MarkerInterval] = set() + + +def _report_pass( + owner: Any, + intervals: tuple[MarkerInterval, ...], + recorded: set[str], + kept: set[MarkerInterval], +) -> None: + state = _DIAGNOSTICS.get(owner) if owner is not None else None + if state is None: + # No owner declared these regions -- a direct caller, or a test. Diagnose the region alone. + state = _OwnerDiagnostics() + state.expected_layers = 1 + + state.completed_layers += 1 + state.recorded_markers |= recorded + state.kept_intervals |= kept + state.declared_intervals |= set(intervals) + + # Wait for a full pass over the owner's layers before concluding anything: only then has every + # layer type had its chance to record a marker and keep an op. + if state.completed_layers < state.expected_layers: + return + + for interval in sorted(state.declared_intervals): + if interval in state.kept_intervals or interval in state.reported: + continue + state.reported.add(interval) + if interval[0] not in state.recorded_markers: + log_rank0.warning( + f"Selective checkpointing: marker {interval[0]!r} of interval {interval} never ran, so this " + f"interval keeps nothing resident and its region is recomputed. Either the model does not record " + f"that marker, or it records it inside a torch.compile'd region, where markers have no effect." + ) + else: + log_rank0.warning( + f"Selective checkpointing: interval {interval} opened but kept nothing resident, so its region is " + f"recomputed. Its contents run inside a torch.compile'd region or consist only of ops that are " + f"never kept, such as collectives; markers delimiting a compiled region do not make it addressable." + ) + + +_MARKER_SESSION: contextvars.ContextVar["MarkerSession | None"] = contextvars.ContextVar( + "xtuner_sac_marker_session", default=None +) + +# Diagnostics accumulate per model and are held weakly, so they neither keep models alive nor let +# one model's conclusions silence another's. +_DIAGNOSTICS: "weakref.WeakKeyDictionary[Any, _OwnerDiagnostics]" = weakref.WeakKeyDictionary() From d582fe18161b6ef64a7814ae8e8fc85b2e792b42 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 28 Jul 2026 04:27:09 +0000 Subject: [PATCH 2/5] [Feature] Add region-level `recompute_cfg` and per-model marker declarations Lets a user keep named activation regions resident instead of recomputing them, inside the layers `fsdp_cfg.recompute_ratio` already selects. The two knobs stay orthogonal: `recompute_ratio` picks the layers, `recompute_cfg` picks the regions inside a selected layer. `XTunerBaseModelConfig.recompute_cfg` is tri-state. `None` and `False` keep today's behaviour of recomputing everything, `True` keeps every region the model declares, and an explicit list of `RecomputeUnit` keeps exactly those. Unlike `compile_cfg`, `None` is not "the model default": `default_recompute_cfg` declares what an architecture *can* keep, not what is worth keeping, so an unset config never changes a run's memory profile. `False` additionally propagates into nested sub-model configs, which is the walk `compile_cfg` already had, now shared between the two switches instead of duplicated. `MoEDecoderLayer` and `DenseDecoderLayer` record paired markers around the attention, router, dispatch, combine, shared-expert and MLP regions, in both the single-batch and the domino EP path. How much of that survives `torch.compile` depends on where the markers sit, which the per-model `default_recompute_cfg` docstrings spell out unit by unit. --- tests/model/test_recompute.py | 168 +++++++++++++++++- tests/model/test_selective_checkpointing.py | 6 +- xtuner/v1/model/base.py | 120 ++++++++++--- xtuner/v1/model/dense/dense.py | 27 ++- xtuner/v1/model/moe/moe.py | 43 +++++ .../decoder_layer/dense_decoder_layer.py | 6 +- .../module/decoder_layer/moe_decoder_layer.py | 28 ++- xtuner/v1/train/trainer.py | 6 + 8 files changed, 373 insertions(+), 31 deletions(-) diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py index aafe2f13f7..4c3c2a1ad6 100644 --- a/tests/model/test_recompute.py +++ b/tests/model/test_recompute.py @@ -5,9 +5,23 @@ test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。 TestDominoEPRecompute test_recompute_matches_baseline_under_domino_ep: domino EP 下重算与不重算的 loss/梯度一致。 +TestRecomputeCfgResolution + test_unset_cfg_keeps_full_recompute: `None` 不改变显存行为，解析为空区间。 + test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。 + test_explicit_units_select_only_their_intervals: 显式 list 只解析出对应区间。 + test_string_units_are_accepted: 配置文件里的字符串能解析成 RecomputeUnit。 + test_unsupported_unit_is_rejected: 模型不支持的 unit 报错并列出支持项。 + test_disable_propagates_into_nested_configs: `False` 递归关闭嵌套子模型配置。 + test_units_round_trip_through_json: enum 序列化成可读字符串并能读回。 +TestMarkerVocabulary + test_declared_intervals_have_markers: default_recompute_cfg 引用的 marker 都真实埋点。 + test_micro_batch_path_records_the_same_markers: 单路与 domino 路埋点集合一致。 """ +import ast +import inspect import os +import textwrap import pytest import torch @@ -17,9 +31,15 @@ from xtuner._testing import DeterministicDDPTestCase from xtuner.v1.config import FSDPConfig from xtuner.v1.loss.ce_loss import CELossConfig -from xtuner.v1.model.moe.moe import MoE, MoEConfig, SequenceContext -from xtuner.v1.model.utils import apply_gradient_checkpointing +from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig +from xtuner.v1.model.dense.dense import DENSE_RECOMPUTE_CFG +from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG, MoE, MoEConfig, SequenceContext +from xtuner.v1.model.utils import ( + RecomputeUnit, + apply_gradient_checkpointing, +) from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.decoder_layer import dense_decoder_layer, moe_decoder_layer from xtuner.v1.module.router import NoAuxRouterConfig @@ -185,3 +205,147 @@ def _run_once(self, ep_size: int, dispatcher: str, recompute_ratio: float): del model, out, loss torch.cuda.empty_cache() return loss_val, grad_norm, all_finite + + +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_intervals(**overrides) -> list[tuple[str, str]]: + with torch.device("meta"): + return MoE(config=_build_tiny_moe_config(**overrides)).recompute_intervals + + +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_intervals(recompute_cfg=None) == [] + + def test_true_selects_every_supported_unit(self): + intervals = _resolve_intervals(recompute_cfg=True) + + expected = [interval for unit_intervals in MOE_RECOMPUTE_CFG.values() for interval in unit_intervals] + assert intervals == expected + + def test_explicit_units_select_only_their_intervals(self): + intervals = _resolve_intervals(recompute_cfg=[RecomputeUnit.SAVE_MOE_DISPATCH]) + + assert intervals == MOE_RECOMPUTE_CFG[RecomputeUnit.SAVE_MOE_DISPATCH] + + def test_string_units_are_accepted(self): + # Configs arrive as JSON/py files where units are written as plain strings. + assert _resolve_intervals(recompute_cfg=["save_attn"]) == MOE_RECOMPUTE_CFG[RecomputeUnit.SAVE_ATTN] + + def test_unsupported_unit_is_rejected(self): + # A hybrid model declares no units, so any selection is a user configuration error rather + # than something to silently drop. + 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.recompute_intervals == [] + assert config.text_config.recompute_cfg is False + assert config.text_config.compile_cfg is None + + 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_MLP]) + + dumped = config.model_dump(mode="json")["recompute_cfg"] + assert dumped == ["save_attn", "save_mlp"] + + restored = _build_tiny_moe_config(recompute_cfg=dumped) + assert restored.recompute_cfg == [RecomputeUnit.SAVE_ATTN, RecomputeUnit.SAVE_MLP] + + +def _recorded_markers(func) -> set[str]: + """Marker names a function passes to ``checkpoint_record``.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + return { + node.args[0].value + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "checkpoint_record" + and node.args + and isinstance(node.args[0], ast.Constant) + } + + +class TestMarkerVocabulary: + @pytest.mark.parametrize( + "interval_map", [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG], ids=["moe", "dense"] + ) + def test_declared_intervals_have_markers(self, interval_map): + # A renamed marker would not fail anywhere at runtime: the interval would simply never open + # and the region would stay recomputed, silently costing the memory the user asked to keep. + recorded: set[str] = set() + for layer_module in (moe_decoder_layer, dense_decoder_layer): + for _, member in inspect.getmembers(layer_module, inspect.isclass): + if member.__module__ != layer_module.__name__: + continue + for _, method in inspect.getmembers(member, inspect.isfunction): + recorded |= _recorded_markers(method) + + declared = {name for intervals in interval_map.values() for interval in intervals for name in interval} + assert declared <= recorded, f"markers referenced but never recorded: {sorted(declared - recorded)}" + + def test_micro_batch_path_records_the_same_markers(self): + # `_micro_batch_forward` re-implements the dispatch/combine chain across four stage loops. + # If the two paths drift, domino EP silently keeps a different set of regions resident. + single = _recorded_markers(moe_decoder_layer.MoEDecoderLayer._forward) + micro_batch = _recorded_markers(moe_decoder_layer.MoEDecoderLayer._micro_batch_forward) + + assert single == micro_batch diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py index 93f46aa1e7..0530ae85fa 100644 --- a/tests/model/test_selective_checkpointing.py +++ b/tests/model/test_selective_checkpointing.py @@ -29,7 +29,7 @@ from xtuner.v1.config import FSDPConfig from xtuner.v1.loss.ce_loss import CELossConfig from xtuner.v1.model.moe.moe import MoE, MoEConfig, SequenceContext -from xtuner.v1.model.utils import apply_selective_checkpointing, checkpoint_record +from xtuner.v1.model.utils import MarkerInterval, apply_selective_checkpointing, checkpoint_record from xtuner.v1.module.attention import MHAConfig from xtuner.v1.module.router import NoAuxRouterConfig from xtuner.v1.utils import selective_checkpointing as contract @@ -265,8 +265,8 @@ class _RegionMoE(MoE): _REGION = ("moe.dispatch", "moe.combine.end") @property - def recompute_intervals(self) -> tuple[tuple[str, str], ...]: - return (self._REGION,) + def recompute_intervals(self) -> list[MarkerInterval]: + return [self._REGION] def fully_shard(self, *args, **kwargs): for layer in self.layers.values(): diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 4f51640923..2608de7ccf 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,7 @@ set_async_save_process_qos, ) -from .utils import MarkerInterval, ModelForwardExtraLogInfo +from .utils import MarkerInterval, ModelForwardExtraLogInfo, RecomputeIntervalMap, RecomputeUnit logger = get_logger() @@ -143,6 +142,17 @@ class XTunerBaseModelConfig(PydanticBaseModel): "`dict[str, TorchCompileOption]`: Customize the compile option", ), ] = None + 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 +548,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 @@ -557,6 +597,7 @@ def __init__(self, config: XTunerBaseModelConfig): self._async_hf_resources: AsyncHFResources | None = None self._compile_cfg = self._resolve_compile_cfg(self.config) + self._recompute_intervals = self._resolve_recompute_cfg(self.config) self._float8_handler: Float8Handler | None = None def set_hf(self, hf_path: str | Path): @@ -985,18 +1026,33 @@ def compile_cfg(self) -> dict[str, TorchCompileOption]: return _compile_cfg @property - def recompute_intervals(self) -> tuple[MarkerInterval, ...]: - """Marker intervals kept resident inside every recomputed layer. + def default_recompute_cfg(self) -> RecomputeIntervalMap: + """Marker intervals this architecture can keep resident, keyed by + semantic unit. + + This is the model author's vocabulary: it declares which :class:`RecomputeUnit` s the architecture supports and + where each one lives, not which of them are worth enabling. A model that has no ``checkpoint_record`` markers, + or whose markers are all inert in the setup it ships with, declares nothing. + + Returns: + RecomputeIntervalMap: Supported units mapped to the marker intervals that implement them. + """ + return {} + + @property + def recompute_intervals(self) -> list[MarkerInterval]: + """Marker intervals to keep resident in the layers selected for + recompute. - This is the whole surface between the config layer and the checkpointing mechanism: the - sharding paths pass whatever this returns to ``apply_selective_checkpointing``. Resolving - the user's ``recompute_cfg`` against a model's ``default_recompute_cfg`` belongs in an - override here; the default keeps nothing, which recomputes each selected layer whole. + Unlike ``compile_cfg``, this deliberately does not merge sub-models: ``compile_cfg`` names free functions and + methods that are patched process-wide, so every sub-model's entries must reach the single patching step, while + intervals are per-layer policy consumed at the sharding site of the sub-model that owns those layers. A compose + model's vision tower and language model each resolve their own. Returns: - tuple[MarkerInterval, ...]: Half-open ``[start, end)`` marker intervals. + list[MarkerInterval]: Intervals selected by ``config.recompute_cfg``. Empty means plain full recompute. """ - return () + return self._recompute_intervals @property def float8_handler(self): @@ -2564,14 +2620,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: @@ -2581,17 +2637,35 @@ def _resolve_compile_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 + def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> list[MarkerInterval]: + selected = config.recompute_cfg + + if selected is False: + _disable_nested_switch(self.config, "recompute_cfg") + return [] + + # `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 region 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 [] + + supported = self.default_recompute_cfg + units = list(supported) if selected is True else selected + + intervals: list[MarkerInterval] = [] + 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." + ) + intervals.extend(supported[unit]) + return intervals def _maybe_enable_compile(self, compile_cfg: dict[str, TorchCompileOption]): if compile_cfg: diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index 38286b574d..e07cc22092 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 RecomputeIntervalMap, RecomputeUnit, apply_selective_checkpointing from xtuner.v1.module import ( GatedDeltaNetConfig, LMHead, @@ -51,6 +51,11 @@ **DEFAULT_FLOAT8_CFG, } +DENSE_RECOMPUTE_CFG: RecomputeIntervalMap = { + RecomputeUnit.SAVE_ATTN: [("attn.begin", "attn.end")], + RecomputeUnit.SAVE_MLP: [("mlp.begin", "mlp.end")], +} + class Dense(BaseModel): config: TransformerConfig @@ -164,6 +169,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) -> RecomputeIntervalMap: + """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: + RecomputeIntervalMap: 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 diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 1538039b38..6e8e6f36f5 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -46,6 +46,8 @@ ) from xtuner.v1.model.utils import ( ModelForwardExtraLogInfo, + RecomputeIntervalMap, + RecomputeUnit, apply_selective_checkpointing, module_dict_repr, ) @@ -113,6 +115,21 @@ MOE_EP_COMPILE_CFG = MOE_NON_EP_COMPILE_CFG.copy() MOE_EP_COMPILE_CFG.pop(MOE_DECODER_LAYER_FORWARD) +MOE_RECOMPUTE_CFG: RecomputeIntervalMap = { + RecomputeUnit.SAVE_ATTN: [("attn.begin", "attn.end")], + RecomputeUnit.SAVE_MOE_GATE: [("moe.gate.begin", "moe.gate.end")], + RecomputeUnit.SAVE_MOE_DISPATCH: [ + ("moe.dispatch.begin", "moe.dispatch.end"), + ("moe.combine.begin", "moe.combine.end"), + ], + # The first `first_k_dense_replace` layers of a MoE model are dense layers, whose MLP is a + # different region from a MoE layer's shared experts. Both belong to the same user-facing unit. + RecomputeUnit.SAVE_MLP: [ + ("moe.shared_experts.begin", "moe.shared_experts.end"), + ("mlp.begin", "mlp.end"), + ], +} + class MoEModelOutputs(ModelOutputs): router_logits: dict[str, torch.Tensor] | None = None @@ -1275,6 +1292,32 @@ def default_compile_cfg(self) -> dict[str, TorchCompileOption]: else: return MOE_NON_EP_COMPILE_CFG + @property + @override + def default_recompute_cfg(self) -> RecomputeIntervalMap: + """Marker intervals this architecture can keep resident, keyed by + semantic unit. + + How much of this survives ``torch.compile`` depends on where each region's markers sit. A marker inside a + compiled region is folded away, so only regions delimited in the uncompiled layer body remain addressable: + + - ``ep_size > 1``: ``SAVE_MOE_DISPATCH`` and the shared-expert half of ``SAVE_MLP`` are delimited in + ``MoEDecoderLayer._forward`` / ``_micro_batch_forward`` and stay effective. ``SAVE_ATTN``, + ``SAVE_MOE_GATE`` and the dense-layer half of ``SAVE_MLP`` are delimited inside compiled methods and take + effect in eager only. + - ``ep_size == 1``: the whole layer forward is compiled as one region, so every unit is eager-only. + + A unit that is inert simply leaves its region recomputed, which is the behaviour of plain full recompute. + + Returns: + RecomputeIntervalMap: 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 diff --git a/xtuner/v1/module/decoder_layer/dense_decoder_layer.py b/xtuner/v1/module/decoder_layer/dense_decoder_layer.py index aa5660ea09..94d1a00aa2 100644 --- a/xtuner/v1/module/decoder_layer/dense_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/dense_decoder_layer.py @@ -9,7 +9,7 @@ from xtuner.v1.module import AttnOutputs, GatedDeltaNetConfig, MHAConfig, MLAConfig, RMSNorm from xtuner.v1.module.rope import RopeScalingConfig from xtuner.v1.ops.act_fn import get_act_fn -from xtuner.v1.utils import ForwardState +from xtuner.v1.utils import ForwardState, checkpoint_record from ..linear import build_linear @@ -155,18 +155,22 @@ def _forward( hidden_states = self.input_layernorm(hidden_states) # Self Attention + checkpoint_record("attn.begin") attn_outputs: AttnOutputs = self.self_attn( hidden_states=hidden_states, position_embeddings=position_embeddings, seq_ctx=seq_ctx, ) hidden_states = attn_outputs["projected_output"] + checkpoint_record("attn.end") hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) + checkpoint_record("mlp.begin") hidden_states = self.mlp(hidden_states) + checkpoint_record("mlp.end") hidden_states = residual + hidden_states return hidden_states diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 37ec22bb12..4746b6d82e 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -36,7 +36,7 @@ from xtuner.v1.module.grouped_linear.moe_group_linear import build_grouped_linear from xtuner.v1.module.rope import RopeScalingConfig from xtuner.v1.ops.act_fn import get_act_fn -from xtuner.v1.utils import ForwardState +from xtuner.v1.utils import ForwardState, checkpoint_record from ..linear import build_linear @@ -420,6 +420,7 @@ def _forward( # ProberList.before_dispatch( # self.layer_idx, hidden_states, router_results["topk_ids"], router_results["topk_weights"] # ) + checkpoint_record("moe.dispatch.begin") pre_dispatched = self.dispatcher.dispatch_preprocess( hidden_states=hidden_states.view(-1, hidden_states.shape[-1]), topk_ids=router_results["topk_ids"], @@ -433,6 +434,7 @@ def _forward( pre_dispatched=pre_dispatched, dispatched=dispatched, ) + checkpoint_record("moe.dispatch.end") # ProberList.after_dispatch( # self.layer_idx, # post_dispatched["hidden_states"], @@ -451,6 +453,7 @@ def _forward( # post_dispatched.get("row_ids_map"), # type: ignore[arg-type] # dispatched["topk_weights"], # ) + checkpoint_record("moe.combine.begin") pre_combined = self.dispatcher.combine_preprocess( hidden_states=experts_out, pre_dispatched=pre_dispatched, @@ -475,16 +478,21 @@ def _forward( ) combined_hidden_states = post_combined["hidden_states"] combined_hidden_states = combined_hidden_states.view(*origin_shape) + checkpoint_record("moe.combine.end") # 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. + checkpoint_record("moe.shared_experts.begin") if self.n_shared_experts > 0: shared_experts_out = self._shared_experts_forward(hidden_states=hidden_states) else: shared_experts_out = None + checkpoint_record("moe.shared_experts.end") hidden_states = self._post_moe_forward( combined_hidden_states=combined_hidden_states, @@ -534,11 +542,13 @@ def _micro_batch_forward( ) pre_moe_forward_out_list.append(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + checkpoint_record("moe.dispatch.begin") pre_dispatched = self.dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=router_results["topk_ids"], async_op=True, ) + checkpoint_record("moe.dispatch.end") pre_dispatched_list.append(pre_dispatched) residual_list.append(residual) router_results_list.append(router_results) @@ -553,6 +563,7 @@ def _micro_batch_forward( router_results_list, pre_dispatched_list, ): + checkpoint_record("moe.dispatch.begin") dispatched = self.dispatcher.dispatch( pre_dispatched=pre_dispatched, topk_weights=router_results["topk_weights"], @@ -564,12 +575,14 @@ def _micro_batch_forward( dispatched=dispatched, async_op=True, ) + checkpoint_record("moe.dispatch.end") experts_out = self.experts( post_dispatched["hidden_states"], post_dispatched["tokens_per_expert"], decoding=False, ) + checkpoint_record("moe.combine.begin") pre_combined = self.dispatcher.combine_preprocess( hidden_states=experts_out, pre_dispatched=pre_dispatched, @@ -577,6 +590,7 @@ def _micro_batch_forward( post_dispatched=post_dispatched, async_op=True, ) + checkpoint_record("moe.combine.end") post_dispatched_list.append(post_dispatched) experts_out_list.append(experts_out) @@ -589,6 +603,7 @@ def _micro_batch_forward( dispatched_list, post_dispatched_list, ): + checkpoint_record("moe.combine.begin") combined = self.dispatcher.combine( pre_combined=pre_combined, pre_dispatched=pre_dispatched, @@ -596,10 +611,14 @@ def _micro_batch_forward( post_dispatched=post_dispatched, async_op=True, ) + checkpoint_record("moe.combine.end") combined_list.append(combined) 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. + checkpoint_record("moe.shared_experts.begin") if self.n_shared_experts > 0: shared_experts_out_list = [] for pre_moe_forward_out in pre_moe_forward_out_list: @@ -609,9 +628,11 @@ def _micro_batch_forward( shared_experts_out_list.append(shared_experts_out) else: shared_experts_out_list = [None] * intra_layer_micro_batch + checkpoint_record("moe.shared_experts.end") hidden_states_out_list: list[torch.Tensor] = [] for i in range(intra_layer_micro_batch): + checkpoint_record("moe.combine.begin") post_combined = self.dispatcher.combine_postprocess( pre_dispatched=pre_dispatched_list[i], dispatched=dispatched_list[i], @@ -620,6 +641,7 @@ def _micro_batch_forward( combined=combined_list[i], async_op=True, ) + checkpoint_record("moe.combine.end") hidden_states = self._post_moe_forward( # hidden_states=pre_moe_forward_out_list[i], combined_hidden_states=post_combined["hidden_states"].view(*pre_moe_forward_out_list[i].shape), @@ -649,6 +671,7 @@ def _pre_moe_forward( hidden_states = self.input_layernorm(hidden_states) # Self Attention + checkpoint_record("attn.begin") if state == ForwardState.TRAINING: attn_outputs: AttnOutputs = self.self_attn( hidden_states=hidden_states, @@ -672,6 +695,7 @@ def _pre_moe_forward( seq_ctx=seq_ctx, past_key_values=past_key_values, ) + checkpoint_record("attn.end") hidden_states = residual + hidden_states # Fully Connected @@ -687,7 +711,9 @@ def _pre_moe_forward( rollout_routed_experts = rollout_routed_experts.to(hidden_states.device) else: rollout_routed_experts = None + checkpoint_record("moe.gate.begin") router_results: RouterResults = self.gate(hidden_states, rollout_routed_experts) + checkpoint_record("moe.gate.end") return residual, hidden_states, router_results def _shared_experts_forward( diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index 1d279156d9..f1fe796bbb 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -1152,6 +1152,12 @@ def build_engine( if engine.model.compile_cfg is not None: log_rank0.info(f"The `compile_cfg` of model is {json.dumps(engine.model.compile_cfg, indent=4)}") + # Only reported when something is actually kept resident: this is the memory knob to look at + # first when a run's peak does not match the one it was tuned for. + if engine.model.recompute_intervals: + log_rank0.info( + f"The `recompute_cfg` of model keeps these regions resident: {engine.model.recompute_intervals}" + ) return engine def build_lr_scheduler(self, lr_cfg: LRConfig, scheduler_step: int) -> torch.optim.lr_scheduler.LRScheduler: From f2671e4bff5e0a7a6fe9e46f7f93d978bf6de5bc Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 28 Jul 2026 05:02:44 +0000 Subject: [PATCH 3/5] [Fix] Assert region coverage and reject unsupported units at construction `test_micro_batch_path_records_the_same_markers` compared marker name sets, so it would have passed even if the domino path wrapped entirely different operations. It now compares which of the layer's own operations each region encloses, verified by mutation: moving `moe.dispatch.end` past the expert GEMMs fails the new assertion and passed the old one. Along the way the two paths are made to agree exactly -- `_forward` now reshapes outside the dispatch and combine regions, as the domino path already did. Intervals resolve in `BaseModel.__init__`, next to `compile_cfg` and by the same rule: `default_recompute_cfg` is answerable from the config alone, so there is nothing to wait for. A config naming a unit the model does not support now fails before the run spends anything on materializing and sharding weights. `recompute_cfg=False` reaching every nested sub-model config gains a test on a shipped compose config, which nests three of them; the previous probe nested one. Verified by mutation: stopping the walk at the outer config fails it. Record why regions use explicit `.begin` / `.end` marker pairs instead of ending each region at the next region's start marker, and warn instead of silently resolving to nothing when `recompute_cfg=True` meets a model that declares no units. --- tests/model/test_recompute.py | 78 ++++++++++++++++--- xtuner/v1/model/base.py | 18 +++++ xtuner/v1/model/dense/dense.py | 1 + xtuner/v1/model/moe/moe.py | 8 ++ .../module/decoder_layer/moe_decoder_layer.py | 13 +++- 5 files changed, 105 insertions(+), 13 deletions(-) diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py index 4c3c2a1ad6..f87c668be8 100644 --- a/tests/model/test_recompute.py +++ b/tests/model/test_recompute.py @@ -10,12 +10,13 @@ test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。 test_explicit_units_select_only_their_intervals: 显式 list 只解析出对应区间。 test_string_units_are_accepted: 配置文件里的字符串能解析成 RecomputeUnit。 - test_unsupported_unit_is_rejected: 模型不支持的 unit 报错并列出支持项。 + 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 序列化成可读字符串并能读回。 TestMarkerVocabulary test_declared_intervals_have_markers: default_recompute_cfg 引用的 marker 都真实埋点。 - test_micro_batch_path_records_the_same_markers: 单路与 domino 路埋点集合一致。 + test_micro_batch_path_covers_the_same_operations: 单路与 domino 路每个区间覆盖的算子一致。 """ import ast @@ -31,7 +32,8 @@ from xtuner._testing import DeterministicDDPTestCase from xtuner.v1.config import FSDPConfig from xtuner.v1.loss.ce_loss import CELossConfig -from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig +from xtuner.v1.model.base import BaseModel, 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, SequenceContext from xtuner.v1.model.utils import ( @@ -282,8 +284,9 @@ def test_string_units_are_accepted(self): assert _resolve_intervals(recompute_cfg=["save_attn"]) == MOE_RECOMPUTE_CFG[RecomputeUnit.SAVE_ATTN] def test_unsupported_unit_is_rejected(self): - # A hybrid model declares no units, so any selection is a user configuration error rather - # than something to silently drop. + # 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])) @@ -298,6 +301,18 @@ def test_disable_propagates_into_nested_configs(self): 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. @@ -324,6 +339,50 @@ def _recorded_markers(func) -> set[str]: } +def _region_coverage(func) -> dict[str, set[str]]: + """Which of the layer's own operations each marker region encloses. + + Regions are keyed by the shared prefix of a ``.begin`` / ``.end`` pair, and an operation is a call on + ``self`` -- ``self.experts(...)``, ``self.dispatcher.dispatch(...)``. Calls on tensors are ignored: a region is + defined by the sub-modules it covers, not by the reshapes threaded between them. + """ + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + # `ast.walk` is breadth-first; the marker state machine needs source order. + nodes = sorted( + (node for node in ast.walk(tree) if isinstance(node, ast.Call)), + key=lambda node: (node.lineno, node.col_offset), + ) + + coverage: dict[str, set[str]] = {} + active: set[str] = set() + for node in nodes: + name = _called_name(node) + if name == "checkpoint_record" and node.args and isinstance(node.args[0], ast.Constant): + region, _, edge = node.args[0].value.rpartition(".") + if edge == "begin": + active.add(region) + coverage.setdefault(region, set()) + elif edge == "end": + active.discard(region) + elif name is not None and name.startswith("self."): + for region in active: + coverage[region].add(name.removeprefix("self.")) + return coverage + + +def _called_name(node: ast.Call) -> str | None: + """Dotted name of a call target, e.g. ``self.dispatcher.dispatch``.""" + parts: list[str] = [] + target: ast.expr = node.func + while isinstance(target, ast.Attribute): + parts.append(target.attr) + target = target.value + if not isinstance(target, ast.Name): + return None + parts.append(target.id) + return ".".join(reversed(parts)) + + class TestMarkerVocabulary: @pytest.mark.parametrize( "interval_map", [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG], ids=["moe", "dense"] @@ -342,10 +401,11 @@ def test_declared_intervals_have_markers(self, interval_map): declared = {name for intervals in interval_map.values() for interval in intervals for name in interval} assert declared <= recorded, f"markers referenced but never recorded: {sorted(declared - recorded)}" - def test_micro_batch_path_records_the_same_markers(self): + def test_micro_batch_path_covers_the_same_operations(self): # `_micro_batch_forward` re-implements the dispatch/combine chain across four stage loops. - # If the two paths drift, domino EP silently keeps a different set of regions resident. - single = _recorded_markers(moe_decoder_layer.MoEDecoderLayer._forward) - micro_batch = _recorded_markers(moe_decoder_layer.MoEDecoderLayer._micro_batch_forward) + # Comparing marker names alone would pass even if the domino path wrapped entirely different + # operations, so compare what each region actually encloses. + single = _region_coverage(moe_decoder_layer.MoEDecoderLayer._forward) + micro_batch = _region_coverage(moe_decoder_layer.MoEDecoderLayer._micro_batch_forward) assert single == micro_batch diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 2608de7ccf..8148fd606e 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -142,6 +142,10 @@ 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( @@ -1034,6 +1038,9 @@ def default_recompute_cfg(self) -> RecomputeIntervalMap: where each one lives, not which of them are worth enabling. A model that has no ``checkpoint_record`` markers, or whose markers are all inert in the setup it ships with, declares nothing. + Like ``default_compile_cfg``, an override must be answerable from ``self.config`` alone: it is read while + ``BaseModel.__init__`` resolves the user's selection, which is before the subclass has built its layers. + Returns: RecomputeIntervalMap: Supported units mapped to the marker intervals that implement them. """ @@ -2640,6 +2647,8 @@ def _resolve_compile_cfg( def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> list[MarkerInterval]: 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 [] @@ -2653,6 +2662,15 @@ def _resolve_recompute_cfg(self, config: XTunerBaseModelConfig) -> list[MarkerIn return [] 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 intervals: list[MarkerInterval] = [] diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py index e07cc22092..05305b5514 100644 --- a/xtuner/v1/model/dense/dense.py +++ b/xtuner/v1/model/dense/dense.py @@ -51,6 +51,7 @@ **DEFAULT_FLOAT8_CFG, } +# Explicit `.begin` / `.end` marker pairs, for the reasons recorded next to `MOE_RECOMPUTE_CFG`. DENSE_RECOMPUTE_CFG: RecomputeIntervalMap = { RecomputeUnit.SAVE_ATTN: [("attn.begin", "attn.end")], RecomputeUnit.SAVE_MLP: [("mlp.begin", "mlp.end")], diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 6e8e6f36f5..910caa82e4 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -115,6 +115,13 @@ MOE_EP_COMPILE_CFG = MOE_NON_EP_COMPILE_CFG.copy() MOE_EP_COMPILE_CFG.pop(MOE_DECODER_LAYER_FORWARD) +# Regions are delimited by an explicit `.begin` / `.end` marker pair rather than by ending each region at +# the marker that starts the next one. Both express the same half-open `[start, end)` interval, but "end = whatever +# comes next" couples every region to its successor: reordering two regions in a forward, or inserting one between +# them, would silently redefine what the earlier region keeps. Explicit pairs also survive `_micro_batch_forward` +# splitting the dispatch/combine chain across four stage loops, where "the next region" differs from the single-batch +# path. A missing `.end` only leaves the region open to the end of the layer, which costs retained memory and never +# gradients, since the kept and recomputed paths are both numerically exact. MOE_RECOMPUTE_CFG: RecomputeIntervalMap = { RecomputeUnit.SAVE_ATTN: [("attn.begin", "attn.end")], RecomputeUnit.SAVE_MOE_GATE: [("moe.gate.begin", "moe.gate.end")], @@ -1316,6 +1323,7 @@ def default_recompute_cfg(self) -> RecomputeIntervalMap: # 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 diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 4746b6d82e..25d1e2e291 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -417,12 +417,15 @@ 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"] # ) checkpoint_record("moe.dispatch.begin") 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( @@ -476,9 +479,8 @@ def _forward( pre_combined=pre_combined, combined=combined, ) - combined_hidden_states = post_combined["hidden_states"] - combined_hidden_states = combined_hidden_states.view(*origin_shape) checkpoint_record("moe.combine.end") + 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) @@ -617,7 +619,10 @@ 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. + # 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. checkpoint_record("moe.shared_experts.begin") if self.n_shared_experts > 0: shared_experts_out_list = [] From 855eae08d0df631d402515d39a087660378de407 Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 28 Jul 2026 05:04:52 +0000 Subject: [PATCH 4/5] [Docs] Say what SAVE_MOE_DISPATCH actually keeps "Keep the expert dispatch / combine communication region" tells a user they are keeping the communication, which is the one thing the unit does not keep: the SAC policy forces every c10d collective to be recomputed, because keeping one would elide it from the recompute pass. What the unit trades memory for is the permutation, padding and unpermutation work on either side of the all-to-all. --- xtuner/v1/utils/selective_checkpointing.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/xtuner/v1/utils/selective_checkpointing.py b/xtuner/v1/utils/selective_checkpointing.py index 4381916861..9bbd03fcb7 100644 --- a/xtuner/v1/utils/selective_checkpointing.py +++ b/xtuner/v1/utils/selective_checkpointing.py @@ -77,7 +77,14 @@ class RecomputeUnit(StrEnum): """Keep the MoE router: gating projection, top-k selection, and routing weights.""" SAVE_MOE_DISPATCH = "save_moe_dispatch" - """Keep the expert dispatch / combine communication region.""" + """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. + """ SAVE_MLP = "save_mlp" """Keep the dense MLP / shared-expert region.""" From 0e87277af63482358f89aa132566d5612ea3a9be Mon Sep 17 00:00:00 2001 From: HAOCHENYE <21724054@zju.edu.cn> Date: Tue, 28 Jul 2026 05:11:34 +0000 Subject: [PATCH 5/5] [Docs] Correct which recompute units survive torch.compile Measured against the SAC policy at ep_size=2: only SAVE_MOE_DISPATCH keeps ops under compile (152), while SAVE_ATTN, SAVE_MOE_GATE and SAVE_MLP keep none. The docstring claimed the shared-expert half of SAVE_MLP stayed effective. The rule was stated as "markers must sit outside compiled regions", which is necessary but not sufficient. SAVE_MLP's markers do run in eager -- they sit in `_forward`, which EP does not compile -- but the region encloses `_shared_experts_forward`, which EP does compile, and ops inside a compiled region execute as fused kernels that never reach the per-op policy. A region is addressable only if its contents run in eager, not merely its endpoints. SAVE_MOE_DISPATCH survives because the dispatcher calls it wraps are the one part of the MoE path that stays uncompiled. --- xtuner/v1/model/moe/moe.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 910caa82e4..2e18dc7711 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1305,16 +1305,21 @@ def default_recompute_cfg(self) -> RecomputeIntervalMap: """Marker intervals this architecture can keep resident, keyed by semantic unit. - How much of this survives ``torch.compile`` depends on where each region's markers sit. A marker inside a - compiled region is folded away, so only regions delimited in the uncompiled layer body remain addressable: - - - ``ep_size > 1``: ``SAVE_MOE_DISPATCH`` and the shared-expert half of ``SAVE_MLP`` are delimited in - ``MoEDecoderLayer._forward`` / ``_micro_batch_forward`` and stay effective. ``SAVE_ATTN``, - ``SAVE_MOE_GATE`` and the dense-layer half of ``SAVE_MLP`` are delimited inside compiled methods and take - effect in eager only. + Under ``torch.compile`` a region is only addressable if the ops it encloses run in eager: a marker inside a + compiled region is folded away, and ops inside one execute as fused kernels that never reach the per-op + checkpoint policy. Both a region's markers *and* its contents therefore have to sit outside compiled code. + + - ``ep_size > 1``: only ``SAVE_MOE_DISPATCH`` survives, because the dispatcher calls it wraps are not + compiled. ``SAVE_ATTN`` and ``SAVE_MOE_GATE`` are marked inside the compiled ``_pre_moe_forward``. + ``SAVE_MLP`` is inert on both of its intervals, for the two different reasons the mechanism allows: its + shared-experts markers do fire in eager but enclose the compiled ``_shared_experts_forward``, while its + ``mlp`` markers never fire at all, sitting in a ``DenseDecoderLayer`` that stays compiled whole even + under EP. - ``ep_size == 1``: the whole layer forward is compiled as one region, so every unit is eager-only. A unit that is inert simply leaves its region recomputed, which is the behaviour of plain full recompute. + Measured per unit against the SAC policy: eager keeps ops for all four; ``ep_size=2`` under compile keeps + 152 ops for ``SAVE_MOE_DISPATCH`` and none for the others. Returns: RecomputeIntervalMap: Supported units mapped to the marker intervals that implement them.