From 7bf4688ab8d1b0dfcb99fc77a9dd94096ba025ea Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 28 Jul 2026 04:32:31 +0000
Subject: [PATCH 1/6] [Feature] Add the region-level selective checkpointing
 engine

Turn the whole-layer checkpoint into a per-op decision: a contextvars marker
session records which `checkpoint_record` intervals are open, and the policy
behind `create_selective_checkpoint_contexts` keeps the ops inside them while
recomputing the rest. The sharding paths call one entry point per selected
layer, which also owns the reentrant fallback for DSA layers.

Two op classes are never kept, whatever the markers say: ops with a mutable
schema, whose cached tensor can be overwritten before backward reads it, and
collectives, whose destination buffer may be allocated outside the interval.
---
 xtuner/v1/model/base.py                       |  16 +-
 .../compose/intern_s1/modeling_vision.py      |  10 +-
 .../model/compose/qwen3_vl/modeling_vision.py |  10 +-
 xtuner/v1/model/dense/dense.py                |  10 +-
 xtuner/v1/model/moe/moe.py                    |  27 ++-
 xtuner/v1/model/utils/__init__.py             |   2 +
 xtuner/v1/model/utils/checkpointing.py        |   2 +-
 .../v1/model/utils/selective_checkpointing.py | 225 +++++++++++++++++-
 8 files changed, 280 insertions(+), 22 deletions(-)

diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py
index 0856a3fd7..4f5164092 100644
--- a/xtuner/v1/model/base.py
+++ b/xtuner/v1/model/base.py
@@ -62,7 +62,7 @@
     set_async_save_process_qos,
 )
 
-from .utils import ModelForwardExtraLogInfo
+from .utils import MarkerInterval, ModelForwardExtraLogInfo
 
 
 logger = get_logger()
@@ -984,6 +984,20 @@ 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.
+
+        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.
+
+        Returns:
+            tuple[MarkerInterval, ...]: Half-open ``[start, end)`` marker intervals.
+        """
+        return ()
+
     @property
     def float8_handler(self):
         if (
diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
index a971efcc6..37cf3e2a2 100644
--- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py
+++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
@@ -34,7 +34,7 @@
     fully_shard,
 )
 from xtuner.v1.ops.attn_imp import attn_impl_mapping, AttnOpOutputs
-from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing
+from xtuner.v1.model.utils import apply_selective_checkpointing
 from xtuner.v1.module import RMSNorm
 from xtuner.v1.ops.others import Dropout
 from xtuner.v1.ops.act_fn import get_act_fn
@@ -407,7 +407,13 @@ def fully_shard(
             layer = self.encoder.layer[layer_idx]
 
             if layer_idx < num_recompute_layers:
-                layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state)
+                layer = apply_selective_checkpointing(
+                    layer,
+                    self.recompute_intervals,
+                    preserve_rng_state=checkpoint_preserve_rng_state,
+                    # The layer's own forward is compiled just below, making it one opaque region.
+                    layer_compiled_as_one_region=bool(self.compile_cfg),
+                )
                 if self.config.drop_path_rate == 0.0 and self.compile_cfg:
                     layer.forward = torch.compile(layer.forward, fullgraph=True)
 
diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
index aa1dd2a83..1fbc3a677 100644
--- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
+++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
@@ -24,7 +24,7 @@
 from torch.distributed.device_mesh import init_device_mesh
 import torch.distributed as dist
 from xtuner.v1.utils.compile import maybe_compile
-from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing
+from xtuner.v1.model.utils import apply_selective_checkpointing
 from xtuner.v1.module import AttnOutputs
 from torch.distributed.device_mesh import DeviceMesh
 from tqdm import tqdm
@@ -338,7 +338,13 @@ def fully_shard(
             layer = self.blocks[layer_idx]
 
             if layer_idx < num_recompute_layers:
-                layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state)
+                layer = apply_selective_checkpointing(
+                    layer,
+                    self.recompute_intervals,
+                    preserve_rng_state=checkpoint_preserve_rng_state,
+                    # The layer's own forward is compiled just below, making it one opaque region.
+                    layer_compiled_as_one_region=bool(self.compile_cfg),
+                )
                 if self.compile_cfg:
                     layer.forward = torch.compile(layer.forward, fullgraph=True)
 
diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py
index 59a9a7df0..f83e2df17 100644
--- a/xtuner/v1/model/dense/dense.py
+++ b/xtuner/v1/model/dense/dense.py
@@ -26,7 +26,7 @@
     TorchCompileOption,
     TransformerConfig,
 )
-from xtuner.v1.model.utils import apply_gradient_checkpointing
+from xtuner.v1.model.utils import apply_selective_checkpointing
 from xtuner.v1.module import (
     GatedDeltaNetConfig,
     LMHead,
@@ -235,7 +235,13 @@ def fully_shard(
             layer = self.layers[str(int(layer_idx))]
             layer_idx = int(layer_idx)
             if layer_idx < num_recompute_layers:
-                layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state)
+                layer = apply_selective_checkpointing(
+                    layer,
+                    self.recompute_intervals,
+                    preserve_rng_state=checkpoint_preserve_rng_state,
+                    # The layer's own forward is compiled just below, making it one opaque region.
+                    layer_compiled_as_one_region=bool(self.compile_cfg),
+                )
 
                 # Linear-attention (GatedDeltaNet) layers write ``seq_ctx.seq_idx`` inside the
                 # checkpoint region; compiling the checkpointed layer with ``fullgraph=True`` turns
diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py
index 316bdfca0..d1c8eac05 100644
--- a/xtuner/v1/model/moe/moe.py
+++ b/xtuner/v1/model/moe/moe.py
@@ -46,7 +46,7 @@
 )
 from xtuner.v1.model.utils import (
     ModelForwardExtraLogInfo,
-    apply_gradient_checkpointing,
+    apply_selective_checkpointing,
     module_dict_repr,
 )
 from xtuner.v1.module import (
@@ -90,9 +90,13 @@
 logger = get_logger()
 
 
+# Compiling the decoder layer as a whole is what makes marker intervals inert, so the key is named
+# rather than spelled out at each use site.
+MOE_DECODER_LAYER_FORWARD = "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward"
+
 MOE_NON_EP_COMPILE_CFG: dict[str, TorchCompileOption] = {
     "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEBlock.forward": TorchCompileOption(fullgraph=True),
-    "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward": TorchCompileOption(fullgraph=True),
+    MOE_DECODER_LAYER_FORWARD: TorchCompileOption(fullgraph=True),
     "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer._pre_moe_forward": TorchCompileOption(
         fullgraph=True
     ),
@@ -107,7 +111,7 @@
 }
 
 MOE_EP_COMPILE_CFG = MOE_NON_EP_COMPILE_CFG.copy()
-MOE_EP_COMPILE_CFG.pop("xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward")
+MOE_EP_COMPILE_CFG.pop(MOE_DECODER_LAYER_FORWARD)
 
 
 class MoEModelOutputs(ModelOutputs):
@@ -1168,7 +1172,11 @@ def fully_shard(
                 layer_idx=layer_idx,
                 mtp_idx=None,
             ):
-                layer = apply_gradient_checkpointing(layer)
+                layer = apply_selective_checkpointing(
+                    layer,
+                    self.recompute_intervals,
+                    layer_compiled_as_one_region=self._compiles_whole_decoder_layer(),
+                )
 
             self.layers[str(layer_idx)] = layer
             if layer_idx >= len(self.layers) - 1 and self.mtp_block is None:
@@ -1220,7 +1228,9 @@ def fully_shard(
                 if self._should_recompute(None, mtp_idx=mtp_idx) or (
                     self.config.mtp_config is not None and self.config.mtp_config.share_weights
                 ):  # share mtp head must recompute
-                    mtp_layer = apply_gradient_checkpointing(mtp_layer)
+                    # Marker intervals are declared against the decoder layer, so an MTP layer
+                    # gets the same mechanism with nothing kept: recomputed whole, as before.
+                    mtp_layer = apply_selective_checkpointing(mtp_layer, ())
                 self.mtp_block.layers[mtp_idx] = mtp_layer
 
                 reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1
@@ -1426,6 +1436,13 @@ def patched_emb_forward(self, input):
             self.sparse,
         )
 
+    def _compiles_whole_decoder_layer(self) -> bool:
+        # Without EP the decoder layer's own forward is compiled, so the layer is one opaque region
+        # and no marker inside it can delimit anything. With EP that entry is dropped and only the
+        # methods below it are compiled, which leaves the dispatcher-call boundaries in eager python
+        # for markers to land on.
+        return MOE_DECODER_LAYER_FORWARD in self.compile_cfg
+
     def _should_recompute(
         self,
         layer_idx: int | None,
diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py
index d645006d8..273ef0668 100644
--- a/xtuner/v1/model/utils/__init__.py
+++ b/xtuner/v1/model/utils/__init__.py
@@ -4,12 +4,14 @@
     MarkerInterval,
     RecomputeIntervalMap,
     RecomputeUnit,
+    apply_selective_checkpointing,
     checkpoint_record,
 )
 
 
 __all__ = [
     "apply_gradient_checkpointing",
+    "apply_selective_checkpointing",
     "module_dict_repr",
     "ModelForwardExtraLogInfo",
     "MarkerInterval",
diff --git a/xtuner/v1/model/utils/checkpointing.py b/xtuner/v1/model/utils/checkpointing.py
index 80aa4bcb7..88628be11 100644
--- a/xtuner/v1/model/utils/checkpointing.py
+++ b/xtuner/v1/model/utils/checkpointing.py
@@ -7,7 +7,7 @@
 from torch.utils.checkpoint import checkpoint
 
 
-__all__ = ["apply_gradient_checkpointing"]
+__all__ = ["apply_gradient_checkpointing", "CheckpointWrapper"]
 
 
 ContextFn = Callable[[], tuple[AbstractContextManager, AbstractContextManager]]
diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py
index 4768ef668..7d05a82db 100644
--- a/xtuner/v1/model/utils/selective_checkpointing.py
+++ b/xtuner/v1/model/utils/selective_checkpointing.py
@@ -1,29 +1,41 @@
-"""Shared contract for region-level selective activation checkpointing (SAC).
+"""Region-level selective activation checkpointing (SAC).
 
-This module holds the vocabulary that the three SAC layers agree on and nothing else:
+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 SAC engine turns the selected units into intervals and drives a per-op checkpoint policy.
+- 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.
 
-Only the vocabulary lives here. The contextvars session behind :func:`checkpoint_record`, the
-policy function, and the config resolution that turns user selections into intervals are owned by
-the SAC engine and the config layer respectively.
+Granularity note, because it decides which units are worth declaring for a given model: markers
+delimit regions only where the marked code runs in eager python. A ``torch.compile``d region
+executes as one unit and its markers are folded away while it is traced, so an interval whose
+endpoints fall inside a compiled region has no effect and that region is recomputed whole.
+Intervals whose endpoints fall *between* compiled regions work in both modes; in eager, all do.
 """
 
