[Feature] Add the region-level selective checkpointing engine (2/3)

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 (

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,14 @@ 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,
+                    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),
+                )
                 if self.config.drop_path_rate == 0.0 and self.compile_cfg:
                     layer.forward = torch.compile(layer.forward, fullgraph=True)
 

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,14 @@ 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,
+                    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),
+                )
                 if self.compile_cfg:
                     layer.forward = torch.compile(layer.forward, fullgraph=True)
 

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,14 @@ 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,
+                    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),
+                )
 
                 # Linear-attention (GatedDeltaNet) layers write ``seq_ctx.seq_idx`` inside the
                 # checkpoint region; compiling the checkpointed layer with ``fullgraph=True`` turns

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,12 @@ def fully_shard(
                 layer_idx=layer_idx,
                 mtp_idx=None,
             ):
-                layer = apply_gradient_checkpointing(layer)
+                layer = apply_selective_checkpointing(
+                    layer,
+                    self.recompute_intervals,
+                    owner=self,
+                    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 +1229,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 +1437,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,

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",

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]]