-from typing import TypeAlias
+import contextvars
+from collections.abc import Sequence
+from functools import partial
+from typing import Any, TypeAlias
 
 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 .checkpointing import CheckpointWrapper, apply_gradient_checkpointing
+
 
 __all__ = [
     "RecomputeUnit",
     "MarkerInterval",
     "RecomputeIntervalMap",
+    "apply_selective_checkpointing",
     "checkpoint_record",
 ]
 
@@ -79,6 +91,55 @@ class RecomputeUnit(StrEnum):
 """
 
 
+def apply_selective_checkpointing(
+    module: nn.Module,
+    intervals: Sequence[MarkerInterval] = (),
+    *,
+    preserve_rng_state: bool = True,
+    layer_compiled_as_one_region: bool = False,
+) -> nn.Module:
+    """Apply to ``module`` the recompute strategy it supports, keeping
+    ``intervals`` resident.
+
+    Ops executed while one of ``intervals`` is open are kept; every other op in the layer is
+    recomputed during backward. An empty ``intervals`` is the degenerate case of the same mechanism
+    -- nothing kept, everything recomputed -- so a sharding path can call this for every layer
+    ``recompute_ratio`` selects, whether or not the model declares any recompute regions.
+
+    Intervals need not be balanced: an ``end`` marker that never runs only widens the kept region.
+    Both the kept and the recomputed path are numerically exact, so marker bookkeeping can cost
+    memory but can never change gradients.
+
+    Args:
+        module (nn.Module): The layer to checkpoint.
+        intervals (Sequence[MarkerInterval]): Half-open ``[start, end)`` marker intervals to keep
+            resident, as resolved from the user's ``recompute_cfg``. Defaults to keeping nothing.
+        preserve_rng_state (bool): Restore the RNG state before recomputing, so dropout and other
+            stochastic ops replay identically. Defaults to True.
+        layer_compiled_as_one_region (bool): Whether the caller compiles this whole layer as a
+            single ``torch.compile`` region. Such a layer never runs its forward in eager python, so
+            no marker ever fires and the intervals silently keep nothing; passing True lets that be
+            reported instead of leaving the user to wonder why memory did not move. Defaults to
+            False.
+
+    Returns:
+        nn.Module: The checkpoint-wrapped layer, transparent to parameter names and ``state_dict``.
+    """
+    intervals = tuple(intervals)
+
+    if not 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")
+
+    # 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.
+    checkpointed_call = partial(_run_checkpointed_region, intervals, preserve_rng_state)
+    return CheckpointWrapper(module, checkpointed_call)
+
+
 def checkpoint_record(name: str) -> None:
     """Mark an addressable semantic boundary in the current forward pass.
 
@@ -100,8 +161,154 @@ def checkpoint_record(name: str) -> None:
     # 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
+    # `_MarkerSession.report_unreached`, which runs in eager python.
     if torch.compiler.is_compiling():
         return
 
-    # The rest is intentionally empty: the SAC engine implements the session behind this marker.
-    # Remove this stub only together with that implementation.
+    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")
+
+# Diagnostics fire once per distinct cause, not once per layer or per step.
+_REPORTED_UNREACHED_INTERVALS: set["MarkerInterval"] = set()
+_REPORTED_UNSUPPORTED_LAYERS: set[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, ...]) -> None:
+        self._intervals = intervals
+        self._open: set[MarkerInterval] = set()
+        self._recorded: set[str] = 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 report_unreached(self) -> None:
+        # An `end` that never runs is legal -- it only widens the kept region. A `start` that never
+        # runs means the interval kept nothing, and the user sees no memory change with no way to
+        # find out why. Two causes produce this and cannot be told apart from here: a marker name
+        # that this architecture does not have, and a marker that exists but sits inside a compiled
+        # region, where markers are folded away.
+        for interval in self._intervals:
+            if interval[0] in self._recorded or interval in _REPORTED_UNREACHED_INTERVALS:
+                continue
+            _REPORTED_UNREACHED_INTERVALS.add(interval)
+            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."
+            )
+
+
+def _run_checkpointed_region(
+    intervals: tuple[MarkerInterval, ...],
+    preserve_rng_state: bool,
+    module: nn.Module,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    # `context_fn` must be a module-level function or a `functools.partial` of one: Dynamo's
+    # checkpoint higher-order op rejects anything else (lambdas, closures, bound methods) with
+    # `NotImplementedError: ... LazyVariableTracker context_fn`. Keep it that way.
+    return checkpoint(
+        partial(_forward_with_marker_session, module, intervals),
+        *args,
+        use_reentrant=False,
+        preserve_rng_state=preserve_rng_state,
+        context_fn=_selective_checkpoint_contexts,
+        **kwargs,
+    )
+
+
+def _forward_with_marker_session(
+    module: nn.Module,
+    intervals: tuple[MarkerInterval, ...],
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    # Calling `module` rather than `module.forward` keeps the module's hooks inside the checkpointed
+    # region, so a hook that maintains per-forward state sees the recompute pass too.
+    #
+    # Skipped under compile for the same reason `checkpoint_record` is: Dynamo cannot trace
+    # ContextVar. Nothing is lost by it -- the markers are erased from the traced graph, so the
+    # policy would see an empty session anyway and recompute the whole region.
+    if torch.compiler.is_compiling():
+        return module(*args, **kwargs)
+
+    session = _MarkerSession(intervals)
+    token = _MARKER_SESSION.set(session)
+    try:
+        output = module(*args, **kwargs)
+    finally:
+        _MARKER_SESSION.reset(token)
+    # Only a pass that ran to the end can say a marker never ran: with `set_checkpoint_early_stop`
+    # on, the recompute pass is cut short by an exception once the last needed tensor is repacked,
+    # and reporting from there would blame markers that simply had not come up yet.
+    session.report_unreached()
+    return output
+
+
+def _selective_checkpoint_contexts() -> tuple[Any, Any]:
+    return create_selective_checkpoint_contexts(_checkpoint_policy)
+
+
+def _checkpoint_policy(ctx: Any, op: Any, *args: Any, **kwargs: Any) -> CheckpointPolicy:
+    session = _MARKER_SESSION.get()
+    if session is None or not session.keeping:
+        return CheckpointPolicy.MUST_RECOMPUTE
+
+    # Two classes of op are never kept, whatever the markers say, because keeping them is unsound
+    # rather than merely wasteful. Selective checkpointing hands the recompute pass the very tensors
+    # the forward cached, and skips the ops that produced them:
+    #
+    # - Ops with a mutable schema write into an argument, so the cached tensor is whatever the last
+    #   writer left behind. Under torch.compile inductor's `out=` extern kernels reuse buffers and
+    #   torch catches this as "Tensor cached during selective activation checkpoint has been
+    #   mutated"; in eager the same mistake is silent.
+    # - Keeping a collective elides it from the recompute pass. That is only correct while the op
+    #   that allocated its destination buffer is kept too, and an interval boundary falling between
+    #   the allocation and the collective would leave the recompute reading an uninitialised buffer
+    #   -- silently, and differently on each rank.
+    if op._schema.is_mutable or op.namespace in _NEVER_KEPT_NAMESPACES:
+        return CheckpointPolicy.MUST_RECOMPUTE
+
+    return CheckpointPolicy.MUST_SAVE
+
+
+def _warn_intervals_unsupported(module: nn.Module, intervals: tuple[MarkerInterval, ...], reason: str) -> None:
+    # Reported here rather than left to `_MarkerSession.report_unreached`, which never speaks for
+    # such a layer: its forward never runs in eager python, so no session ever opens.
+    layer = type(module).__name__
+    if layer in _REPORTED_UNSUPPORTED_LAYERS:
+        return
+    _REPORTED_UNSUPPORTED_LAYERS.add(layer)
+    log_rank0.warning(
+        f"Selective checkpointing: {layer} cannot keep recompute regions resident because {reason}, so the "
+        f"intervals {list(intervals)} have no effect and every selected layer is recomputed whole."
+    )

From eb6ba04003023bd20804935c127c595e051f1357 Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 28 Jul 2026 04:45:32 +0000
Subject: [PATCH 2/6] [Feature] Route the sharding paths through selective
 checkpointing

Every layer `recompute_ratio` selects now goes through one entry point,
which picks the recompute strategy the layer supports: region-level
selective checkpointing, or the legacy reentrant path for layers carrying
the DSA top-k lifecycle. `BaseModel.recompute_intervals` is the whole
surface the config layer needs; keeping nothing reproduces today's
whole-layer recompute.

Add the regression tests: kept regions reproduce full-recompute gradients
bit for bit and keep them alive, unbalanced and overlapping intervals are
safe, in-place writes inside a kept region are rejected rather than
silently doubled, and a kept region matches full recompute under domino EP
both eager and compiled.
---
 tests/model/test_selective_checkpointing.py   | 307 ++++++++++++++++++
 .../v1/model/utils/selective_checkpointing.py |  59 +++-
 2 files changed, 353 insertions(+), 13 deletions(-)
 create mode 100644 tests/model/test_selective_checkpointing.py

diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py
new file mode 100644
index 000000000..aec5633c6
--- /dev/null
+++ b/tests/model/test_selective_checkpointing.py
@@ -0,0 +1,307 @@
+"""Region-level selective checkpointing regression tests.
+
+TestKeptRegions
+    test_kept_region_reproduces_full_recompute: 留驻区间与全重算的输出/梯度逐位相同且梯度存活。
+    test_unbalanced_interval_is_safe: end marker 不执行时只多留驻，不影响梯度。
+    test_overlapping_intervals_keep_the_union: 区间重叠时按并集留驻。
+    test_marker_outside_session_is_noop: 不在会话内时埋点不做任何事。
+TestUnsupportedRegions
+    test_in_place_op_in_kept_region_is_rejected: 留驻区间内的 in-place 写会明确报错而不是静默改梯度。
+    test_dsa_layer_falls_back_to_reentrant: 带 DSA top-k 生命周期的层退回 legacy reentrant。
+TestRegionRecomputeUnderDominoEP
+    test_kept_region_matches_full_recompute_under_domino_ep: domino EP 下留驻区间与全重算数值一致。
+    test_kept_region_matches_full_recompute_under_compile: compile 下同上，且不触发 cached-tensor-mutated。
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.tensor import DTensor
+
+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_selective_checkpointing, checkpoint_record
+from xtuner.v1.module.attention import MHAConfig
+from xtuner.v1.module.router import NoAuxRouterConfig
+
+
+class _MarkedBlock(nn.Module):
+    """Two linear stages separated by markers, so a region can be kept or recomputed."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.first = nn.Linear(4, 4)
+        self.second = nn.Linear(4, 4)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        checkpoint_record("first.start")
+        hidden = torch.tanh(self.first(x))
+        checkpoint_record("second.start")
+        hidden = torch.tanh(self.second(hidden))
+        checkpoint_record("second.end")
+        return hidden
+
+
+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
+
+
+def _run_block(module: nn.Module, intervals) -> tuple[torch.Tensor, list[torch.Tensor]]:
+    torch.manual_seed(0)
+    for parameter in module.parameters():
+        nn.init.normal_(parameter, std=0.5)
+    wrapped = apply_selective_checkpointing(module, intervals)
+
+    inputs = torch.arange(12, dtype=torch.float32).reshape(3, 4) / 12
+    inputs.requires_grad_(True)
+    output = wrapped(inputs)
+    output.sum().backward()
+    return output.detach().clone(), [parameter.grad.clone() for parameter in module.parameters()]
+
+
+class TestKeptRegions:
+    def test_kept_region_reproduces_full_recompute(self):
+        # 留驻与重算两条路都必须是精确值，不是近似：任何差异都说明 save-list 与重算对不上。
+        recomputed_out, recomputed_grads = _run_block(_MarkedBlock(), ())
+        kept_out, kept_grads = _run_block(_MarkedBlock(), [("first.start", "second.start")])
+
+        torch.testing.assert_close(kept_out, recomputed_out, atol=0.0, rtol=0.0)
+        for kept, recomputed in zip(kept_grads, recomputed_grads):
+            torch.testing.assert_close(kept, recomputed, atol=0.0, rtol=0.0)
+
+        # 梯度断掉时 loss 依然有限、上面的比较也依然成立（两边都是 None/零），所以单独断言存活。
+        assert len(kept_grads) == 4
+        for grad in kept_grads:
+            assert grad is not None
+            assert torch.count_nonzero(grad) > 0
+
+    def test_unbalanced_interval_is_safe(self):
+        # end marker 落在没走到的分支上，只应该多留驻一段显存，绝不影响梯度。
+        recomputed_out, recomputed_grads = _run_block(_MarkedBlock(), ())
+        kept_out, kept_grads = _run_block(_MarkedBlock(), [("first.start", "never.reached")])
+
+        torch.testing.assert_close(kept_out, recomputed_out, atol=0.0, rtol=0.0)
+        for kept, recomputed in zip(kept_grads, recomputed_grads):
+            torch.testing.assert_close(kept, recomputed, atol=0.0, rtol=0.0)
+
+    def test_overlapping_intervals_keep_the_union(self):
+        # 重叠区间用「活跃集合非空即留驻」定义，不做配对断言，因此嵌套/交叉都只是并集。
+        recomputed_out, recomputed_grads = _run_block(_MarkedBlock(), ())
+        kept_out, kept_grads = _run_block(
+            _MarkedBlock(),
+            [("first.start", "second.start"), ("first.start", "second.end")],
+        )
+
+        torch.testing.assert_close(kept_out, recomputed_out, atol=0.0, rtol=0.0)
+        for kept, recomputed in zip(kept_grads, recomputed_grads):
+            torch.testing.assert_close(kept, recomputed, atol=0.0, rtol=0.0)
+
+    def test_marker_outside_session_is_noop(self):
+        # 模型可以先埋点、后开启 recompute，埋点本身不能有任何可观察行为。
+        checkpoint_record("first.start")
+
+        module = _MarkedBlock()
+        inputs = torch.zeros(3, 4, requires_grad=True)
+        module(inputs).sum().backward()
+
+        assert inputs.grad is not None
+
+
+class TestUnsupportedRegions:
+    def test_in_place_op_in_kept_region_is_rejected(self):
+        # 留驻的张量会原样交给重算那趟，read-modify-write 会在重算时二次累加：loss 有限、
+        # 没有任何报错，梯度却和重算路径不同。必须报错而不是静默算错。
+        with pytest.raises(RuntimeError, match="in-place op"):
+            _run_block(_InPlaceBlock(), [("region.start", "region.end")])
+
+    def test_in_place_op_outside_kept_region_is_fine(self):
+        out, grads = _run_block(_InPlaceBlock(), ())
+        assert torch.count_nonzero(grads[0]) > 0
+        assert torch.isfinite(out).all()
+
+    def test_dsa_layer_falls_back_to_reentrant(self):
+        # DSA 的 top-k 生命周期靠 grad 是否开启区分 checkpoint 的两趟，只有 reentrant 满足；
+        # 走到非 reentrant 上时 cache 永不释放，且不会报错。
+        module = _MarkedBlock()
+        module._dsa_topk_decoder_lifecycle_hooks_registered = True
+
+        wrapped = apply_selective_checkpointing(module, [("first.start", "second.start")])
+
+        inputs = torch.zeros(3, 4, requires_grad=True)
+        wrapped(inputs).sum().backward()
+        assert inputs.grad is not None
+
+
+def _build_moe_config(ep_size: int, dispatcher: str, compile_model: bool) -> MoEConfig:
+    router_config = NoAuxRouterConfig(
+        scoring_func="sigmoid",
+        router_scaling_factor=1.0,
+        n_group=1,
+        topk_group=1,
+        norm_topk_prob=True,
+    )
+    attention_config = MHAConfig(num_attention_heads=8, num_key_value_heads=8, head_dim=128)
+    return MoEConfig(
+        vocab_size=10240,
+        max_position_embeddings=2048,
+        pad_token_id=0,
+        eos_token_id=0,
+        num_hidden_layers=4,
+        hidden_size=1024,
+        intermediate_size=2048,
+        rms_norm_eps=1e-6,
+        rope_theta=1e6,
+        hidden_act="silu",
+        attention=attention_config,
+        tie_word_embeddings=False,
+        n_routed_experts=32,
+        n_shared_experts=1,
+        num_experts_per_tok=8,
+        first_k_dense_replace=1,
+        hidden_factor=1.0,
+        moe_intermediate_size=512,
+        router=router_config,
+        ep_size=ep_size,
+        dispatcher=dispatcher,
+        compile_cfg=None if compile_model else False,
+    )
+
+
+class _RegionMoE(MoE):
+    """A MoE whose decoder layers keep their dispatch/combine region resident.
+
+    The markers stand in for the ones the model layer records: they bracket the dispatcher calls,
+    which are the boundaries that survive ``torch.compile`` under EP.
+    """
+
+    _REGION = ("moe.dispatch", "moe.combine.end")
+
+    @property
+    def recompute_intervals(self) -> tuple[tuple[str, str], ...]:
+        return (self._REGION,)
+
+    def fully_shard(self, *args, **kwargs):
+        for layer in self.layers.values():
+            _record_markers_around_dispatch(layer)
+        return super().fully_shard(*args, **kwargs)
+
+
+def _record_markers_around_dispatch(layer: nn.Module) -> None:
+    dispatcher = getattr(layer, "dispatcher", None)
+    if dispatcher is None:  # dense layers of `first_k_dense_replace`
+        return
+
+    def mark(method_name: str, marker: str, before: bool):
+        original = getattr(dispatcher, method_name)
+
+        def marked(*args, **kwargs):
+            if before:
+                checkpoint_record(marker)
+                return original(*args, **kwargs)
+            result = original(*args, **kwargs)
+            checkpoint_record(marker)
+            return result
+
+        setattr(dispatcher, method_name, marked)
+
+    mark("dispatch", _RegionMoE._REGION[0], before=True)
+    mark("combine_postprocess", _RegionMoE._REGION[1], before=False)
+
+
+class TestRegionRecomputeUnderDominoEP(DeterministicDDPTestCase):
+    """Keeping a region must not change what the layer computes.
+
+    Under domino EP the layer interleaves micro-batches across four stage loops with async
+    dispatch/combine handles, so the forward and the recompute pass must still emit the same
+    per-overload op sequence for the save list to line up. A mismatch shows up as a wrong gradient
+    rather than an error.
+    """
+
+    @property
+    def world_size(self) -> int:
+        return int(os.getenv("XTUNER_TEST_WORLD_SIZE", "2"))
+
+    @pytest.mark.gpu
+    def test_kept_region_matches_full_recompute_under_domino_ep(self):
+        self.create_pg("cuda")
+        self._assert_region_matches_baseline(compile_model=False)
+
+    @pytest.mark.gpu
+    def test_kept_region_matches_full_recompute_under_compile(self):
+        # compile 下 policy 看到的是 inductor 的 `out=` extern kernel。把它们留驻会踩
+        # "Tensor cached during selective activation checkpoint has been mutated"，
+        # 所以这条用例同时钉住 policy 里「mutable schema 一律重算」这条规则。
+        self.create_pg("cuda")
+        self._assert_region_matches_baseline(compile_model=True)
+
+    def _assert_region_matches_baseline(self, compile_model: bool):
+        ep_size = self.world_size
+        loss_ref, grad_norm_ref, nonzero_ref = self._run_once(ep_size, compile_model, keep_region=False)
+        loss_kept, grad_norm_kept, nonzero_kept = self._run_once(ep_size, compile_model, keep_region=True)
+
+        self.assertGreater(nonzero_ref, 0)
+        self.assertEqual(nonzero_kept, nonzero_ref)
+        self.assertTrue(
+            torch.allclose(loss_kept, loss_ref, atol=5e-3, rtol=0.0),
+            f"kept-region loss {loss_kept.item()} diverged from full recompute {loss_ref.item()}",
+        )
+        rel = abs(grad_norm_kept - grad_norm_ref) / (grad_norm_ref + 1e-8)
+        self.assertLess(rel, 5e-2, f"kept-region grad-norm rel diff {rel} too large")
+
+    def _run_once(self, ep_size: int, compile_model: bool, keep_region: bool):
+        num_mb = 2
+        seq_len = 512
+        config = _build_moe_config(ep_size, "all2all", compile_model)
+        model_cls = _RegionMoE if keep_region else MoE
+        with torch.device("meta"):
+            model = model_cls(config=config)._to_device_dtype(dtype=torch.bfloat16, skip_buffers_dtype=True)
+        model.fully_shard(
+            fsdp_config=FSDPConfig(ep_size=ep_size, recompute_ratio=1.0, torch_compile=compile_model)
+        )
+
+        torch.manual_seed(42)
+        torch.cuda.manual_seed_all(42)
+        model.init_weights()
+
+        loss_cfg = CELossConfig()
+        seq_ctx_list = []
+        loss_ctx_list = []
+        gen = torch.Generator(device="cuda").manual_seed(1234)
+        for _ in range(num_mb):
+            input_ids = torch.randint(0, config.vocab_size, (1, seq_len + 1), device="cuda", generator=gen)
+            seq_ctx_list.append(SequenceContext.from_input_ids(input_ids=(input_ids[:, :-1],)))
+            loss_ctx_list.append(loss_cfg.build(data={"shifted_labels": input_ids[:, 1:]}, sp_mesh=None))
+        loss_ctx_list = loss_cfg.loss_ctx_cls.build_batches(loss_ctx_list)
+
+        out = model(seq_ctx=seq_ctx_list, loss_ctx=[{"lm": lc} for lc in loss_ctx_list])
+        loss = out["loss"]
+        loss.backward()
+
+        # FSDP hands back DTensor gradients, which do not implement every aten op; compare the local
+        # shards instead.
+        grads = [
+            p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad
+            for p in model.parameters()
+            if p.grad is not None
+        ]
+        nonzero = sum(1 for g in grads if torch.count_nonzero(g) > 0)
+        grad_norm = torch.nn.utils.get_total_norm(grads)
+        dist.barrier()
+        return loss.detach().float().cpu(), float(grad_norm), nonzero
diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py
index 7d05a82db..94e183481 100644
--- a/xtuner/v1/model/utils/selective_checkpointing.py
+++ b/xtuner/v1/model/utils/selective_checkpointing.py
@@ -180,6 +180,10 @@ def checkpoint_record(name: str) -> 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`.
+_VALUE_PRESERVING_MUTATING_OPS = frozenset({torch.ops.aten.record_stream.default})
+
 # Diagnostics fire once per distinct cause, not once per layer or per step.
 _REPORTED_UNREACHED_INTERVALS: set["MarkerInterval"] = set()
 _REPORTED_UNSUPPORTED_LAYERS: set[str] = set()
@@ -283,24 +287,53 @@ def _checkpoint_policy(ctx: Any, op: Any, *args: Any, **kwargs: Any) -> Checkpoi
     if session is None or not session.keeping:
         return CheckpointPolicy.MUST_RECOMPUTE
 
-    # Two classes of op are never kept, whatever the markers say, because keeping them is unsound
-    # rather than merely wasteful. Selective checkpointing hands the recompute pass the very tensors
-    # the forward cached, and skips the ops that produced them:
-    #
-    # - Ops with a mutable schema write into an argument, so the cached tensor is whatever the last
-    #   writer left behind. Under torch.compile inductor's `out=` extern kernels reuse buffers and
-    #   torch catches this as "Tensor cached during selective activation checkpoint has been
-    #   mutated"; in eager the same mistake is silent.
-    # - Keeping a collective elides it from the recompute pass. That is only correct while the op
-    #   that allocated its destination buffer is kept too, and an interval boundary falling between
-    #   the allocation and the collective would leave the recompute reading an uninitialised buffer
-    #   -- silently, and differently on each rank.
-    if op._schema.is_mutable or op.namespace in _NEVER_KEPT_NAMESPACES:
+    # Keeping a collective would elide it from the recompute pass. That is only correct while the op
+    # that allocated its destination buffer is kept too, and an interval boundary falling between
+    # the allocation and the collective would leave the recompute reading an uninitialised buffer --
+    # silently, and differently on each rank.
+    if op.namespace in _NEVER_KEPT_NAMESPACES:
+        return CheckpointPolicy.MUST_RECOMPUTE
+
+    if op._schema.is_mutable:
+        _reject_non_replayable_op_in_kept_region(op)
+        # Keeping a mutating op is unsound from the other side too: it writes into an argument, so
+        # what the recompute pass gets back from the cache is whatever the last writer left behind.
+        # Under torch.compile inductor's `out=` extern kernels reuse buffers and torch catches this
+        # as "Tensor cached during selective activation checkpoint has been mutated".
         return CheckpointPolicy.MUST_RECOMPUTE
 
     return CheckpointPolicy.MUST_SAVE
 
 
+def _reject_non_replayable_op_in_kept_region(op: Any) -> None:
+    # A kept region must survive being replayed on top of its own results. The recompute pass gets
+    # the *forward's* tensors back from the cache, so a read-modify-write op such as `add_` applies
+    # its update a second time to a value that already includes it: gradients then differ from the
+    # recomputed path with a finite loss and no error anywhere. Refuse rather than compute silently
+    # wrong gradients.
+    #
+    # Writing through the `out` argument is the exception, and it is what the policy sees most of
+    # the time: inductor's extern kernels are `out=` variants. Those overwrite the destination with
+    # a value that does not depend on what was there, so replaying them is idempotent. An op that
+    # writes through any other argument is rejected even when it happens to be idempotent
+    # (`copy_`), because the schema cannot tell the two apart.
+    if op in _VALUE_PRESERVING_MUTATING_OPS or _writes_only_through_out(op):
+        return
+    raise RuntimeError(
+        f"Selective checkpointing cannot keep a region containing the in-place op {op}. Move the "
+        f"`checkpoint_record` boundary so that the in-place write falls outside the kept interval, or "
+        f"rewrite it out of place."
+    )
+
+
+def _writes_only_through_out(op: Any) -> bool:
+    return all(
+        argument.name == "out"
+        for argument in op._schema.arguments
+        if argument.alias_info is not None and argument.alias_info.is_write
+    )
+
+
 def _warn_intervals_unsupported(module: nn.Module, intervals: tuple[MarkerInterval, ...], reason: str) -> None:
     # Reported here rather than left to `_MarkerSession.report_unreached`, which never speaks for
     # such a layer: its forward never runs in eager python, so no session ever opens.

From 2d22d8b42289218c9299fd2a468953c804d4b73c Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 28 Jul 2026 05:07:30 +0000
Subject: [PATCH 3/6] [Fix] Scope selective checkpointing diagnostics per model

The warnings exist so that a `recompute_cfg` which keeps nothing does not
look like a silent no-op, and they are deduplicated because every layer of
a model reaches the same diagnosis. Keyed globally on the interval or the
layer name, though, the second model in a process -- an RL reference model,
a compose model's other tower -- was silenced by the first, which is the
same silent no-op wearing a different hat. Key on the diagnosis instead.
---
 tests/model/test_selective_checkpointing.py   | 19 ++++++++++
 .../v1/model/utils/selective_checkpointing.py | 36 ++++++++++++-------
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py
index aec5633c6..1d3b09cc4 100644
--- a/tests/model/test_selective_checkpointing.py
+++ b/tests/model/test_selective_checkpointing.py
@@ -6,6 +6,7 @@
     test_overlapping_intervals_keep_the_union: 区间重叠时按并集留驻。
     test_marker_outside_session_is_noop: 不在会话内时埋点不做任何事。
 TestUnsupportedRegions
+    test_a_second_model_still_gets_its_own_diagnosis: 诊断去重不跨模型，第二个模型仍会告警。
     test_in_place_op_in_kept_region_is_rejected: 留驻区间内的 in-place 写会明确报错而不是静默改梯度。
     test_dsa_layer_falls_back_to_reentrant: 带 DSA top-k 生命周期的层退回 legacy reentrant。
 TestRegionRecomputeUnderDominoEP
@@ -28,6 +29,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 engine
 
 
 class _MarkedBlock(nn.Module):
@@ -47,6 +49,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return hidden
 
 
+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."""
 
@@ -136,6 +142,19 @@ def test_in_place_op_outside_kept_region_is_fine(self):
         assert torch.count_nonzero(grads[0]) > 0
         assert torch.isfinite(out).all()
 
+    def test_a_second_model_still_gets_its_own_diagnosis(self, monkeypatch):
+        # 诊断做去重是对的（一个模型几十层会喊几十遍），但去重不能跨模型：一个进程里
+        # 同时存在 actor / reference 或 compose 模型的两个塔时，第二个模型的告警被第一个
+        # 静默掉，用户就又回到「配了 SAVE_ATTN 却什么都没发生」的处境。
+        warnings: list[str] = []
+        monkeypatch.setattr(engine.log_rank0, "warning", warnings.append)
+
+        for module in (_MarkedBlock(), _OtherMarkedBlock()):
+            wrapped = apply_selective_checkpointing(module, [("never.reached", "second.start")])
+            wrapped(torch.zeros(3, 4, requires_grad=True)).sum().backward()
+
+        assert len(warnings) == 2, warnings
+
     def test_dsa_layer_falls_back_to_reentrant(self):
         # DSA 的 top-k 生命周期靠 grad 是否开启区分 checkpoint 的两趟，只有 reentrant 满足；
         # 走到非 reentrant 上时 cache 永不释放，且不会报错。
diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py
index 94e183481..316f16264 100644
--- a/xtuner/v1/model/utils/selective_checkpointing.py
+++ b/xtuner/v1/model/utils/selective_checkpointing.py
@@ -184,17 +184,19 @@ def checkpoint_record(name: str) -> None:
 # with a mutable schema is rejected by `_reject_in_place_op_in_kept_region`.
 _VALUE_PRESERVING_MUTATING_OPS = frozenset({torch.ops.aten.record_stream.default})
 
-# Diagnostics fire once per distinct cause, not once per layer or per step.
-_REPORTED_UNREACHED_INTERVALS: set["MarkerInterval"] = set()
-_REPORTED_UNSUPPORTED_LAYERS: set[str] = set()
+# Diagnostics fire once per distinct cause, not once per layer or per step. The cause includes
+# which layers it is about, so a second model in the same process still reports its own.
+_REPORTED_UNREACHED_INTERVALS: set[tuple[type | None, "MarkerInterval"]] = set()
+_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, ...]) -> None:
+    def __init__(self, intervals: tuple[MarkerInterval, ...], scope: type | None = None) -> None:
         self._intervals = intervals
+        self._scope = scope
         self._open: set[MarkerInterval] = set()
         self._recorded: set[str] = set()
 
@@ -219,10 +221,15 @@ def report_unreached(self) -> None:
         # find out why. Two causes produce this and cannot be told apart from here: a marker name
         # that this architecture does not have, and a marker that exists but sits inside a compiled
         # region, where markers are folded away.
+        #
+        # Deduplicated per scope rather than globally: every layer of a model reaches this with the
+        # same diagnosis, but a second model in the same process -- an RL reference model, a compose
+        # model's other tower -- is a diagnosis of its own and must not be silenced by the first.
         for interval in self._intervals:
-            if interval[0] in self._recorded or interval in _REPORTED_UNREACHED_INTERVALS:
+            reported = (self._scope, interval)
+            if interval[0] in self._recorded or reported in _REPORTED_UNREACHED_INTERVALS:
                 continue
-            _REPORTED_UNREACHED_INTERVALS.add(interval)
+            _REPORTED_UNREACHED_INTERVALS.add(reported)
             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 "
@@ -265,7 +272,7 @@ def _forward_with_marker_session(
     if torch.compiler.is_compiling():
         return module(*args, **kwargs)
 
-    session = _MarkerSession(intervals)
+    session = _MarkerSession(intervals, scope=type(module))
     token = _MARKER_SESSION.set(session)
     try:
         output = module(*args, **kwargs)
@@ -337,11 +344,16 @@ def _writes_only_through_out(op: Any) -> bool:
 def _warn_intervals_unsupported(module: nn.Module, intervals: tuple[MarkerInterval, ...], reason: str) -> None:
     # Reported here rather than left to `_MarkerSession.report_unreached`, which never speaks for
     # such a layer: its forward never runs in eager python, so no session ever opens.
-    layer = type(module).__name__
-    if layer in _REPORTED_UNSUPPORTED_LAYERS:
+    #
+    # Deduplicated on the whole diagnosis, not on the layer alone: every layer of a model produces
+    # the same one, but a second model in the same process with different intervals -- an RL
+    # reference model, a compose model's other tower -- has a diagnosis of its own to report.
+    layer = type(module)
+    diagnosis = (layer, intervals, reason)
+    if diagnosis in _REPORTED_UNSUPPORTED_DIAGNOSES:
         return
-    _REPORTED_UNSUPPORTED_LAYERS.add(layer)
+    _REPORTED_UNSUPPORTED_DIAGNOSES.add(diagnosis)
     log_rank0.warning(
-        f"Selective checkpointing: {layer} cannot keep recompute regions resident because {reason}, so the "
-        f"intervals {list(intervals)} have no effect and every selected layer is recomputed whole."
+        f"Selective checkpointing: {layer.__name__} cannot keep recompute regions resident because {reason}, so "
+        f"the intervals {list(intervals)} have no effect and every selected layer is recomputed whole."
     )

From ba674ef7a44e95203a4313eb88d5a15848b920bc Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 28 Jul 2026 05:36:44 +0000
Subject: [PATCH 4/6] [Fix] Report regions that open but keep nothing, per
 model

The compile-granularity rule was wrong: a region is addressable only if its
contents run in eager, not merely its endpoints. `SAVE_MLP` under EP falsifies
the old rule -- both markers fire in the uncompiled layer body while the region
encloses the compiled shared-expert forward, so nothing is kept. Neither
diagnostic could fire for it: the markers ran, and the layer is not compiled as
one region. That is a silent no-op of the same shape as the three this stack
already found.

The session now learns from the policy whether anything was actually kept while
an interval was open, which is a statement about contents rather than
endpoints, and reports intervals that opened and kept nothing.

Diagnostics also aggregate over the owning model instead of firing per layer.
Interval maps legitimately span dense and MoE layers, so a marker missing from
one layer type is normal, and warning per layer fired on every correctly
configured model -- which teaches users to ignore the warnings that matter.
---
 tests/model/test_selective_checkpointing.py   |  67 ++++++++-
 .../compose/intern_s1/modeling_vision.py      |   1 +
 .../model/compose/qwen3_vl/modeling_vision.py |   1 +
 xtuner/v1/model/dense/dense.py                |   1 +
 xtuner/v1/model/moe/moe.py                    |   1 +
 .../v1/model/utils/selective_checkpointing.py | 127 +++++++++++++-----
 6 files changed, 164 insertions(+), 34 deletions(-)

diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py
index 1d3b09cc4..d52057f25 100644
--- a/tests/model/test_selective_checkpointing.py
+++ b/tests/model/test_selective_checkpointing.py
@@ -29,7 +29,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 engine
+from xtuner.v1.model.utils import selective_checkpointing as contract
 
 
 class _MarkedBlock(nn.Module):
@@ -69,6 +69,40 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return accumulator
 
 
+class _EmptyRegionBlock(nn.Module):
+    """A layer whose declared region encloses no op the policy can keep.
+
+    Stands in for a region whose contents run inside a compiled kernel: both markers fire, so the
+    interval opens, yet nothing between them ever reaches the per-op policy.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(4, 4)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        hidden = self.linear(x)
+        checkpoint_record("empty.begin")
+        checkpoint_record("empty.end")
+        return hidden
+
+
+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
+
+
 def _run_block(module: nn.Module, intervals) -> tuple[torch.Tensor, list[torch.Tensor]]:
     torch.manual_seed(0)
     for parameter in module.parameters():
@@ -147,7 +181,7 @@ def test_a_second_model_still_gets_its_own_diagnosis(self, monkeypatch):
         # 同时存在 actor / reference 或 compose 模型的两个塔时，第二个模型的告警被第一个
         # 静默掉，用户就又回到「配了 SAVE_ATTN 却什么都没发生」的处境。
         warnings: list[str] = []
-        monkeypatch.setattr(engine.log_rank0, "warning", warnings.append)
+        monkeypatch.setattr(contract.log_rank0, "warning", warnings.append)
 
         for module in (_MarkedBlock(), _OtherMarkedBlock()):
             wrapped = apply_selective_checkpointing(module, [("never.reached", "second.start")])
@@ -155,6 +189,35 @@ def test_a_second_model_still_gets_its_own_diagnosis(self, monkeypatch):
 
         assert len(warnings) == 2, warnings
 
+    def test_open_interval_that_keeps_nothing_is_reported(self, monkeypatch):
+        # 这是第四次「静默无操作」：markers 都跑了，report_unreached 因此不响；区间里的 op 全在
+        # 编译区域里执行、根本到不了 policy，于是什么都没留驻，用户却收不到任何提示。
+        warnings: list[str] = []
+        monkeypatch.setattr(contract.log_rank0, "warning", warnings.append)
+
+        wrapped = apply_selective_checkpointing(_EmptyRegionBlock(), [("empty.begin", "empty.end")])
+        wrapped(torch.zeros(3, 4, requires_grad=True)).sum().backward()
+
+        assert len(warnings) == 1, warnings
+        assert "kept nothing resident" in warnings[0]
+
+    def test_interval_working_in_another_layer_is_not_reported(self, monkeypatch):
+        # 区间图会同时覆盖 dense 与 MoE 两种层（一个 MoE 模型两者都有），所以某个 marker 在
+        # 其中一种层里不出现是正常的。按层告警会让每个配置正确的模型每步都刷告警。
+        warnings: list[str] = []
+        monkeypatch.setattr(contract.log_rank0, "warning", warnings.append)
+        owner = nn.Module()
+        intervals = [("first.start", "second.start")]
+
+        layers = [
+            apply_selective_checkpointing(_EmptyRegionBlock(), intervals, owner=owner),
+            apply_selective_checkpointing(_MarkedBlock(), intervals, owner=owner),
+        ]
+        for layer in layers:
+            layer(torch.zeros(3, 4, requires_grad=True)).sum().backward()
+
+        assert warnings == []
+
     def test_dsa_layer_falls_back_to_reentrant(self):
         # DSA 的 top-k 生命周期靠 grad 是否开启区分 checkpoint 的两趟，只有 reentrant 满足；
         # 走到非 reentrant 上时 cache 永不释放，且不会报错。
diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
index 37cf3e2a2..17b718809 100644
--- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py
+++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
@@ -410,6 +410,7 @@ def fully_shard(
                 layer = apply_selective_checkpointing(
                     layer,
                     self.recompute_intervals,
+                    owner=self,
                     preserve_rng_state=checkpoint_preserve_rng_state,
                     # The layer's own forward is compiled just below, making it one opaque region.
                     layer_compiled_as_one_region=bool(self.compile_cfg),
diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
index 1fbc3a677..152be5a94 100644
--- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
+++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
@@ -341,6 +341,7 @@ def fully_shard(
                 layer = apply_selective_checkpointing(
                     layer,
                     self.recompute_intervals,
+                    owner=self,
                     preserve_rng_state=checkpoint_preserve_rng_state,
                     # The layer's own forward is compiled just below, making it one opaque region.
                     layer_compiled_as_one_region=bool(self.compile_cfg),
diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py
index f83e2df17..38286b574 100644
--- a/xtuner/v1/model/dense/dense.py
+++ b/xtuner/v1/model/dense/dense.py
@@ -238,6 +238,7 @@ def fully_shard(
                 layer = apply_selective_checkpointing(
                     layer,
                     self.recompute_intervals,
+                    owner=self,
                     preserve_rng_state=checkpoint_preserve_rng_state,
                     # The layer's own forward is compiled just below, making it one opaque region.
                     layer_compiled_as_one_region=bool(self.compile_cfg),
diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py
index d1c8eac05..1538039b3 100644
--- a/xtuner/v1/model/moe/moe.py
+++ b/xtuner/v1/model/moe/moe.py
@@ -1175,6 +1175,7 @@ def fully_shard(
                 layer = apply_selective_checkpointing(
                     layer,
                     self.recompute_intervals,
+                    owner=self,
                     layer_compiled_as_one_region=self._compiles_whole_decoder_layer(),
                 )
 
diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py
index 316f16264..1d481b845 100644
--- a/xtuner/v1/model/utils/selective_checkpointing.py
+++ b/xtuner/v1/model/utils/selective_checkpointing.py
@@ -9,14 +9,16 @@
 - 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.
 
-Granularity note, because it decides which units are worth declaring for a given model: markers
-delimit regions only where the marked code runs in eager python. A ``torch.compile``d region
-executes as one unit and its markers are folded away while it is traced, so an interval whose
-endpoints fall inside a compiled region has no effect and that region is recomputed whole.
-Intervals whose endpoints fall *between* compiled regions work in both modes; in eager, all do.
+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
+``torch.compile``d region executes as fused kernels whose ops never reach the per-op policy, so an
+interval enclosing one keeps nothing even when both its markers fire -- ``SAVE_MLP`` under EP is the
+case to remember, its markers sitting in the uncompiled layer body while the region encloses the
+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
@@ -95,6 +97,7 @@ def apply_selective_checkpointing(
     module: nn.Module,
     intervals: Sequence[MarkerInterval] = (),
     *,
+    owner: nn.Module | None = None,
     preserve_rng_state: bool = True,
     layer_compiled_as_one_region: bool = False,
 ) -> nn.Module:
@@ -114,6 +117,9 @@ def apply_selective_checkpointing(
         module (nn.Module): The layer to checkpoint.
         intervals (Sequence[MarkerInterval]): Half-open ``[start, end)`` marker intervals to keep
             resident, as resolved from the user's ``recompute_cfg``. Defaults to keeping nothing.
+        owner (nn.Module | None): The model these layers belong to. Diagnostics are aggregated over
+            it, so that an interval which keeps nothing in one layer type but works in another is
+            not reported. Defaults to None, which diagnoses each layer on its own.
         preserve_rng_state (bool): Restore the RNG state before recomputing, so dropout and other
             stochastic ops replay identically. Defaults to True.
         layer_compiled_as_one_region (bool): Whether the caller compiles this whole layer as a
@@ -136,7 +142,8 @@ def apply_selective_checkpointing(
     # 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.
-    checkpointed_call = partial(_run_checkpointed_region, intervals, preserve_rng_state)
+    _declare_selective_regions(owner, intervals)
+    checkpointed_call = partial(_run_checkpointed_region, intervals, preserve_rng_state, owner)
     return CheckpointWrapper(module, checkpointed_call)
 
 
@@ -184,9 +191,9 @@ def checkpoint_record(name: str) -> None:
 # with a mutable schema is rejected by `_reject_in_place_op_in_kept_region`.
 _VALUE_PRESERVING_MUTATING_OPS = frozenset({torch.ops.aten.record_stream.default})
 
-# Diagnostics fire once per distinct cause, not once per layer or per step. The cause includes
-# which layers it is about, so a second model in the same process still reports its own.
-_REPORTED_UNREACHED_INTERVALS: set[tuple[type | None, "MarkerInterval"]] = set()
+# 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()
 
 
@@ -194,11 +201,12 @@ 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, ...], scope: type | None = None) -> None:
+    def __init__(self, intervals: tuple[MarkerInterval, ...], owner: nn.Module | None = None) -> None:
         self._intervals = intervals
-        self._scope = scope
+        self._owner = owner
         self._open: set[MarkerInterval] = set()
         self._recorded: set[str] = set()
+        self._kept: set[MarkerInterval] = set()
 
     @property
     def keeping(self) -> bool:
@@ -215,31 +223,84 @@ def record(self, name: str) -> None:
             elif name == end:
                 self._open.discard(interval)
 
-    def report_unreached(self) -> None:
-        # An `end` that never runs is legal -- it only widens the kept region. A `start` that never
-        # runs means the interval kept nothing, and the user sees no memory change with no way to
-        # find out why. Two causes produce this and cannot be told apart from here: a marker name
-        # that this architecture does not have, and a marker that exists but sits inside a compiled
-        # region, where markers are folded away.
-        #
-        # Deduplicated per scope rather than globally: every layer of a model reaches this with the
-        # same diagnosis, but a second model in the same process -- an RL reference model, a compose
-        # model's other tower -- is a diagnosis of its own and must not be silenced by the first.
-        for interval in self._intervals:
-            reported = (self._scope, interval)
-            if interval[0] in self._recorded or reported in _REPORTED_UNREACHED_INTERVALS:
-                continue
-            _REPORTED_UNREACHED_INTERVALS.add(reported)
+    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."
+            )
 
 
 def _run_checkpointed_region(
     intervals: tuple[MarkerInterval, ...],
     preserve_rng_state: bool,
+    owner: nn.Module | None,
     module: nn.Module,
     *args: Any,
     **kwargs: Any,
@@ -248,7 +309,7 @@ def _run_checkpointed_region(
     # checkpoint higher-order op rejects anything else (lambdas, closures, bound methods) with
     # `NotImplementedError: ... LazyVariableTracker context_fn`. Keep it that way.
     return checkpoint(
-        partial(_forward_with_marker_session, module, intervals),
+        partial(_forward_with_marker_session, module, intervals, owner),
         *args,
         use_reentrant=False,
         preserve_rng_state=preserve_rng_state,
@@ -260,6 +321,7 @@ def _run_checkpointed_region(
 def _forward_with_marker_session(
     module: nn.Module,
     intervals: tuple[MarkerInterval, ...],
+    owner: nn.Module | None,
     *args: Any,
     **kwargs: Any,
 ) -> Any:
@@ -272,16 +334,16 @@ def _forward_with_marker_session(
     if torch.compiler.is_compiling():
         return module(*args, **kwargs)
 
-    session = _MarkerSession(intervals, scope=type(module))
+    session = _MarkerSession(intervals, owner)
     token = _MARKER_SESSION.set(session)
     try:
         output = module(*args, **kwargs)
     finally:
         _MARKER_SESSION.reset(token)
-    # Only a pass that ran to the end can say a marker never ran: with `set_checkpoint_early_stop`
-    # on, the recompute pass is cut short by an exception once the last needed tensor is repacked,
-    # and reporting from there would blame markers that simply had not come up yet.
-    session.report_unreached()
+    # 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.
+    session.finish()
     return output
 
 
@@ -309,6 +371,7 @@ def _checkpoint_policy(ctx: Any, op: Any, *args: Any, **kwargs: Any) -> Checkpoi
         # as "Tensor cached during selective activation checkpoint has been mutated".
         return CheckpointPolicy.MUST_RECOMPUTE
 
+    session.note_kept()
     return CheckpointPolicy.MUST_SAVE
 
 

From ff2100a0137bbb033e2eed151fe5f819bd1cd8fb Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 28 Jul 2026 09:33:32 +0000
Subject: [PATCH 5/6] [Test] Drop the DSA reentrant fallback test

The legacy reentrant path no longer exists, so a layer carrying the DSA
top-k lifecycle takes the ordinary selective checkpointing path like any
other layer. The test asserted a fallback that was removed with the path
itself.
---
 tests/model/test_selective_checkpointing.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/tests/model/test_selective_checkpointing.py b/tests/model/test_selective_checkpointing.py
index d52057f25..e3b8577f4 100644
--- a/tests/model/test_selective_checkpointing.py
+++ b/tests/model/test_selective_checkpointing.py
@@ -8,7 +8,6 @@
 TestUnsupportedRegions
     test_a_second_model_still_gets_its_own_diagnosis: 诊断去重不跨模型，第二个模型仍会告警。
     test_in_place_op_in_kept_region_is_rejected: 留驻区间内的 in-place 写会明确报错而不是静默改梯度。
-    test_dsa_layer_falls_back_to_reentrant: 带 DSA top-k 生命周期的层退回 legacy reentrant。
 TestRegionRecomputeUnderDominoEP
     test_kept_region_matches_full_recompute_under_domino_ep: domino EP 下留驻区间与全重算数值一致。
     test_kept_region_matches_full_recompute_under_compile: compile 下同上，且不触发 cached-tensor-mutated。
@@ -218,18 +217,6 @@ def test_interval_working_in_another_layer_is_not_reported(self, monkeypatch):
 
         assert warnings == []
 
-    def test_dsa_layer_falls_back_to_reentrant(self):
-        # DSA 的 top-k 生命周期靠 grad 是否开启区分 checkpoint 的两趟，只有 reentrant 满足；
-        # 走到非 reentrant 上时 cache 永不释放，且不会报错。
-        module = _MarkedBlock()
-        module._dsa_topk_decoder_lifecycle_hooks_registered = True
-
-        wrapped = apply_selective_checkpointing(module, [("first.start", "second.start")])
-
-        inputs = torch.zeros(3, 4, requires_grad=True)
-        wrapped(inputs).sum().backward()
-        assert inputs.grad is not None
-
 
 def _build_moe_config(ep_size: int, dispatcher: str, compile_model: bool) -> MoEConfig:
     router_config = NoAuxRouterConfig(

From c6f7d8713cd8a0216c417fccc274965fc1020745 Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 28 Jul 2026 10:06:00 +0000
Subject: [PATCH 6/6] [Docs] Point the diagnostics comments at the function
 that reports

The two comments named `_MarkerSession.report_unreached`, a method that never existed
under that name: unreached markers are reported from `_report_pass`, once the owner's
layers have all completed a pass.
---
 xtuner/v1/model/utils/selective_checkpointing.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/xtuner/v1/model/utils/selective_checkpointing.py b/xtuner/v1/model/utils/selective_checkpointing.py
index 1d481b845..fd2741578 100644
--- a/xtuner/v1/model/utils/selective_checkpointing.py
+++ b/xtuner/v1/model/utils/selective_checkpointing.py
@@ -171,7 +171,7 @@ def checkpoint_record(name: str) -> None:
     #
     # 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
-    # `_MarkerSession.report_unreached`, which runs in eager python.
+    # `_report_pass`, which runs in eager python once the owner's layers have all been through.
     if torch.compiler.is_compiling():
         return
 
@@ -405,8 +405,8 @@ def _writes_only_through_out(op: Any) -> bool:
 
 
 def _warn_intervals_unsupported(module: nn.Module, intervals: tuple[MarkerInterval, ...], reason: str) -> None:
-    # Reported here rather than left to `_MarkerSession.report_unreached`, which never speaks for
-    # such a layer: its forward never runs in eager python, so no session ever opens.
+    # Reported here rather than left to `_report_pass`, which never speaks for such a layer: its
+    # forward never runs in eager python, so no session ever opens and no pass ever completes.
     #
     # Deduplicated on the whole diagnosis, not on the layer alone: every layer of a model produces
     # the same one, but a second model in the same process with different intervals -- an RL