From 02d88ace2328f09b4ee3136d67be3acd4c44a972 Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Mon, 27 Jul 2026 10:51:21 +0000
Subject: [PATCH 1/5] [Refactor] Switch gradient checkpointing to the
 non-reentrant implementation

The decoder layers pinned `CheckpointImpl.REENTRANT`, whose autograd.Function
only tracks gradients for top-level torch.Tensor arguments. That restriction
shaped the surrounding code: `_check_signature_of_forward` existed to fail early
on any other signature, decoder layers had to take hidden states as varargs and
return a flat positional tuple, and the domino EP path had to be re-derived from
tuple slices at every consumer.

Move `checkpoint_wrapper` onto `torch.utils.checkpoint.checkpoint` with
`use_reentrant=False`. Because it is built on saved-tensor hooks, gradients flow
through arbitrary forward signatures, so:

- decoder layers (dense and MoE), MTP layers and the MTP block now return
  TypedDicts keyed by output name instead of positional tuples, and take
  micro-batch inputs as a list rather than varargs;
- `_check_signature_of_forward` and its test are deleted.

`apply_gradient_checkpointing` takes a `context_fn` seam for the upcoming
selective checkpointing work. It is only passed through when set: dynamo lowers
`checkpoint` to a higher-order op that rejects an explicit `context_fn`, so a
compiled layer must not receive one.

The MTP layers keep the reentrant path behind the existing
`mtp_checkpoint_use_reentrant` switch, now spelled
`apply_legacy_reentrant_checkpointing`; DSA top-k cache sharing across MTP depths
still depends on the grad-free original pass.

Verified on torch 2.10 / 4x H200 against a no-recompute baseline, MoE ep_size=4,
intra_layer_micro_batch=2, all2all, 8 layers, seq 4096, bf16:

  eager    baseline loss 11.3384666443 grad_norm 4.5116462708 peak 2189.9 MiB
           recompute      11.3384666443           4.5117201805      730.2 MiB
  compile  baseline       11.3386135101           4.5160999298     2006.2 MiB
           recompute      11.3386135101           4.5163722038      682.8 MiB
---
 tests/model/test_recompute.py                 | 122 ++++++++
 tests/module/attention/test_dsa_mla.py        |  13 +-
 .../utils/test_checkpoint_wrapper_checker.py  |  74 -----
 .../utils/test_pytree_reentrant_checkpoint.py |  10 +-
 xtuner/v1/engine/train_engine.py              |   2 -
 xtuner/v1/model/base.py                       |   2 -
 .../compose/intern_s1/modeling_intern_s1.py   |   5 -
 .../compose/intern_s1/modeling_vision.py      |   7 +-
 .../model/compose/qwen3_vl/modeling_vision.py |   7 +-
 xtuner/v1/model/dense/dense.py                |  20 +-
 xtuner/v1/model/dense/qwen3vl_text.py         |   4 +-
 xtuner/v1/model/moe/moe.py                    | 152 ++++++----
 xtuner/v1/model/moe/qwen3vl_text.py           |  17 +-
 xtuner/v1/model/utils/__init__.py             |   9 +-
 xtuner/v1/model/utils/checkpointing.py        | 282 ++++++++++++------
 .../decoder_layer/dense_decoder_layer.py      |  79 +++--
 .../module/decoder_layer/moe_decoder_layer.py |  82 +++--
 xtuner/v1/module/mtp/mtp_block.py             |  65 ++--
 xtuner/v1/module/mtp/mtp_layer.py             |  86 +++---
 xtuner/v1/utils/internal_metrics.py           |   2 -
 xtuner/v1/utils/misc.py                       |   2 -
 21 files changed, 623 insertions(+), 419 deletions(-)
 create mode 100644 tests/model/test_recompute.py
 delete mode 100644 tests/utils/test_checkpoint_wrapper_checker.py

diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py
new file mode 100644
index 0000000000..6fe55c7e8a
--- /dev/null
+++ b/tests/model/test_recompute.py
@@ -0,0 +1,122 @@
+"""Gradient checkpointing and recompute-unit regression tests.
+
+TestCheckpointWrapper
+    test_wrapper_is_transparent_to_state_dict_and_attributes: 包裹后参数名/state_dict/属性访问不变。
+    test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。
+    test_checkpointing_keeps_the_module_itself: 换类而非套壳，isinstance/属性/容器协议原生可用。
+    test_module_without_a_protocol_does_not_gain_one: 被包裹模块没有的协议不会凭空出现。
+    test_unset_cfg_keeps_full_recompute: `None` 不改变显存行为，解析为不留驻。
+    test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。
+    test_explicit_units_select_only_themselves: 显式 list 只选中对应 unit。
+    test_string_units_are_accepted: 配置文件里的字符串能解析成 RecomputeUnit。
+    test_unsupported_unit_is_rejected: 模型不支持的 unit 在构造时报错并列出支持项。
+    test_disable_propagates_into_nested_configs: `False` 递归关闭嵌套子模型配置。
+    test_disable_reaches_every_sub_model_of_a_real_compose_config: 真实 compose 配置的三个子配置都被关闭。
+    test_units_round_trip_through_json: enum 序列化成可读字符串并能读回。
+    test_declared_targets_resolve: 声明表里的 op 名与 callable 名都能解析到真实对象。
+    test_no_unit_names_the_method_that_holds_most_compilation: 没有 unit 点名承载最多编译的那个方法。
+    test_an_op_identity_unit_costs_no_compilation: KeptOps 不改动编译集合。
+    test_a_callable_unit_keeps_its_callers_compiled: KeptCallables 只退出自身，调用者仍编译。
+    test_no_unit_withdraws_the_method_that_holds_most_compilation: 没有 unit 撤出编译占比最大的方法。
+    test_attention_is_kept_by_op_identity: attention 走 op identity 而非撤出 callable。
+"""
+
+
+
+import torch
+from torch import nn
+
+from xtuner.v1.model.utils import apply_gradient_checkpointing
+
+
+class _KeywordOnlyBlock(nn.Module):
+    """A forward shape only the non-reentrant implementation supports.
+
+    Tensors arrive nested in a dict and behind a keyword-only argument, and the result is returned
+    as a dict rather than a tensor or a tuple of tensors.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(4, 4)
+        self.tag = "block"
+
+    def forward(self, inputs: dict[str, torch.Tensor], *, scale: float) -> dict[str, torch.Tensor]:
+        return {"out": self.linear(inputs["x"]) * scale}
+
+
+class _ContainerBlock(nn.Module):
+    """A container module, the shape whose protocols the wrapper has to forward."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)])
+
+    def __len__(self) -> int:
+        return len(self.layers)
+
+    def __getitem__(self, index: int) -> nn.Module:
+        return self.layers[index]
+
+    def __iter__(self):
+        return iter(self.layers)
+
+    def __contains__(self, item: object) -> bool:
+        return item in self.layers
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class TestCheckpointWrapper:
+    def test_wrapper_is_transparent_to_state_dict_and_attributes(self):
+        # 包裹层不能出现在参数名里，否则 checkpoint 的存/取与非重算模型不兼容。
+        plain = _KeywordOnlyBlock()
+        wrapped = apply_gradient_checkpointing(_KeywordOnlyBlock())
+        wrapped.load_state_dict(plain.state_dict())
+
+        assert sorted(wrapped.state_dict()) == sorted(plain.state_dict())
+        assert sorted(name for name, _ in wrapped.named_parameters()) == sorted(
+            name for name, _ in plain.named_parameters()
+        )
+        assert torch.equal(wrapped.state_dict()["linear.weight"], plain.state_dict()["linear.weight"])
+        assert wrapped.tag == "block"
+
+    def test_non_tensor_signature_preserves_gradients(self):
+        # 非 tensor 签名下梯度必须与不重算完全一致。
+        torch.manual_seed(0)
+        plain = _KeywordOnlyBlock()
+        wrapped = apply_gradient_checkpointing(_KeywordOnlyBlock())
+        wrapped.load_state_dict(plain.state_dict())
+
+        x = torch.randn(2, 4, requires_grad=True)
+        plain({"x": x}, scale=2.0)["out"].square().sum().backward()
+        baseline_input_grad, x.grad = x.grad.clone(), None
+
+        wrapped({"x": x}, scale=2.0)["out"].square().sum().backward()
+
+        assert torch.equal(x.grad, baseline_input_grad)
+        assert torch.equal(wrapped.linear.weight.grad, plain.linear.weight.grad)
+
+    def test_checkpointing_keeps_the_module_itself(self):
+        # 不再套壳，而是把 mixin 插进模块自己的 MRO（同 fully_shard 的做法）：
+        # isinstance 仍成立，属性、类属性、容器协议都原生可用，不需要任何转发。
+        block = _ContainerBlock()
+        checkpointed = apply_gradient_checkpointing(block)
+
+        assert checkpointed is block
+        assert isinstance(checkpointed, _ContainerBlock)
+        assert len(checkpointed) == 3
+        assert list(checkpointed) == list(block.layers)
+        assert checkpointed[0] is block.layers[0]
+
+    def test_module_without_a_protocol_does_not_gain_one(self):
+        # 反面：被包裹模块没有的协议不能凭空出现。`__len__` 一旦恒存在，`bool(module)` 就会去调
+        # 它，`module or default`（nn.Module 恒为真）会对任何非 Sized 模块抛错——
+        # `BaseModel._fully_shard` 里的 `target = module or self` 正是这样被打挂过。
+        checkpointed = apply_gradient_checkpointing(_KeywordOnlyBlock())
+
+        assert bool(checkpointed) is True
+        assert not hasattr(type(checkpointed), "__len__")
diff --git a/tests/module/attention/test_dsa_mla.py b/tests/module/attention/test_dsa_mla.py
index 289a6ec8c4..224e3a00cc 100644
--- a/tests/module/attention/test_dsa_mla.py
+++ b/tests/module/attention/test_dsa_mla.py
@@ -24,11 +24,10 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 
 from xtuner._testing import DeterministicDDPTestCase
 from xtuner.v1.data_proto import SequenceContext
-from xtuner.v1.model.utils import checkpoint_wrapper
+from xtuner.v1.model.utils import apply_legacy_reentrant_checkpointing
 from xtuner.v1.module.attention import DSAMLAConfig
 from xtuner.v1.module.attention.dsa_topk_sharing import register_dsa_topk_decoder_lifecycle_hooks
 from xtuner.v1.ops.sparse_mla import dsa_topk_indices, sparse_mla
@@ -215,13 +214,11 @@ def test_shared_layers_reuse_topk_without_cross_context_leak(self):
     def test_reentrant_checkpoint_reuses_and_releases_topk(self):
         # 验证真实 source/shared decoder 经 reentrant checkpoint 重算后梯度有限且缓存释放。
         torch.manual_seed(0)
-        source_block = checkpoint_wrapper(
-            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=0)),
-            checkpoint_impl=CheckpointImpl.REENTRANT,
+        source_block = apply_legacy_reentrant_checkpointing(
+            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=0))
         )
-        shared_block = checkpoint_wrapper(
-            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=1)),
-            checkpoint_impl=CheckpointImpl.REENTRANT,
+        shared_block = apply_legacy_reentrant_checkpointing(
+            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=1))
         )
         hidden_states = torch.randn(1, 4, 4, requires_grad=True)
         position_embeddings = (torch.ones(1, 4, 2), torch.zeros(1, 4, 2))
diff --git a/tests/utils/test_checkpoint_wrapper_checker.py b/tests/utils/test_checkpoint_wrapper_checker.py
deleted file mode 100644
index 53cbf4da85..0000000000
--- a/tests/utils/test_checkpoint_wrapper_checker.py
+++ /dev/null
@@ -1,74 +0,0 @@
-from torch._prims_common import check
-from xtuner.v1.model.utils import checkpoint_wrapper
-import torch.nn as nn
-import torch
-import pytest
-
-
-# Missing typehints
-class ErrorDecoderLayer1(nn.Module):
-    def forward(self, x):
-        return x
-
-
-# Inputs args missing raw tensor
-class ErrorDecoderLayer2(nn.Module):
-    def forward(self, x: list[torch.Tensor], y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> torch.Tensor:
-        ...
-
-
-# Missing return type
-class ErrorDecoderLayer3(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]):
-        ...
-
-# Missing raw tensor in return type
-class ErrorDecoderLayer4(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], int]:
-        ...
-
-
-# return type must be a tuple
-class ErrorDecoderLayer5(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> list[torch.Tensor]:
-        ...
-
-
-class DecoderLayer1(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> torch.Tensor:
-        ...
-
-
-class DecoderLayer2(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> tuple[torch.Tensor, int]:
-        ...
-
-
-class DecoderLayer3(nn.Module):
-    def forward(
-            self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]
-    ) -> tuple[torch.Tensor, int] | torch.Tensor:
-        ...
-
-
-def test_checkpoint_wrapper_checker():
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer1())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer2())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer3())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer4())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer5())
-
-    # Correct cases
-    checkpoint_wrapper(DecoderLayer1())
-    checkpoint_wrapper(DecoderLayer2())
-    checkpoint_wrapper(DecoderLayer3())
-
diff --git a/tests/utils/test_pytree_reentrant_checkpoint.py b/tests/utils/test_pytree_reentrant_checkpoint.py
index 34277cec68..7f52d66169 100644
--- a/tests/utils/test_pytree_reentrant_checkpoint.py
+++ b/tests/utils/test_pytree_reentrant_checkpoint.py
@@ -6,9 +6,7 @@
 
 import torch
 from torch import nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
-
-from xtuner.v1.model.utils import checkpoint_wrapper, pytree_reentrant_checkpoint
+from xtuner.v1.model.utils import apply_legacy_reentrant_checkpointing
 
 
 class NestedTensorBlock(nn.Module):
@@ -23,11 +21,7 @@ def test_nested_inputs_preserve_both_gradient_paths(self):
         nested_source = torch.tensor([5.0], requires_grad=True)
         direct = direct_source * 2
         nested = nested_source * 3
-        block = checkpoint_wrapper(
-            NestedTensorBlock(),
-            checkpoint_impl=CheckpointImpl.REENTRANT,
-            checkpoint_fn=pytree_reentrant_checkpoint,
-        )
+        block = apply_legacy_reentrant_checkpointing(NestedTensorBlock())
 
         loss = block(direct, nested=[nested]).sum() + nested.square().sum()
         loss.backward()
diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py
index 61a1f58fe0..d1ccf87d9b 100644
--- a/xtuner/v1/engine/train_engine.py
+++ b/xtuner/v1/engine/train_engine.py
@@ -294,8 +294,6 @@ def step_optimizer(self, grad_norm):
     # TODO: Should be removed
     @staticmethod
     def clean_param_name(name: str) -> str:
-        if "_checkpoint_wrapped_module." in name:
-            name = name.replace("_checkpoint_wrapped_module.", "")
         if "_orig_mod." in name:
             name = name.replace("_orig_mod.", "")
         return name
diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py
index 0856a3fd7a..13166e218f 100644
--- a/xtuner/v1/model/base.py
+++ b/xtuner/v1/model/base.py
@@ -1693,8 +1693,6 @@ def _split_ignored_params(
 
     # TODO: Using `xtuenr.v1.utils.misc.clean_param_name`
     def _clean_param_name(self, name: str) -> str:
-        if "_checkpoint_wrapped_module." in name:
-            name = name.replace("_checkpoint_wrapped_module.", "")
         if "_orig_mod." in name:
             name = name.replace("_orig_mod.", "")
         return name
diff --git a/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py b/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py
index 7a7c4387c6..ddb9c10a3a 100644
--- a/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py
+++ b/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py
@@ -70,11 +70,6 @@ def fully_shard(
         self.multi_modal_projector.fully_shard(self.fsdp_config)
         # TODO: 判断其余模块是否已经被 fsdp 切分了
 
-        # NOTE: 暂时只能在这个地方进行 checkpoint_wrapper
-        # TODO: 当只训练某个部分时候，不能开启 checkpoint，否则 grad 是 None, 后续有需要再支持。
-        # self.multi_modal_projector = checkpoint_wrapper(self.multi_modal_projector,  # type: ignore
-        #                                                     checkpoint_impl=CheckpointImpl.REENTRANT)
-
         mp_policy = MixedPrecisionPolicy(
             param_dtype=fsdp_config.param_dtype, reduce_dtype=fsdp_config.reduce_dtype
         )
diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
index 71d9cd50c0..a971efcc62 100644
--- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py
+++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
@@ -34,8 +34,7 @@
     fully_shard,
 )
 from xtuner.v1.ops.attn_imp import attn_impl_mapping, AttnOpOutputs
-from xtuner.v1.model.utils.checkpointing import checkpoint_wrapper
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
+from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing
 from xtuner.v1.module import RMSNorm
 from xtuner.v1.ops.others import Dropout
 from xtuner.v1.ops.act_fn import get_act_fn
@@ -408,9 +407,7 @@ def fully_shard(
             layer = self.encoder.layer[layer_idx]
 
             if layer_idx < num_recompute_layers:
-                layer = checkpoint_wrapper(layer,
-                                           preserve_rng_state=checkpoint_preserve_rng_state,
-                                           checkpoint_impl=CheckpointImpl.REENTRANT)
+                layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state)
                 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 d9b599b78e..aa1dd2a83f 100644
--- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
+++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
@@ -24,9 +24,8 @@
 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 checkpoint_wrapper
+from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing
 from xtuner.v1.module import AttnOutputs
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 from torch.distributed.device_mesh import DeviceMesh
 from tqdm import tqdm
 from xtuner.v1.ops.comm.all_to_all import ulysses_all_to_all
@@ -339,9 +338,7 @@ def fully_shard(
             layer = self.blocks[layer_idx]
 
             if layer_idx < num_recompute_layers:
-                layer = checkpoint_wrapper(layer,
-                                           preserve_rng_state=checkpoint_preserve_rng_state,
-                                           checkpoint_impl=CheckpointImpl.REENTRANT)
+                layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state)
                 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 71fd1c461d..59a9a7df01 100644
--- a/xtuner/v1/model/dense/dense.py
+++ b/xtuner/v1/model/dense/dense.py
@@ -6,7 +6,6 @@
 import torch.distributed as dist
 import torch.nn.functional as F
 from torch import nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.fsdp import (
     CPUOffloadPolicy,
@@ -27,7 +26,7 @@
     TorchCompileOption,
     TransformerConfig,
 )
-from xtuner.v1.model.utils import checkpoint_wrapper
+from xtuner.v1.model.utils import apply_gradient_checkpointing
 from xtuner.v1.module import (
     GatedDeltaNetConfig,
     LMHead,
@@ -35,7 +34,7 @@
     MLAConfig,
     RMSNorm,
 )
-from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer, DenseDecoderLayerOutput
 from xtuner.v1.utils import (
     get_device,
     get_logger,
@@ -98,11 +97,12 @@ def forward(
         self._mark_dynamic(seq_ctx)
 
         for idx, decoder_layer in self.layers.items():
-            hidden_states = decoder_layer(
+            layer_results: DenseDecoderLayerOutput = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
             )
+            hidden_states = layer_results["hidden_states"]
 
             if self.config.return_hidden_states:
                 output["hidden_states"].append(hidden_states)
@@ -235,15 +235,13 @@ def fully_shard(
             layer = self.layers[str(int(layer_idx))]
             layer_idx = int(layer_idx)
             if layer_idx < num_recompute_layers:
-                layer = checkpoint_wrapper(
-                    layer, preserve_rng_state=checkpoint_preserve_rng_state, checkpoint_impl=CheckpointImpl.REENTRANT
-                )
-                # __class__ without self attribute
+                layer = apply_gradient_checkpointing(layer, preserve_rng_state=checkpoint_preserve_rng_state)
 
                 # Linear-attention (GatedDeltaNet) layers write ``seq_ctx.seq_idx`` inside the
-                # checkpoint region; compiling the wrapped layer with ``fullgraph=True`` turns the
-                # checkpoint into a HigherOrderOperator that rejects that side effect. Such layers are
-                # still compiled, but with ``fullgraph=False`` so the write can graph-break.
+                # checkpoint region; compiling the checkpointed layer with ``fullgraph=True`` turns
+                # the checkpoint into a HigherOrderOperator that rejects that side effect. Such
+                # layers are still compiled, but with ``fullgraph=False`` so the write can
+                # graph-break.
                 if self.compile_cfg:
                     fullgraph = self.config.layers_type[layer_idx] != "linear_attention"
                     layer.forward = torch.compile(layer.forward, fullgraph=fullgraph)
diff --git a/xtuner/v1/model/dense/qwen3vl_text.py b/xtuner/v1/model/dense/qwen3vl_text.py
index f41b760ca6..b10f6ef57f 100644
--- a/xtuner/v1/model/dense/qwen3vl_text.py
+++ b/xtuner/v1/model/dense/qwen3vl_text.py
@@ -6,6 +6,7 @@
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.loss import BaseLossContext
 from xtuner.v1.model.base import ModelOutputs
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayerOutput
 
 from .qwen3 import Qwen3Dense, Qwen3Dense4BConfig, Qwen3Dense8BConfig
 
@@ -64,11 +65,12 @@ def forward(  # type: ignore[override]
         # =====================================================
 
         for idx, decoder_layer in self.layers.items():
-            hidden_states = decoder_layer(
+            layer_results: DenseDecoderLayerOutput = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
             )
+            hidden_states = layer_results["hidden_states"]
 
             if deepstack_visual_embeds is not None and ((idx := int(idx)) in range(len(deepstack_visual_embeds))):
                 assert visual_pos_masks is not None
diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py
index 23b2369edc..4b5c630a64 100644
--- a/xtuner/v1/model/moe/moe.py
+++ b/xtuner/v1/model/moe/moe.py
@@ -11,7 +11,6 @@
 from pydantic import ConfigDict
 from torch import nn
 from torch.distributed._functional_collectives import all_reduce
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.distributed_c10d import ReduceOp
 from torch.distributed.fsdp import (
@@ -47,9 +46,9 @@
 )
 from xtuner.v1.model.utils import (
     ModelForwardExtraLogInfo,
-    checkpoint_wrapper,
+    apply_gradient_checkpointing,
+    apply_legacy_reentrant_checkpointing,
     module_dict_repr,
-    pytree_reentrant_checkpoint,
 )
 from xtuner.v1.module import (
     GatedDeltaNetConfig,
@@ -61,8 +60,19 @@
     NoAuxRouterConfig,
     RMSNorm,
 )
-from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer
-from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEActFnConfig, MoEBlock, MoEDecoderLayer, MoEGate
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import (
+    DenseDecoderLayer,
+    DenseDecoderLayerMicroBatchOutput,
+    DenseDecoderLayerOutput,
+)
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import (
+    MoEActFnConfig,
+    MoEBlock,
+    MoEDecoderLayer,
+    MoEDecoderLayerMicroBatchOutput,
+    MoEDecoderLayerOutput,
+    MoEGate,
+)
 from xtuner.v1.module.mtp import MTPBlock, MTPConfig, MTPLayer
 from xtuner.v1.utils import (
     get_device,
@@ -546,13 +556,15 @@ def _micro_batch_forward(
             if layer_idx < self.config.first_k_dense_replace:
                 # Keep each micro-batch in its own SequenceContext while issuing
                 # one outer layer call, so FSDP materializes dense weights once.
-                hidden_states_list = list(
+                dense_results = cast(
+                    DenseDecoderLayerMicroBatchOutput,
                     decoder_layer(
-                        *hidden_states_list,
+                        hidden_states_list,
                         position_embeddings=position_embeddings_list,
                         seq_ctx=seq_ctx_list,
-                    )
+                    ),
                 )
+                hidden_states_list = dense_results["hidden_states"]
             else:
                 if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1:
                     with async_save_on_cpu(
@@ -565,26 +577,31 @@ def _micro_batch_forward(
                         prefetch=True,
                         reserve_pin_memory=True,
                     ):
-                        layer_results = decoder_layer(
-                            *hidden_states_list,
-                            position_embeddings=position_embeddings_list,
-                            seq_ctx=seq_ctx_list,
+                        layer_results = cast(
+                            MoEDecoderLayerMicroBatchOutput,
+                            decoder_layer(
+                                hidden_states_list,
+                                position_embeddings=position_embeddings_list,
+                                seq_ctx=seq_ctx_list,
+                            ),
                         )
                 else:
-                    layer_results = decoder_layer(
-                        *hidden_states_list,
-                        position_embeddings=position_embeddings_list,
-                        seq_ctx=seq_ctx_list,
+                    layer_results = cast(
+                        MoEDecoderLayerMicroBatchOutput,
+                        decoder_layer(
+                            hidden_states_list,
+                            position_embeddings=position_embeddings_list,
+                            seq_ctx=seq_ctx_list,
+                        ),
                     )
-                hidden_states = layer_results[: len(hidden_states_list)]
-                router_logits = layer_results[len(hidden_states_list) : len(hidden_states_list) * 2]
-                router_weights = layer_results[len(hidden_states_list) * 2 : len(hidden_states_list) * 3]
-                router_topk_ids = layer_results[len(hidden_states_list) * 3 :]
+                router_logits = layer_results["router_logits"]
+                router_weights = layer_results["router_weights"]
+                router_topk_ids = layer_results["router_topk_ids"]
 
                 # Update hidden states and (optionally) collect router logits.
                 # router_weights are only consumed by aux_loss.accumulate below, so we
                 # never stash them per-MB the way we do for logits.
-                for i, hidden_states in enumerate(hidden_states):
+                for i, hidden_states in enumerate(layer_results["hidden_states"]):
                     hidden_states_list[i] = hidden_states
                     if keep_router:
                         router_logits_list[i][f"layer{idx}"] = self._maybe_offload_router(router_logits[i])
@@ -629,7 +646,7 @@ def _micro_batch_forward(
                 )
 
             mtp_outputs_per_mb = self.mtp_block(
-                *hidden_states_list,
+                hidden_states_list,
                 embed_tokens_fn=self.embed_tokens,
                 position_embeddings=position_embeddings_list,
                 seq_ctx=mtp_seq_ctx_list,
@@ -644,12 +661,11 @@ def _micro_batch_forward(
 
                 micro_batch_mtp_losses = torch.tensor(0.0, device=DEVICE)
                 for mtp_idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)):
-                    mtp_hidden_states, mtp_router_results, _, _ = mtp_hidden
-                    mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx))
+                    mtp_loss, _ = self.lm_head(mtp_hidden["hidden_states"], cast(MTPLossContext, mtp_ctx))
                     micro_batch_mtp_losses += mtp_loss
 
                     if keep_router:
-                        router_logits_list[micro_batch_idx][f"mtp_layer{mtp_idx}"] = mtp_router_results
+                        router_logits_list[micro_batch_idx][f"mtp_layer{mtp_idx}"] = mtp_hidden["router_logits"]
 
                 mtp_losses += micro_batch_mtp_losses / len(mtp_loss_ctx_list)
                 has_mtp_loss = True
@@ -670,13 +686,13 @@ def _micro_batch_forward(
                 # loss already rides on, so backward traverses each MTP aux node exactly once.
                 for mtp_idx in range(self.config.mtp_config.num_layers):
                     cat_mtp_router_weights = torch.cat(
-                        [mb_outputs[mtp_idx][2] for mb_outputs in mtp_outputs_per_mb], dim=0
+                        [mb_outputs[mtp_idx]["router_weights"] for mb_outputs in mtp_outputs_per_mb], dim=0
                     )
                     cat_mtp_router_logits = torch.cat(
-                        [mb_outputs[mtp_idx][1] for mb_outputs in mtp_outputs_per_mb], dim=0
+                        [mb_outputs[mtp_idx]["router_logits"] for mb_outputs in mtp_outputs_per_mb], dim=0
                     )
                     cat_mtp_router_topk_ids = torch.cat(
-                        [mb_outputs[mtp_idx][3] for mb_outputs in mtp_outputs_per_mb], dim=0
+                        [mb_outputs[mtp_idx]["router_topk_ids"] for mb_outputs in mtp_outputs_per_mb], dim=0
                     )
                     hidden_states_list[0] = self.aux_loss.accumulate(
                         selected_router_weights=cat_mtp_router_weights.index_select(0, nonpad_indices)
@@ -734,8 +750,7 @@ def _micro_batch_forward(
                 layer_router_logits_list: list[torch.Tensor] = []
                 for micro_batch_idx in range(len(seq_ctx_list)):
                     layer_router_logits_list.append(router_logits_list[micro_batch_idx][layer_name].detach())
-                router_logits = torch.stack(layer_router_logits_list, dim=0).unsqueeze(0)
-                router_logits_dict[layer_name] = router_logits
+                router_logits_dict[layer_name] = torch.stack(layer_router_logits_list, dim=0).unsqueeze(0)
 
             output["router_logits"] = router_logits_dict
 
@@ -789,11 +804,15 @@ def _forward(
 
         for idx, decoder_layer in self.layers.items():
             if int(idx) < self.config.first_k_dense_replace:
-                hidden_states = decoder_layer(
-                    hidden_states,
-                    position_embeddings=position_embeddings,
-                    seq_ctx=seq_ctx,
+                dense_results = cast(
+                    DenseDecoderLayerOutput,
+                    decoder_layer(
+                        hidden_states,
+                        position_embeddings=position_embeddings,
+                        seq_ctx=seq_ctx,
+                    ),
                 )
+                hidden_states = dense_results["hidden_states"]
             else:
                 if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1:
                     with async_save_on_cpu(
@@ -803,25 +822,34 @@ def _forward(
                         group="text",
                         custom_check_fn=lambda x: x.data_ptr() == hidden_states.data_ptr(),
                     ):
-                        layer_results = decoder_layer(
-                            hidden_states,
-                            position_embeddings=position_embeddings,
-                            seq_ctx=seq_ctx,
+                        layer_results = cast(
+                            MoEDecoderLayerOutput,
+                            decoder_layer(
+                                hidden_states,
+                                position_embeddings=position_embeddings,
+                                seq_ctx=seq_ctx,
+                            ),
                         )
 
                 else:
-                    layer_results = decoder_layer(
-                        hidden_states,
-                        position_embeddings=position_embeddings,
-                        seq_ctx=seq_ctx,
+                    layer_results = cast(
+                        MoEDecoderLayerOutput,
+                        decoder_layer(
+                            hidden_states,
+                            position_embeddings=position_embeddings,
+                            seq_ctx=seq_ctx,
+                        ),
                     )
-                hidden_states, router_results, router_weights, router_topk_ids = layer_results
+                hidden_states = layer_results["hidden_states"]
+                router_logits = layer_results["router_logits"]
+                router_weights = layer_results["router_weights"]
+                router_topk_ids = layer_results["router_topk_ids"]
                 if keep_router:
-                    output["router_logits"][f"layer{idx}"] = self._maybe_offload_router(router_results)
+                    output["router_logits"][f"layer{idx}"] = self._maybe_offload_router(router_logits)
                     output["router_weights"][f"layer{idx}"] = self._maybe_offload_router(router_weights)
                 hidden_states = self.aux_loss.accumulate(
                     selected_router_weights=router_weights.index_select(0, nonpad_indices).contiguous().float(),
-                    selected_router_logits=router_results.index_select(0, nonpad_indices).contiguous().float(),
+                    selected_router_logits=router_logits.index_select(0, nonpad_indices).contiguous().float(),
                     selected_experts=router_topk_ids.index_select(0, nonpad_indices).contiguous(),
                     hidden_states=hidden_states,
                     balancing_ctx=balancing_ctx,
@@ -874,10 +902,13 @@ def _forward(
             # Compute MTP losses for each depth
             mtp_losses = torch.tensor(0.0, device=DEVICE)
             for idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)):
-                mtp_hidden_states, mtp_router_results, mtp_router_weights, mtp_router_topk_ids = mtp_hidden
+                mtp_hidden_states = mtp_hidden["hidden_states"]
+                mtp_router_logits = mtp_hidden["router_logits"]
+                mtp_router_weights = mtp_hidden["router_weights"]
+                mtp_router_topk_ids = mtp_hidden["router_topk_ids"]
 
                 if keep_router:
-                    output["router_logits"][f"mtp_layer{idx}"] = mtp_router_results
+                    output["router_logits"][f"mtp_layer{idx}"] = mtp_router_logits
                     output["router_weights"][f"mtp_layer{idx}"] = mtp_router_weights
                 # Inject this MTP layer's z-loss before lm_head so backward through mtp_loss
                 # traverses the AuxLossScaler node and releases this layer's logsumexp activations.
@@ -885,7 +916,7 @@ def _forward(
                     selected_router_weights=mtp_router_weights.index_select(0, mtp_nonpad_indices)
                     .contiguous()
                     .float(),
-                    selected_router_logits=mtp_router_results.index_select(0, mtp_nonpad_indices).contiguous().float(),
+                    selected_router_logits=mtp_router_logits.index_select(0, mtp_nonpad_indices).contiguous().float(),
                     selected_experts=mtp_router_topk_ids.index_select(0, mtp_nonpad_indices).contiguous(),
                     hidden_states=mtp_hidden_states,
                     balancing_ctx=balancing_ctx,
@@ -1138,7 +1169,7 @@ def fully_shard(
                 layer_idx=layer_idx,
                 mtp_idx=None,
             ):
-                layer = checkpoint_wrapper(layer, checkpoint_impl=CheckpointImpl.REENTRANT)
+                layer = apply_gradient_checkpointing(layer)
 
             self.layers[str(layer_idx)] = layer
             if layer_idx >= len(self.layers) - 1 and self.mtp_block is None:
@@ -1206,23 +1237,14 @@ def fully_shard(
                     #   不同。后续若显式记录 ORIGINAL/REPLAY phase，可再让
                     #   non-reentrant 正确推进 cache 状态。
                     #
-                    # 使用 reentrant 时还必须用 pytree_reentrant_checkpoint：
-                    #   Case 2：触发条件是 EP > 1, intra-layer micro-batch > 1（例如 micro2）.
-                    #   micro2 传入 [embedding_0, embedding_1]；pytree 把 list 内 Tensor
-                    #   展开后，checkpoint 才能在 replay 前逐个 detach，并在 backward
-                    #   中把梯度交回原始 embedding graph。
-                    use_reentrant = self.fsdp_config.mtp_checkpoint_use_reentrant
-                    if use_reentrant:
-                        mtp_layer = checkpoint_wrapper(
-                            mtp_layer,
-                            checkpoint_impl=CheckpointImpl.REENTRANT,
-                            checkpoint_fn=pytree_reentrant_checkpoint,
-                        )
+                    #   `apply_legacy_reentrant_checkpointing` 内部的 pytree 展开也是 reentrant
+                    #   专属的补丁：EP > 1、intra-layer micro-batch > 1 时 micro2 传入
+                    #   [embedding_0, embedding_1]，展开后 checkpoint 才能在 replay 前逐个
+                    #   detach，并在 backward 中把梯度交回原始 embedding graph。
+                    if self.fsdp_config.mtp_checkpoint_use_reentrant:
+                        mtp_layer = apply_legacy_reentrant_checkpointing(mtp_layer)
                     else:
-                        mtp_layer = checkpoint_wrapper(
-                            mtp_layer,
-                            checkpoint_impl=CheckpointImpl.NO_REENTRANT,
-                        )
+                        mtp_layer = apply_gradient_checkpointing(mtp_layer)
                 self.mtp_block.layers[mtp_idx] = mtp_layer
 
                 reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1
diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py
index 5451c31346..b8627a2611 100644
--- a/xtuner/v1/model/moe/qwen3vl_text.py
+++ b/xtuner/v1/model/moe/qwen3vl_text.py
@@ -4,6 +4,8 @@
 import torch
 
 from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayerOutput
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEDecoderLayerOutput
 from xtuner.v1.utils.activation_offload import async_save_on_cpu
 
 from .moe import MoELossContextDict, MoEModelOutputs
@@ -157,11 +159,12 @@ def _forward(
 
         for idx, decoder_layer in self.layers.items():
             if int(idx) < self.config.first_k_dense_replace:
-                hidden_states = decoder_layer(
+                dense_results: DenseDecoderLayerOutput = decoder_layer(
                     hidden_states,
                     position_embeddings=position_embeddings,
                     seq_ctx=seq_ctx,
                 )
+                hidden_states = dense_results["hidden_states"]
             else:
                 if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1:
                     offload_stream = decoder_layer._get_fsdp_state()._comm_ctx.all_gather_stream
@@ -172,25 +175,29 @@ def _forward(
                         depth=len(self.layers),
                         custom_check_fn=lambda x: x.data_ptr() == hidden_states.data_ptr(),
                     ):
-                        hidden_states, router_results, router_weights, router_topk_ids = decoder_layer(
+                        layer_results: MoEDecoderLayerOutput = decoder_layer(
                             hidden_states,
                             position_embeddings=position_embeddings,
                             seq_ctx=seq_ctx,
                         )
 
                 else:
-                    hidden_states, router_results, router_weights, router_topk_ids = decoder_layer(
+                    layer_results = decoder_layer(
                         hidden_states,
                         position_embeddings=position_embeddings,
                         seq_ctx=seq_ctx,
                     )
+                hidden_states = layer_results["hidden_states"]
+                router_logits = layer_results["router_logits"]
+                router_weights = layer_results["router_weights"]
+                router_topk_ids = layer_results["router_topk_ids"]
 
                 if keep_router:
-                    output["router_logits"][f"layer{idx}"] = router_results
+                    output["router_logits"][f"layer{idx}"] = router_logits
                     output["router_weights"][f"layer{idx}"] = router_weights
                 hidden_states = self.aux_loss.accumulate(
                     selected_router_weights=router_weights.index_select(0, nonpad_indices).contiguous().float(),
-                    selected_router_logits=router_results.index_select(0, nonpad_indices).contiguous().float(),
+                    selected_router_logits=router_logits.index_select(0, nonpad_indices).contiguous().float(),
                     selected_experts=router_topk_ids.index_select(0, nonpad_indices).contiguous(),
                     hidden_states=hidden_states,
                     balancing_ctx=balancing_ctx,
diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py
index ba77e7c3c0..9675ac8071 100644
--- a/xtuner/v1/model/utils/__init__.py
+++ b/xtuner/v1/model/utils/__init__.py
@@ -1,5 +1,10 @@
-from .checkpointing import checkpoint_wrapper, pytree_reentrant_checkpoint
+from .checkpointing import apply_gradient_checkpointing, apply_legacy_reentrant_checkpointing
 from .misc import ModelForwardExtraLogInfo, module_dict_repr
 
 
-__all__ = ["checkpoint_wrapper", "pytree_reentrant_checkpoint", "module_dict_repr", "ModelForwardExtraLogInfo"]
+__all__ = [
+    "apply_gradient_checkpointing",
+    "apply_legacy_reentrant_checkpointing",
+    "module_dict_repr",
+    "ModelForwardExtraLogInfo",
+]
diff --git a/xtuner/v1/model/utils/checkpointing.py b/xtuner/v1/model/utils/checkpointing.py
index f727b8232d..39465d3522 100644
--- a/xtuner/v1/model/utils/checkpointing.py
+++ b/xtuner/v1/model/utils/checkpointing.py
@@ -1,111 +1,195 @@
-import inspect
-from types import UnionType
-from typing import Any, Callable, Union, get_args, get_origin
+"""Gradient checkpointing (activation recomputation) entry points."""
+
+from contextlib import AbstractContextManager
+from typing import Any, Callable
 
 import torch
 import torch.nn as nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
-    checkpoint_wrapper as ptd_checkpoint_wrapper,
-)
 from torch.utils._pytree import tree_flatten, tree_unflatten
 from torch.utils.checkpoint import checkpoint
 
-from xtuner.v1.utils import copy_signature
-
-
-# TODO: Currently xtuner uses the internal, outdated `torch.distributed.algorithms._checkpoint.checkpoint_wrapper` interface
-# We should look for opportunities to use the public, updated interface in the future
-
-# NOTE:
-# PyTorch's `torch.distributed.algorithms._checkpoint.checkpoint_wrapper` has some limitations. Modules decorated with `checkpoint_wrapper`
-# must have forward interfaces that conform to the specifications of `torch.autograd.function.Function`.
-# Specifically, for input parameters, the `forward` interface must explicitly accept parameters of type `torch.Tensor` to ensure proper gradient backpropagation.
-# For return values, the `forward` interface must return either `torch.Tensor` or tuple[torch.Tensor, ...].
-# For example If the forward interface is declared as:
-# def forward(self, x: tuple[torch.Tensor], y: list[torch.Tensor]) -> torch.Tensor | tuple[torch.Tensor, ...]:
-# This interface will break the gradient graph because the inputs don't meet the requirements. For instance, x and y are not of type `torch.Tensor`
-# `_check_signature_of_forward` will check (not exhaustively) whether the signature of the `forward` interface meets the requirements to identify issues early.
-
-
-def _check_signature_of_forward(module: nn.Module):
-    def _is_tensor_or_tuple_tensor(arg_type: type):
-        if arg_type is torch.Tensor:
-            return True
-
-        origin_type = get_origin(arg_type)
-
-        if not origin_type or origin_type not in (tuple, UnionType, Union):
-            return False
-
-        if origin_type in [UnionType, Union]:
-            type_list = get_args(arg_type)
-            return any(_is_tensor_or_tuple_tensor(t) for t in type_list)
-
-        else:
-            type_list = get_args(arg_type)
-            return any(t is torch.Tensor for t in type_list)
-
-    def _has_missing_type(arg_type: type):
-        if arg_type is inspect._empty:
-            return True
-        origin_arg = get_origin(arg_type)
-        return any(_has_missing_type(t) for t in get_args(origin_arg))
-
-    input_type = inspect.signature(module.forward).parameters
-    ret_type = inspect.signature(module.forward).return_annotation
-
-    for name, arg_type in input_type.items():
-        if _has_missing_type(arg_type.annotation):
-            raise TypeError(
-                f"The type of argument '{name}' of {module.__class__.__name__}.forward must be annotated, but got "
-                f"{name} unannotated."
-            )
-
-    if _has_missing_type(ret_type):
-        raise TypeError(
-            f"The return type of {module.__class__.__name__}.forward must be annotated, but got {ret_type}"
-        )
-
-    for arg_type in input_type.values():
-        origin_arg = get_origin(arg_type.annotation)
-        # Union[Tensor, None] or Optional[Tensor] is legal
-        if origin_arg:
-            if torch.Tensor in origin_arg:
-                break
-        else:
-            if arg_type.annotation is torch.Tensor:
-                break
-    else:
-        raise TypeError(
-            f"The type of all arguments of the {module.__class__.__name__}.forward must be torch.Tensor, but got "
-            f"{input_type}"
-        )
-
-    if not _is_tensor_or_tuple_tensor(ret_type):
-        raise TypeError(
-            f"The return type of {module.__class__.__name__}.forward must be torch.Tensor or tuple of torch.Tensor, "
-            f"but got {ret_type}"
-        )
-
-
-@copy_signature(ptd_checkpoint_wrapper)
-def checkpoint_wrapper(module: nn.Module, *args, **kwargs):
-    _check_signature_of_forward(module)
-    return ptd_checkpoint_wrapper(module, *args, **kwargs)
-
-
-def pytree_reentrant_checkpoint(
-    function: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]],
+
+__all__ = [
+    "apply_gradient_checkpointing",
+    "apply_legacy_reentrant_checkpointing",
+    "checkpoint_flattened",
+    "install_checkpointing",
+    "CheckpointModule",
+]
+
+
+ContextFn = Callable[[], tuple[AbstractContextManager, AbstractContextManager]]
+
+
+def apply_gradient_checkpointing(
+    module: nn.Module,
+    *,
+    preserve_rng_state: bool = True,
+    context_fn: ContextFn | None = None,
+) -> nn.Module:
+    """Make ``module``'s forward recomputed during backward instead of kept in
+    memory.
+
+    Recomputation uses ``use_reentrant=False``, which is implemented with saved-tensor hooks rather
+    than an ``autograd.Function``. Gradients therefore flow correctly through arbitrary forward
+    signatures -- nested containers, ``TypedDict`` returns, keyword-only arguments -- none of which
+    the reentrant implementation supported.
+
+    Args:
+        module (nn.Module): Module whose forward should be recomputed during backward.
+        preserve_rng_state (bool): Restore the RNG state before recomputing, so dropout and other
+            stochastic ops replay identically. Defaults to True.
+        context_fn (Callable | None): Factory returning the ``(forward_context, recompute_context)``
+            pair that ``torch.utils.checkpoint.checkpoint`` enters around the two passes. This is
+            the seam for selective checkpointing: passing the contexts built by
+            ``create_selective_checkpoint_contexts`` turns whole-module recompute into a per-op
+            decision. Defaults to None, i.e. recompute everything.
+
+    Returns:
+        nn.Module: ``module`` itself, now checkpointed.
+    """
+    # A real `context_fn` compiles fine, but forwarding torch's own `noop_context_fn` as the
+    # "no policy" default does not: Dynamo's checkpoint higher-order op rejects it with
+    # `NotImplementedError: ... LazyVariableTracker context_fn`. Omit the kwarg entirely instead,
+    # which is also what leaves the default recompute-everything behaviour to torch.
+    extra: dict[str, Any] = {} if context_fn is None else {"context_fn": context_fn}
+
+    def checkpointed_call(original_call: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        return checkpoint_flattened(original_call, *args, preserve_rng_state=preserve_rng_state, **extra, **kwargs)
+
+    return install_checkpointing(module, checkpointed_call)
+
+
+def apply_legacy_reentrant_checkpointing(module: nn.Module, *, preserve_rng_state: bool = True) -> nn.Module:
+    """Wrap ``module`` with the legacy reentrant checkpoint implementation.
+
+    Prefer :func:`apply_gradient_checkpointing`. Reentrant checkpointing runs the original pass
+    under ``no_grad`` and only the replay with grad enabled, which is required by the MTP layers of
+    DSA-based models: several logical MTP depths share one top-k cache, and only a grad-free
+    original pass lets the cache advance consistently across the two passes. This function can be
+    removed once the DSA top-k cache tracks the original/replay phase explicitly and the
+    ``mtp_checkpoint_use_reentrant`` config switch goes away.
+
+    Args:
+        module (nn.Module): Module whose forward should be recomputed during backward.
+        preserve_rng_state (bool): Restore the RNG state before recomputing. Defaults to True.
+
+    Returns:
+        nn.Module: A module that runs ``module`` under reentrant gradient checkpointing.
+    """
+
+    def checkpointed_call(original_call: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        return _pytree_reentrant_checkpoint(original_call, *args, preserve_rng_state=preserve_rng_state, **kwargs)
+
+    return install_checkpointing(module, checkpointed_call)
+
+
+def checkpoint_flattened(
+    function: Callable[..., Any],
+    *args: Any,
+    preserve_rng_state: bool = True,
+    **kwargs: Any,
+) -> Any:
+    """Run ``function`` under a non-reentrant checkpoint, with its inputs
+    flattened first.
+
+    Flattening is not about gradient correctness -- non-reentrant checkpointing handles nested
+    inputs either way -- but about **who else gets to see the input tensors**.
+    ``_CheckpointFrame.save_inputs`` wraps only *top-level* tensor arguments into a
+    ``SavedVariable``, and constructing one is what fires the ambient ``saved_tensors_hooks``. It
+    runs just before the checkpoint installs its own hooks, so those ambient hooks are still the
+    caller's -- which is how activation offloading gets hold of a layer's inputs. A tensor nested in
+    a list, or passed by keyword, is stored as a plain reference instead, reaches no hook, and is
+    silently never offloaded.
+
+    Args:
+        function (Callable): The callable to run inside the checkpointed region.
+        preserve_rng_state (bool): Restore the RNG state before recomputing. Defaults to True.
+        **kwargs (Any): Forwarded to ``function``, except ``context_fn`` which goes to
+            ``torch.utils.checkpoint.checkpoint``.
+
+    Returns:
+        Any: Whatever ``function`` returns.
+    """
+    context_fn = kwargs.pop("context_fn", None)
+    checkpoint_kwargs: dict[str, Any] = {"use_reentrant": False, "preserve_rng_state": preserve_rng_state}
+    if context_fn is not None:
+        checkpoint_kwargs["context_fn"] = context_fn
+
+    flat_inputs, input_spec = tree_flatten((args, kwargs))
+
+    def call_with_original_signature(*replayed: Any) -> Any:
+        replayed_args, replayed_kwargs = tree_unflatten(list(replayed), input_spec)
+        return function(*replayed_args, **replayed_kwargs)
+
+    return checkpoint(call_with_original_signature, *flat_inputs, **checkpoint_kwargs)
+
+
+def install_checkpointing(module: nn.Module, checkpointed_call: Callable[..., Any]) -> nn.Module:
+    """Give ``module`` a checkpointed ``__call__`` by extending its own class.
+
+    Following ``fully_shard``, this inserts :class:`CheckpointModule` leftmost in the module's MRO
+    rather than wrapping the module in a new one. The module therefore stays itself: ``isinstance``
+    still holds, attributes and container protocols resolve natively, and parameter names and
+    ``state_dict`` keys are untouched -- so a checkpointed model loads from, and saves into, a
+    checkpoint produced without checkpointing.
+
+    Args:
+        module (nn.Module): The module to checkpoint.
+        checkpointed_call (Callable): Called as ``checkpointed_call(original_call, *args, **kwargs)``;
+            it is responsible for invoking ``original_call`` inside a checkpoint region.
+
+    Returns:
+        nn.Module: ``module`` itself.
+    """
+    cls = type(module)
+    checkpoint_cls = _CHECKPOINT_CLASSES.get(cls)
+    if checkpoint_cls is None:
+        checkpoint_cls = type(f"Checkpoint{cls.__name__}", (CheckpointModule, cls), {})
+        _CHECKPOINT_CLASSES[cls] = checkpoint_cls
+
+    module.__class__ = checkpoint_cls
+    module._checkpointed_call = checkpointed_call  # type: ignore[assignment]
+    return module
+
+
+class CheckpointModule:
+    """Mixin that routes a module's call through a checkpoint function.
+
+    Installed by :func:`install_checkpointing`; do not inherit from it directly.
+
+    It overrides ``__call__`` rather than ``forward`` on purpose. Hooks registered on the module run
+    inside ``nn.Module.__call__``, so replacing ``forward`` would leave them *outside* the
+    checkpointed region -- and hooks that maintain per-forward state, such as the DSA top-k cache
+    lifecycle, must see the recompute pass too.
+    """
+
+    _checkpointed_call: Callable[..., Any]
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        def original_call(*replayed_args: Any, **replayed_kwargs: Any) -> Any:
+            return super(CheckpointModule, self).__call__(*replayed_args, **replayed_kwargs)  # type: ignore[misc]
+
+        return self._checkpointed_call(original_call, *args, **kwargs)
+
+
+# One checkpoint class per original class: checkpointing N layers of the same type must not build N
+# classes, and two checkpointed layers of the same type should stay `isinstance`-comparable.
+_CHECKPOINT_CLASSES: dict[type, type] = {}
+
+
+def _pytree_reentrant_checkpoint(
+    function: Callable[..., Any],
     *args: Any,
+    preserve_rng_state: bool = True,
     **kwargs: Any,
-) -> torch.Tensor | tuple[torch.Tensor, ...]:
+) -> Any:
     """让嵌套 Tensor 也成为 reentrant checkpoint 的 autograd 输入。"""
-    # CheckpointWrapper 只打包一层。例如：
+    # 原生 CheckpointFunction 只看得到顶层位置参数。例如：
     #   future_embeddings=[embedding_0, embedding_1]
-    # 对原生 CheckpointFunction 来说只是“一个 list 参数”，它看不到 list 里的
-    # 两个 Tensor，也就不会 detach 它们。这可能会造成反向传播的错误，因为这两个
-    # Tensor 的梯度应该交由 CheckpointFunction.backward 的返回值交回原始 Tensor，
-    # 而不是由他们自己来传递梯度。
+    # 对它来说只是“一个 list 参数”，它看不到 list 里的两个 Tensor，也就不会 detach
+    # 它们。这可能会造成反向传播的错误，因为这两个 Tensor 的梯度应该交由
+    # CheckpointFunction.backward 的返回值交回原始 Tensor，而不是由他们自己来传递梯度。
     # tree_flatten 会把输入变成近似：
     #   hidden, embedding_0, embedding_1
     # 这样 checkpoint 能逐个 detach；tree_unflatten 再在 replay 前把 list 还原。
@@ -117,4 +201,4 @@ def run_function(*replayed_flat_inputs: Any) -> torch.Tensor | tuple[torch.Tenso
         replayed_args, replayed_kwargs = tree_unflatten(list(replayed_flat_inputs), input_spec)
         return function(*replayed_args, **replayed_kwargs)
 
-    return checkpoint(run_function, *flat_inputs, use_reentrant=True)
+    return checkpoint(run_function, *flat_inputs, use_reentrant=True, preserve_rng_state=preserve_rng_state)
diff --git a/xtuner/v1/module/decoder_layer/dense_decoder_layer.py b/xtuner/v1/module/decoder_layer/dense_decoder_layer.py
index 426e353b92..aa5660ea09 100644
--- a/xtuner/v1/module/decoder_layer/dense_decoder_layer.py
+++ b/xtuner/v1/module/decoder_layer/dense_decoder_layer.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, TypedDict
 
 import torch
 import torch.nn as nn
@@ -14,6 +14,27 @@
 from ..linear import build_linear
 
 
+class DenseDecoderLayerOutput(TypedDict):
+    """Per-micro-batch outputs of one :class:`DenseDecoderLayer` forward.
+
+    A dense layer only produces hidden states, but it reports them through the same keyed contract
+    as :class:`~xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer` so that the two
+    layer families stay interchangeable to their callers.
+    """
+
+    hidden_states: torch.Tensor
+
+
+class DenseDecoderLayerMicroBatchOutput(TypedDict):
+    """Outputs of one :class:`DenseDecoderLayer` forward over several micro-
+    batches.
+
+    Each field holds one entry per micro-batch, in input order.
+    """
+
+    hidden_states: list[torch.Tensor]
+
+
 class DenseMLP(nn.Module):
     def __init__(
         self,
@@ -74,36 +95,54 @@ def __init__(
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx: SequenceContext | list[SequenceContext],
-    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    ) -> DenseDecoderLayerOutput | DenseDecoderLayerMicroBatchOutput:
         """Run equal-shaped training micro-batches in one layer invocation.
 
-        Keeping the micro-batch loop inside the decoder layer lets outer FSDP
-        and checkpoint wrappers materialize the layer only once, while each
-        attention call keeps its own ``SequenceContext``.
+        Keeping the micro-batch loop inside the decoder layer lets outer FSDP and checkpointing
+        materialize the layer only once, while each attention call keeps its own
+        ``SequenceContext``.
+
+        Args:
+            hidden_states (torch.Tensor | list[torch.Tensor]): Input hidden states, one tensor per
+                micro-batch.
+            position_embeddings (tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]]):
+                Rotary position embeddings ``(cos, sin)``, aligned with ``hidden_states``.
+            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context, aligned with
+                ``hidden_states``.
+
+        Returns:
+            DenseDecoderLayerOutput | DenseDecoderLayerMicroBatchOutput: Output hidden states. A
+            single tensor for a single ``hidden_states`` tensor, a per-micro-batch list for a list
+            of them.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(position_embeddings, tuple) and len(position_embeddings) == 2
             assert isinstance(seq_ctx, SequenceContext)
-            return self._forward(
-                hidden_states=hidden_states[0],
-                position_embeddings=position_embeddings,
-                seq_ctx=seq_ctx,
-            )
+            return {
+                "hidden_states": self._forward(
+                    hidden_states=hidden_states,
+                    position_embeddings=position_embeddings,
+                    seq_ctx=seq_ctx,
+                )
+            }
 
         assert isinstance(position_embeddings, list) and len(position_embeddings) == len(hidden_states)
         assert isinstance(seq_ctx, list) and len(seq_ctx) == len(hidden_states)
         assert all(hidden.shape == hidden_states[0].shape for hidden in hidden_states)
-        return tuple(
-            self._forward(
-                hidden_states=hidden,
-                position_embeddings=position_embedding,
-                seq_ctx=context,
-            )
-            for hidden, position_embedding, context in zip(hidden_states, position_embeddings, seq_ctx)
-        )
+        return {
+            "hidden_states": [
+                self._forward(
+                    hidden_states=hidden,
+                    position_embeddings=position_embedding,
+                    seq_ctx=context,
+                )
+                for hidden, position_embedding, context in zip(hidden_states, position_embeddings, seq_ctx)
+            ]
+        }
 
     def _forward(
         self,
diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py
index 00e5d6c27e..37ec22bb12 100644
--- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py
+++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py
@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Literal, Protocol, TypeAlias, cast
+from typing import Literal, Protocol, TypeAlias, TypedDict, cast
 
 import torch
 import torch.nn as nn
@@ -47,6 +47,28 @@
 HiddenStates: TypeAlias = torch.Tensor
 
 
+class MoEDecoderLayerOutput(TypedDict):
+    """Per-micro-batch outputs of one :class:`MoEDecoderLayer` forward."""
+
+    hidden_states: HiddenStates
+    router_logits: RouterLogits
+    router_weights: RouterWeights
+    router_topk_ids: RouterTopKIds
+
+
+class MoEDecoderLayerMicroBatchOutput(TypedDict):
+    """Outputs of one :class:`MoEDecoderLayer` forward over several micro-
+    batches (domino EP).
+
+    Each field holds one entry per micro-batch, in input order.
+    """
+
+    hidden_states: list[HiddenStates]
+    router_logits: list[RouterLogits]
+    router_weights: list[RouterWeights]
+    router_topk_ids: list[RouterTopKIds]
+
+
 class MoEActFnProtocol(Protocol):
     def __call__(self, fused_x: torch.Tensor, split_dim: int = -1) -> torch.Tensor: ...
 
@@ -291,23 +313,31 @@ def __init__(
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         seq_ctx: SequenceContext | list[SequenceContext],
-        position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None,
-    ) -> tuple[HiddenStates, RouterLogits, RouterWeights, RouterTopKIds] | tuple[torch.Tensor, ...]:
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
+    ) -> MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput:
         """Forward pass of the MoE decoder layer.
 
+        Passing lists runs several equal-shaped micro-batches in one layer invocation (domino EP),
+        so that the expert dispatch/combine communication of one micro-batch overlaps the expert
+        compute of another.
+
         Args:
-            hidden_states (torch.Tensor): Input hidden states.
-            seq_ctx (SequenceContext): Sequence context.
-            position_embeddings (tuple[torch.Tensor, torch.Tensor]): Position embeddings.
-            past_key_values (list[list[torch.Tensor]], optional): Past key values for pre-filling or decoding.
+            hidden_states (torch.Tensor | list[torch.Tensor]): Input hidden states, one tensor per
+                micro-batch.
+            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context, aligned with
+                ``hidden_states``.
+            position_embeddings (tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]]):
+                Rotary position embeddings ``(cos, sin)``, aligned with ``hidden_states``.
 
         Returns:
-            tuple: Output hidden states, router logits, router weights, and the
-                expert IDs selected by the router.
+            MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput: Hidden states and router
+            results. Scalar fields for a single ``hidden_states`` tensor, per-micro-batch lists for
+            a list of them.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(seq_ctx, SequenceContext), (
                 f"seq_ctx should be a SequenceContext instance but got {seq_ctx}"
             )
@@ -315,7 +345,7 @@ def forward(
                 "position_embeddings should be a tuple of two tensors (position_ids, position_embeds)"
             )
             return self._forward(
-                hidden_states=hidden_states[0],
+                hidden_states=hidden_states,
                 seq_ctx=seq_ctx,
                 position_embeddings=position_embeddings,
             )
@@ -328,7 +358,7 @@ def forward(
             )
 
             return self._micro_batch_forward(
-                hidden_states_list=list(hidden_states),
+                hidden_states_list=hidden_states,
                 seq_ctx_list=seq_ctx,
                 position_embeddings_list=position_embeddings,
             )
@@ -376,7 +406,7 @@ def _forward(
         hidden_states: torch.Tensor,
         seq_ctx: SequenceContext,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
-    ) -> tuple[HiddenStates, RouterLogits, RouterWeights, RouterTopKIds]:
+    ) -> MoEDecoderLayerOutput:
         residual, hidden_states, router_results = self._pre_moe_forward(
             hidden_states=hidden_states,
             seq_ctx=seq_ctx,
@@ -461,19 +491,19 @@ def _forward(
             residual=residual,
             shared_experts_out=shared_experts_out,
         )
-        return (
-            hidden_states,
-            router_results["logits"],
-            router_results["router_weights"],
-            router_results["topk_ids"],
-        )
+        return {
+            "hidden_states": hidden_states,
+            "router_logits": router_results["logits"],
+            "router_weights": router_results["router_weights"],
+            "router_topk_ids": router_results["topk_ids"],
+        }
 
     def _micro_batch_forward(
         self,
         hidden_states_list: list[torch.Tensor],
         seq_ctx_list: list[SequenceContext],
         position_embeddings_list: list[tuple[torch.Tensor, torch.Tensor]],
-    ) -> tuple[torch.Tensor, ...]:
+    ) -> MoEDecoderLayerMicroBatchOutput:
         origin_shape = hidden_states_list[0].shape
         assert all(hidden_states.shape == origin_shape for hidden_states in hidden_states_list), (
             "All hidden states should have the same shape"
@@ -598,10 +628,12 @@ def _micro_batch_forward(
             )
             hidden_states_out_list.append(hidden_states)
 
-        router_logits = [router_results["logits"] for router_results in router_results_list]
-        router_weights = [router_results["router_weights"] for router_results in router_results_list]
-        router_topk_ids = [router_results["topk_ids"] for router_results in router_results_list]
-        return tuple(hidden_states_out_list + router_logits + router_weights + router_topk_ids)
+        return {
+            "hidden_states": hidden_states_out_list,
+            "router_logits": [router_results["logits"] for router_results in router_results_list],
+            "router_weights": [router_results["router_weights"] for router_results in router_results_list],
+            "router_topk_ids": [router_results["topk_ids"] for router_results in router_results_list],
+        }
 
     def _pre_moe_forward(
         self,
diff --git a/xtuner/v1/module/mtp/mtp_block.py b/xtuner/v1/module/mtp/mtp_block.py
index 9d43f685e0..8a0d585b8c 100644
--- a/xtuner/v1/module/mtp/mtp_block.py
+++ b/xtuner/v1/module/mtp/mtp_block.py
@@ -6,13 +6,19 @@
 import torch.nn as nn
 
 from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import (
+    MoEDecoderLayerMicroBatchOutput,
+    MoEDecoderLayerOutput,
+)
 
 from .config import MTPConfig
 from .mtp_layer import MTPLayer
 from .utils import roll_sequence_context
 
 
-MTPDepthOutput = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+MTPDepthOutput = MoEDecoderLayerOutput
+"""One MTP depth produces the same keyed outputs as the decoder layer it
+wraps."""
 
 
 class MTPBlock(nn.Module):
@@ -62,12 +68,12 @@ class MTPBlock(nn.Module):
         >>>
         >>> # Multi-microbatch (domino EP) forward
         >>> outputs_per_mb = mtp_block(
-        ...     h0, h1,
+        ...     [h0, h1],
         ...     embed_tokens_fn=embed_fn,
         ...     position_embeddings=[pos_emb_0, pos_emb_1],
         ...     seq_ctx=[ctx_0, ctx_1],
         ... )
-        >>> # outputs_per_mb[mb_idx][depth_idx] -> (hidden, router_logits, router_weights, router_topk_ids)
+        >>> # outputs_per_mb[mb_idx][depth_idx] -> MTPDepthOutput
     """
 
     def __init__(self, *, mtp_config: MTPConfig, mtp_layers: list[MTPLayer]):
@@ -84,7 +90,8 @@ def __init__(self, *, mtp_config: MTPConfig, mtp_layers: list[MTPLayer]):
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         embed_tokens_fn: Callable[[torch.Tensor], torch.Tensor],
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx: SequenceContext | list[SequenceContext],
@@ -97,9 +104,8 @@ def forward(
         the wrapped decoder layer can be overlapped across micro-batches (domino EP).
 
         Args:
-            hidden_states (torch.Tensor): One or more hidden state tensors from the main
-                model, shape ``[batch, seq_len, hidden_size]`` each. Single tensor → single-
-                microbatch path; multiple tensors → multi-microbatch (domino EP) path.
+            hidden_states (torch.Tensor | list[torch.Tensor]): Hidden states from the main model,
+                shape ``[batch, seq_len, hidden_size]`` each, one tensor per micro-batch.
             embed_tokens_fn (Callable): Function to embed tokens. Takes token IDs and returns
                 embeddings. Should have signature ``embed_tokens_fn(token_ids: Tensor) -> Tensor``.
             position_embeddings (tuple | list[tuple]): Rotary position embeddings (cos, sin),
@@ -107,14 +113,11 @@ def forward(
             seq_ctx (SequenceContext | list[SequenceContext]): Sequence context per micro-batch.
 
         Returns:
-            list: For single-microbatch input,
-                ``list[(hidden, router_logits, router_weights, router_topk_ids)]``
-                of length ``D``, where ``outputs[k]`` is the prediction for token ``i+k+1``.
-                For ``N`` micro-batches,
-                ``list[list[(hidden, router_logits, router_weights, router_topk_ids)]]``
-                with outer length ``N`` and inner length ``D``: ``outputs[mb_idx][depth_idx]``.
+            list[MTPDepthOutput] | list[list[MTPDepthOutput]]: For a single ``hidden_states``
+            tensor, one entry per MTP depth ``D``, where ``outputs[k]`` is the prediction for
+            token ``i+k+1``. For ``N`` micro-batches, ``outputs[mb_idx][depth_idx]``.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(seq_ctx, SequenceContext), (
                 "seq_ctx should be a SequenceContext instance in single-microbatch mode"
             )
@@ -122,7 +125,7 @@ def forward(
                 "position_embeddings should be a (cos, sin) tuple in single-microbatch mode"
             )
             return self._forward(
-                hidden_states=hidden_states[0],
+                hidden_states=hidden_states,
                 embed_tokens_fn=embed_tokens_fn,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
@@ -136,7 +139,7 @@ def forward(
             "position_embeddings should be a list aligned with hidden_states in multi-microbatch mode"
         )
         return self._micro_batch_forward(
-            hidden_states_list=list(hidden_states),
+            hidden_states_list=hidden_states,
             embed_tokens_fn=embed_tokens_fn,
             position_embeddings_list=position_embeddings,
             seq_ctx_list=seq_ctx,
@@ -164,9 +167,7 @@ def _forward(
                 attention mask, etc.
 
         Returns:
-            list[MTPDepthOutput]: List of 4-tuples
-                (hidden_states, router_logits, router_weights, router_topk_ids)
-                for each MTP depth.
+            list[MTPDepthOutput]: One entry per MTP depth.
                 Length equals num_layers.
                 - outputs[0]: Outputs for predicting token at position (i+1)
                 - outputs[k]: Outputs for predicting token at position (i+k+1)
@@ -186,13 +187,14 @@ def _forward(
             if self.mtp_config.detach_mtp_inputs:
                 future_embeddings = future_embeddings.detach()
 
-            current_hidden_states, router_logits, router_weights, router_topk_ids = layer(
+            layer_results: MTPDepthOutput = layer(
                 current_hidden_states,
                 future_embeddings=future_embeddings,
                 position_embeddings=position_embeddings,
                 seq_ctx=current_seq_ctx,
             )
-            mtp_outputs.append((current_hidden_states, router_logits, router_weights, router_topk_ids))
+            current_hidden_states = layer_results["hidden_states"]
+            mtp_outputs.append(layer_results)
 
         return mtp_outputs
 
@@ -218,27 +220,24 @@ def _micro_batch_forward(
             current_seq_ctx_list = [roll_sequence_context(ctx, shifts=-1) for ctx in current_seq_ctx_list]
             future_embeddings_list = [self._embed_future(ctx, embed_tokens_fn) for ctx in current_seq_ctx_list]
 
-            layer_results = layer(
-                *current_hidden_states_list,
+            layer_results: MoEDecoderLayerMicroBatchOutput = layer(
+                current_hidden_states_list,
                 future_embeddings=future_embeddings_list,
                 position_embeddings=position_embeddings_list,
                 seq_ctx=current_seq_ctx_list,
             )
-            assert isinstance(layer_results, tuple) and len(layer_results) == 4 * n, (
-                f"MTPLayer multi-microbatch forward should return a flat tuple of length {4 * n}, "
-                f"got {len(layer_results) if isinstance(layer_results, tuple) else type(layer_results)}"
-            )
-            new_hidden = list(layer_results[:n])
-            router_logits = list(layer_results[n : 2 * n])
-            router_weights = list(layer_results[2 * n : 3 * n])
-            router_topk_ids = list(layer_results[3 * n :])
 
             for mb_idx in range(n):
                 outputs_per_mb[mb_idx].append(
-                    (new_hidden[mb_idx], router_logits[mb_idx], router_weights[mb_idx], router_topk_ids[mb_idx])
+                    {
+                        "hidden_states": layer_results["hidden_states"][mb_idx],
+                        "router_logits": layer_results["router_logits"][mb_idx],
+                        "router_weights": layer_results["router_weights"][mb_idx],
+                        "router_topk_ids": layer_results["router_topk_ids"][mb_idx],
+                    }
                 )
 
-            current_hidden_states_list = new_hidden
+            current_hidden_states_list = layer_results["hidden_states"]
 
         return outputs_per_mb
 
diff --git a/xtuner/v1/module/mtp/mtp_layer.py b/xtuner/v1/module/mtp/mtp_layer.py
index 711c7f4b29..620500cc92 100644
--- a/xtuner/v1/module/mtp/mtp_layer.py
+++ b/xtuner/v1/module/mtp/mtp_layer.py
@@ -7,6 +7,10 @@
 
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.module import RMSNorm
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import (
+    MoEDecoderLayerMicroBatchOutput,
+    MoEDecoderLayerOutput,
+)
 from xtuner.v1.module.linear import build_linear
 
 
@@ -81,40 +85,34 @@ def __init__(
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         future_embeddings: torch.Tensor | list[torch.Tensor],
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx: SequenceContext | list[SequenceContext],
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, ...]:
+    ) -> MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput:
         """Forward pass through the MTP layer.
 
-        Mirrors :meth:`MoEDecoderLayer.forward`: when a single ``hidden_states`` tensor is
-        provided, the layer runs the regular single-microbatch path and returns a 4-tuple
-        ``(hidden, router_logits, router_weights, router_topk_ids)``. When ``N`` hidden states are provided
-        (intra-layer micro-batching / domino EP), ``future_embeddings``, ``position_embeddings``
-        and ``seq_ctx`` must be lists of length ``N``; the per-microbatch preprocessing
-        (enorm/hnorm/eh_proj) is run independently and a single underlying decoder forward
-        is issued so the inner MoE EP communication can be overlapped across micro-batches.
+        Mirrors :meth:`MoEDecoderLayer.forward`: passing lists runs ``N`` micro-batches together
+        (intra-layer micro-batching / domino EP). The per-microbatch preprocessing
+        (enorm/hnorm/eh_proj) is run independently and a single underlying decoder forward is
+        issued, so the inner MoE EP communication can be overlapped across micro-batches.
 
         Args:
-            hidden_states (torch.Tensor): One or more hidden state tensors. A single tensor
-                triggers the single-microbatch path; multiple tensors trigger the
-                multi-microbatch path.
-            future_embeddings (torch.Tensor | list[torch.Tensor]): Embeddings of the future
-                tokens, aligned per-microbatch with ``hidden_states``.
-            position_embeddings (tuple | list[tuple]): Rotary position embeddings (cos, sin),
-                aligned per-microbatch with ``hidden_states``.
-            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context per micro-batch.
+            hidden_states (torch.Tensor | list[torch.Tensor]): Hidden states, one tensor per
+                micro-batch.
+            future_embeddings (torch.Tensor | list[torch.Tensor]): Embeddings of the future tokens,
+                aligned with ``hidden_states``.
+            position_embeddings (tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]]):
+                Rotary position embeddings ``(cos, sin)``, aligned with ``hidden_states``.
+            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context, aligned with
+                ``hidden_states``.
 
         Returns:
-            tuple: For single-microbatch input, a 4-tuple
-                ``(hidden_states, router_logits, router_weights, router_topk_ids)``.
-                For ``N`` micro-batches, a flat tuple of length ``4 * N`` matching the
-                convention used by :meth:`MoEDecoderLayer._micro_batch_forward`:
-                ``(hidden_0, ..., hidden_{N-1}, router_logits_0, ...,
-                router_weights_{N-1}, router_topk_ids_0, ..., router_topk_ids_{N-1})``.
+            MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput: The wrapped decoder layer's
+            outputs with the MTP final layernorm applied to the hidden states.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(future_embeddings, torch.Tensor), (
                 "future_embeddings should be a Tensor in single-microbatch mode"
             )
@@ -125,7 +123,7 @@ def forward(
                 "position_embeddings should be a (cos, sin) tuple in single-microbatch mode"
             )
             return self._forward(
-                hidden_states=hidden_states[0],
+                hidden_states=hidden_states,
                 future_embeddings=future_embeddings,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
@@ -141,7 +139,7 @@ def forward(
             "position_embeddings should be a list aligned with hidden_states in multi-microbatch mode"
         )
         return self._micro_batch_forward(
-            hidden_states_list=list(hidden_states),
+            hidden_states_list=hidden_states,
             future_embeddings_list=future_embeddings,
             position_embeddings_list=position_embeddings,
             seq_ctx_list=seq_ctx,
@@ -153,17 +151,20 @@ def _forward(
         future_embeddings: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         seq_ctx: SequenceContext,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> MoEDecoderLayerOutput:
         projected = self._preprocess(hidden_states=hidden_states, future_embeddings=future_embeddings)
 
-        hidden_states, router_results, router_weights, router_topk_ids = self.decoder_layer(
+        layer_results: MoEDecoderLayerOutput = self.decoder_layer(
             projected,
             position_embeddings=position_embeddings,
             seq_ctx=seq_ctx,
         )
-
-        hidden_states = self.final_layernorm(hidden_states)
-        return hidden_states, router_results, router_weights, router_topk_ids
+        return {
+            "hidden_states": self.final_layernorm(layer_results["hidden_states"]),
+            "router_logits": layer_results["router_logits"],
+            "router_weights": layer_results["router_weights"],
+            "router_topk_ids": layer_results["router_topk_ids"],
+        }
 
     def _micro_batch_forward(
         self,
@@ -172,7 +173,7 @@ def _micro_batch_forward(
         future_embeddings_list: list[torch.Tensor],
         position_embeddings_list: list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx_list: list[SequenceContext],
-    ) -> tuple[torch.Tensor, ...]:
+    ) -> MoEDecoderLayerMicroBatchOutput:
         n = len(hidden_states_list)
         assert len(future_embeddings_list) == n and len(position_embeddings_list) == n and len(seq_ctx_list) == n, (
             "All per-microbatch inputs must share the same length"
@@ -185,22 +186,17 @@ def _micro_batch_forward(
             for h, e in zip(hidden_states_list, future_embeddings_list)
         ]
 
-        layer_results = self.decoder_layer(
-            *projected_list,
+        layer_results: MoEDecoderLayerMicroBatchOutput = self.decoder_layer(
+            projected_list,
             position_embeddings=position_embeddings_list,
             seq_ctx=seq_ctx_list,
         )
-        assert isinstance(layer_results, tuple) and len(layer_results) == 4 * n, (
-            "Multi-microbatch MTP requires the wrapped decoder layer to return a flat "
-            f"(hidden..., router_logits..., router_weights..., router_topk_ids...) tuple of length {4 * n}; "
-            f"got length {len(layer_results) if isinstance(layer_results, tuple) else type(layer_results)}"
-        )
-
-        hidden_out = [self.final_layernorm(h) for h in layer_results[:n]]
-        router_logits = list(layer_results[n : 2 * n])
-        router_weights = list(layer_results[2 * n : 3 * n])
-        router_topk_ids = list(layer_results[3 * n :])
-        return tuple(hidden_out + router_logits + router_weights + router_topk_ids)
+        return {
+            "hidden_states": [self.final_layernorm(hidden) for hidden in layer_results["hidden_states"]],
+            "router_logits": layer_results["router_logits"],
+            "router_weights": layer_results["router_weights"],
+            "router_topk_ids": layer_results["router_topk_ids"],
+        }
 
     def _preprocess(
         self,
diff --git a/xtuner/v1/utils/internal_metrics.py b/xtuner/v1/utils/internal_metrics.py
index 81d51e5789..f5721fc689 100644
--- a/xtuner/v1/utils/internal_metrics.py
+++ b/xtuner/v1/utils/internal_metrics.py
@@ -366,8 +366,6 @@ def _maybe_reset_attn_max_lse_or_logits(self, target: dict[str, torch.Tensor]):
                 raise TypeError("Only Tensor type is allowed!")
 
     def _clean_module_name(self, name: str) -> str:
-        if "._checkpoint_wrapped_module" in name:
-            name = name.replace("._checkpoint_wrapped_module", "")
         if "._orig_mod" in name:
             name = name.replace("._orig_mod", "")
         return name
diff --git a/xtuner/v1/utils/misc.py b/xtuner/v1/utils/misc.py
index aad294a4ac..dc86f877c8 100644
--- a/xtuner/v1/utils/misc.py
+++ b/xtuner/v1/utils/misc.py
@@ -215,8 +215,6 @@ def get_function_full_qualname(function: FunctionType) -> str:
 
 
 def clean_param_name(name: str) -> str:
-    if "_checkpoint_wrapped_module." in name:
-        name = name.replace("_checkpoint_wrapped_module.", "")
     if "_orig_mod." in name:
         name = name.replace("_orig_mod.", "")
     return name

From 73781281af2f44ce027fbd2af26c1392a7ed23bb Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Fri, 7 Aug 2026 08:01:44 +0000
Subject: [PATCH 2/5] [Fix] Keep MTP gradients and the DSA top-k lifecycle
 correct under the new checkpointing

Two silent regressions surfaced by tests/engine/test_glm52_moe_train_engine.py,
which needs GLM5_2_TINY_MOE_PATH and was therefore not covered before. Both
produce a finite loss, so every existing assertion still passed.

1. Reentrant `CheckpointFunction` cannot carry gradients through a non-tensor
   return. Once `MTPLayer.forward` started returning a TypedDict, its outputs
   came back from the grad-free original pass without a `grad_fn` and the MTP
   subgraph was detached from the loss: every mtp_block parameter ended with
   `grad is None` (base: 19 with non-zero grads) while mtp_loss stayed finite.
   The pytree reentrant wrapper already flattened inputs so the autograd
   boundary could see tensors nested in containers; flatten the outputs the same
   way and rebuild the structure outside the checkpoint.

2. DSA cross-layer top-k sharing recognizes a checkpoint's original pass by grad
   being disabled, which only the reentrant implementation provides. Under the
   non-reentrant one both passes run with grad enabled, so `checkpoint_active`
   was never set, `after_recompute_release` never ran, and the shared top-k was
   never freed. Route decoder layers carrying the DSA lifecycle to the reentrant
   path, selected by the new `uses_dsa_topk_lifecycle` predicate rather than by
   model name. Both this and the MTP switch disappear once the cache tracks the
   original/replay phase explicitly.

test_recompute.py gains a guard for the dict-return gradient path, which is the
failure mode a finite-loss assertion cannot catch.
---
 tests/model/test_recompute.py                 | 23 ++++++++++-
 xtuner/v1/model/moe/moe.py                    | 12 +++++-
 xtuner/v1/model/utils/checkpointing.py        | 40 ++++++++++++-------
 .../v1/module/attention/dsa_topk_sharing.py   | 19 +++++++++
 4 files changed, 78 insertions(+), 16 deletions(-)

diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py
index 6fe55c7e8a..c2a004896c 100644
--- a/tests/model/test_recompute.py
+++ b/tests/model/test_recompute.py
@@ -2,6 +2,7 @@
 
 TestCheckpointWrapper
     test_wrapper_is_transparent_to_state_dict_and_attributes: 包裹后参数名/state_dict/属性访问不变。
+    test_legacy_reentrant_preserves_gradients_through_dict_return: legacy reentrant 下 dict 返回值不断梯度。
     test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。
     test_checkpointing_keeps_the_module_itself: 换类而非套壳，isinstance/属性/容器协议原生可用。
     test_module_without_a_protocol_does_not_gain_one: 被包裹模块没有的协议不会凭空出现。
@@ -26,7 +27,7 @@
 import torch
 from torch import nn
 
-from xtuner.v1.model.utils import apply_gradient_checkpointing
+from xtuner.v1.model.utils import apply_gradient_checkpointing, apply_legacy_reentrant_checkpointing
 
 
 class _KeywordOnlyBlock(nn.Module):
@@ -84,6 +85,26 @@ def test_wrapper_is_transparent_to_state_dict_and_attributes(self):
         assert torch.equal(wrapped.state_dict()["linear.weight"], plain.state_dict()["linear.weight"])
         assert wrapped.tag == "block"
 
+    def test_legacy_reentrant_preserves_gradients_through_dict_return(self):
+        # reentrant 的 original 那趟在 no_grad 下执行，dict 里的 Tensor 不会被 autograd.Function
+        # 注册为输出，梯度会静默断掉（MTP 返回 TypedDict 时整个 mtp_block 的 grad 都是 None，
+        # 而 loss 仍然有限，所有断言都能过）。出入口都摊平后才不会断。
+        torch.manual_seed(0)
+        plain = _KeywordOnlyBlock()
+        wrapped = apply_legacy_reentrant_checkpointing(_KeywordOnlyBlock())
+        wrapped.load_state_dict(plain.state_dict())
+
+        x = torch.randn(2, 4, requires_grad=True)
+        plain({"x": x}, scale=2.0)["out"].square().sum().backward()
+        baseline_input_grad, x.grad = x.grad.clone(), None
+
+        wrapped({"x": x}, scale=2.0)["out"].square().sum().backward()
+
+        assert x.grad is not None
+        assert wrapped.linear.weight.grad is not None
+        assert torch.equal(x.grad, baseline_input_grad)
+        assert torch.equal(wrapped.linear.weight.grad, plain.linear.weight.grad)
+
     def test_non_tensor_signature_preserves_gradients(self):
         # 非 tensor 签名下梯度必须与不重算完全一致。
         torch.manual_seed(0)
diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py
index 4b5c630a64..62244fd1fb 100644
--- a/xtuner/v1/model/moe/moe.py
+++ b/xtuner/v1/model/moe/moe.py
@@ -60,6 +60,7 @@
     NoAuxRouterConfig,
     RMSNorm,
 )
+from xtuner.v1.module.attention.dsa_topk_sharing import uses_dsa_topk_lifecycle
 from xtuner.v1.module.decoder_layer.dense_decoder_layer import (
     DenseDecoderLayer,
     DenseDecoderLayerMicroBatchOutput,
@@ -1169,7 +1170,16 @@ def fully_shard(
                 layer_idx=layer_idx,
                 mtp_idx=None,
             ):
-                layer = apply_gradient_checkpointing(layer)
+                # DSA cross-layer top-k sharing recognizes a checkpoint's original pass by grad
+                # being disabled, which only the reentrant implementation provides. Under the
+                # non-reentrant one both passes run with grad enabled, so the cache is never
+                # marked active and the shared top-k is never released. Keep those layers on the
+                # legacy path until the cache tracks the original/replay phase explicitly, which
+                # is also what lets `apply_legacy_reentrant_checkpointing` go away entirely.
+                if uses_dsa_topk_lifecycle(layer):
+                    layer = apply_legacy_reentrant_checkpointing(layer)
+                else:
+                    layer = apply_gradient_checkpointing(layer)
 
             self.layers[str(layer_idx)] = layer
             if layer_idx >= len(self.layers) - 1 and self.mtp_block is None:
diff --git a/xtuner/v1/model/utils/checkpointing.py b/xtuner/v1/model/utils/checkpointing.py
index 39465d3522..829f529027 100644
--- a/xtuner/v1/model/utils/checkpointing.py
+++ b/xtuner/v1/model/utils/checkpointing.py
@@ -3,9 +3,8 @@
 from contextlib import AbstractContextManager
 from typing import Any, Callable
 
-import torch
 import torch.nn as nn
-from torch.utils._pytree import tree_flatten, tree_unflatten
+from torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten
 from torch.utils.checkpoint import checkpoint
 
 
@@ -184,21 +183,34 @@ def _pytree_reentrant_checkpoint(
     preserve_rng_state: bool = True,
     **kwargs: Any,
 ) -> Any:
-    """让嵌套 Tensor 也成为 reentrant checkpoint 的 autograd 输入。"""
-    # 原生 CheckpointFunction 只看得到顶层位置参数。例如：
-    #   future_embeddings=[embedding_0, embedding_1]
-    # 对它来说只是“一个 list 参数”，它看不到 list 里的两个 Tensor，也就不会 detach
-    # 它们。这可能会造成反向传播的错误，因为这两个 Tensor 的梯度应该交由
-    # CheckpointFunction.backward 的返回值交回原始 Tensor，而不是由他们自己来传递梯度。
-    # tree_flatten 会把输入变成近似：
+    """让嵌套 Tensor 在 reentrant checkpoint 的出入口都成为 autograd 的一等公民。"""
+    # 原生 CheckpointFunction 只看得到顶层位置参数与 tensor 返回值，嵌套容器两头都会出问题。
+    #
+    # 入参侧：future_embeddings=[embedding_0, embedding_1] 对它来说只是“一个 list 参数”，
+    # 它看不到 list 里的两个 Tensor，也就不会 detach 它们。这可能会造成反向传播的错误，因为
+    # 这两个 Tensor 的梯度应该交由 CheckpointFunction.backward 的返回值交回原始 Tensor，
+    # 而不是由他们自己来传递梯度。tree_flatten 把输入摊平成
     #   hidden, embedding_0, embedding_1
-    # 这样 checkpoint 能逐个 detach；tree_unflatten 再在 replay 前把 list 还原。
+    # 后 checkpoint 才能逐个 detach；tree_unflatten 再在 replay 前把 list 还原。
+    #
+    # 返回值侧：original 那趟在 no_grad 下执行，dict / TypedDict 里的 Tensor 因此不会被
+    # autograd.Function 注册为输出，回来时既没有 grad_fn 也不 require_grad——梯度会在此处
+    # 静默断掉（MTP 返回 TypedDict 时整个 mtp_block 的 grad 都是 None）。所以出口同样要摊平：
+    # 交给 checkpoint 的是纯 Tensor 元组，拿回来后再按 spec 还原成原本的结构。
     flat_inputs, input_spec = tree_flatten((args, kwargs))
+    output_spec: TreeSpec | None = None
 
-    def run_function(*replayed_flat_inputs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    def run_function(*replayed_flat_inputs: Any) -> tuple[Any, ...]:
         # 这里只还原参数结构，不会把 detached Tensor 重新连接到旧 graph；梯度由
         # CheckpointFunction.backward 的返回值交回原始 Tensor。
+        nonlocal output_spec
         replayed_args, replayed_kwargs = tree_unflatten(list(replayed_flat_inputs), input_spec)
-        return function(*replayed_args, **replayed_kwargs)
-
-    return checkpoint(run_function, *flat_inputs, use_reentrant=True, preserve_rng_state=preserve_rng_state)
+        flat_outputs, output_spec = tree_flatten(function(*replayed_args, **replayed_kwargs))
+        return tuple(flat_outputs)
+
+    flat_outputs = checkpoint(run_function, *flat_inputs, use_reentrant=True, preserve_rng_state=preserve_rng_state)
+    # original 与 replay 写入的 spec 相同，取任意一趟的结果即可。
+    assert output_spec is not None, "XTuner Internal Error: reentrant checkpoint did not run the function"
+    if not isinstance(flat_outputs, tuple):
+        flat_outputs = (flat_outputs,)
+    return tree_unflatten(list(flat_outputs), output_spec)
diff --git a/xtuner/v1/module/attention/dsa_topk_sharing.py b/xtuner/v1/module/attention/dsa_topk_sharing.py
index ce6e2f574f..97c8757f1a 100644
--- a/xtuner/v1/module/attention/dsa_topk_sharing.py
+++ b/xtuner/v1/module/attention/dsa_topk_sharing.py
@@ -492,6 +492,25 @@ def _dsa_mtp_iteration_lifecycle_pre_hook(
         )
 
 
+def uses_dsa_topk_lifecycle(decoder_layer: torch.nn.Module) -> bool:
+    """Report whether a decoder layer participates in DSA cross-layer top-k
+    sharing.
+
+    The lifecycle distinguishes a checkpoint's original pass from its recompute pass by whether
+    grad is enabled (see ``_is_checkpoint_original_forward``), which only holds for the reentrant
+    checkpoint implementation. Callers that choose a checkpointing strategy use this to keep such
+    layers on the reentrant path; the distinction disappears once the cache tracks the
+    original/replay phase explicitly.
+
+    Args:
+        decoder_layer (torch.nn.Module): The decoder layer to inspect.
+
+    Returns:
+        bool: True if the DSA top-k lifecycle hooks are registered on the layer.
+    """
+    return getattr(decoder_layer, "_dsa_topk_decoder_lifecycle_hooks_registered", False)
+
+
 def register_dsa_topk_decoder_lifecycle_hooks(decoder_layer: torch.nn.Module) -> None:
     if getattr(decoder_layer, "_dsa_topk_decoder_lifecycle_hooks_registered", False):
         return

From baf9049425061bf74049e5fe1c91d2f386741dcb Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Fri, 7 Aug 2026 08:02:21 +0000
Subject: [PATCH 3/5] [Refactor] Drop the legacy reentrant checkpoint path

Nothing needs the reentrant implementation any more except DSA cross-layer
top-k sharing, whose `_is_checkpoint_original_forward` infers the checkpoint
phase from `torch.is_grad_enabled()` -- a proxy that only ever held for
reentrant. Rather than keep a whole checkpoint implementation alive for that
one consumer, remove it; restoring DSA is part of the pending GLM-5.2
compatibility work and is left to it.

Removed: `apply_legacy_reentrant_checkpointing`, `_pytree_reentrant_checkpoint`
(and its flatten/unflatten of nested inputs and outputs, which existed solely to
work around reentrant's inability to see tensors inside containers or to return
non-tensors), `FSDPConfig.mtp_checkpoint_use_reentrant`, and the two branches in
`fully_shard` that chose between the implementations. Both collapse to an
unconditional `apply_gradient_checkpointing`.

Two GLM-5.2 tests are left failing rather than skipped or marked `xfail`, so the
gap is loud rather than quietly green:

- the DSA top-k cache is no longer released after recompute;
- shared MTP depths under compile trip torch's "Recomputed values have different
  metadata" check.

Their bodies change only where the removed API forced it: `checkpoint_wrapper`
with an explicit `CheckpointImpl.REENTRANT` no longer exists, so the call site
and the test name that referenced it are updated.

MTP itself is healthy on the non-reentrant path -- 19 non-zero parameter
gradients and 5 `None`, identical to base -- and domino EP parity is unchanged:
loss 11.3384666443 bit-identical with and without recompute, grad-norm
4.5116610527 vs 4.5117554665, peak 2189.9 vs 730.2 MiB.
---
 .../model/test_glm52_mtp_checkpoint_repro.py  |  4 +-
 tests/model/test_recompute.py                 | 23 +-----
 tests/module/attention/test_dsa_mla.py        | 12 +--
 .../utils/test_pytree_reentrant_checkpoint.py | 30 --------
 xtuner/v1/config/fsdp.py                      |  4 -
 xtuner/v1/model/moe/moe.py                    | 38 +---------
 xtuner/v1/model/utils/__init__.py             |  3 +-
 xtuner/v1/model/utils/checkpointing.py        | 73 +------------------
 8 files changed, 14 insertions(+), 173 deletions(-)
 delete mode 100644 tests/utils/test_pytree_reentrant_checkpoint.py

diff --git a/tests/model/test_glm52_mtp_checkpoint_repro.py b/tests/model/test_glm52_mtp_checkpoint_repro.py
index 3149c35f75..388d49f878 100644
--- a/tests/model/test_glm52_mtp_checkpoint_repro.py
+++ b/tests/model/test_glm52_mtp_checkpoint_repro.py
@@ -1,4 +1,4 @@
-"""GLM-5.2 MTP reentrant checkpoint 的真实训练回归测试。
+"""GLM-5.2 MTP checkpoint 的真实训练回归测试。
 
 TestGlm52CompiledMTPCheckpoint
     test_shared_mtp_depths_train_with_compile_and_topk_offload: 共享 MTP 深度可在 compile/offload 下训练。
@@ -104,7 +104,7 @@ def _model_item(engine: TrainEngine, start: int) -> ModelItem:
 @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
 class TestGlm52CompiledMTPCheckpoint(DeterministicDDPTestCase):
     def test_shared_mtp_depths_train_with_compile_and_topk_offload(self):
-        # 验证默认 reentrant checkpoint 可训练共享 MTP 深度且 loss 有限。
+        # 验证共享 MTP 深度可在 compile/offload 下训练且 loss 有限。
         self.create_pg("cuda")
         engine = _build_engine(
             intra_layer_micro_batch=1,
diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py
index c2a004896c..6fe55c7e8a 100644
--- a/tests/model/test_recompute.py
+++ b/tests/model/test_recompute.py
@@ -2,7 +2,6 @@
 
 TestCheckpointWrapper
     test_wrapper_is_transparent_to_state_dict_and_attributes: 包裹后参数名/state_dict/属性访问不变。
-    test_legacy_reentrant_preserves_gradients_through_dict_return: legacy reentrant 下 dict 返回值不断梯度。
     test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。
     test_checkpointing_keeps_the_module_itself: 换类而非套壳，isinstance/属性/容器协议原生可用。
     test_module_without_a_protocol_does_not_gain_one: 被包裹模块没有的协议不会凭空出现。
@@ -27,7 +26,7 @@
 import torch
 from torch import nn
 
-from xtuner.v1.model.utils import apply_gradient_checkpointing, apply_legacy_reentrant_checkpointing
+from xtuner.v1.model.utils import apply_gradient_checkpointing
 
 
 class _KeywordOnlyBlock(nn.Module):
@@ -85,26 +84,6 @@ def test_wrapper_is_transparent_to_state_dict_and_attributes(self):
         assert torch.equal(wrapped.state_dict()["linear.weight"], plain.state_dict()["linear.weight"])
         assert wrapped.tag == "block"
 
-    def test_legacy_reentrant_preserves_gradients_through_dict_return(self):
-        # reentrant 的 original 那趟在 no_grad 下执行，dict 里的 Tensor 不会被 autograd.Function
-        # 注册为输出，梯度会静默断掉（MTP 返回 TypedDict 时整个 mtp_block 的 grad 都是 None，
-        # 而 loss 仍然有限，所有断言都能过）。出入口都摊平后才不会断。
-        torch.manual_seed(0)
-        plain = _KeywordOnlyBlock()
-        wrapped = apply_legacy_reentrant_checkpointing(_KeywordOnlyBlock())
-        wrapped.load_state_dict(plain.state_dict())
-
-        x = torch.randn(2, 4, requires_grad=True)
-        plain({"x": x}, scale=2.0)["out"].square().sum().backward()
-        baseline_input_grad, x.grad = x.grad.clone(), None
-
-        wrapped({"x": x}, scale=2.0)["out"].square().sum().backward()
-
-        assert x.grad is not None
-        assert wrapped.linear.weight.grad is not None
-        assert torch.equal(x.grad, baseline_input_grad)
-        assert torch.equal(wrapped.linear.weight.grad, plain.linear.weight.grad)
-
     def test_non_tensor_signature_preserves_gradients(self):
         # 非 tensor 签名下梯度必须与不重算完全一致。
         torch.manual_seed(0)
diff --git a/tests/module/attention/test_dsa_mla.py b/tests/module/attention/test_dsa_mla.py
index 224e3a00cc..071cb1de90 100644
--- a/tests/module/attention/test_dsa_mla.py
+++ b/tests/module/attention/test_dsa_mla.py
@@ -5,7 +5,7 @@
 TestDSAAttention
     test_packed_inputs_respect_causal_boundaries_and_backward: packed attention 遵守分段因果边界并可反传。
     test_shared_layers_reuse_topk_without_cross_context_leak: shared layer 复用当前样本 top-k 且不跨样本泄漏。
-    test_reentrant_checkpoint_reuses_and_releases_topk: checkpoint 重算复用并最终释放 top-k。
+    test_checkpoint_reuses_and_releases_topk: checkpoint 重算复用并最终释放 top-k。
 TestAcceleratedSparseMLA
     test_tilelang_forward_backward_matches_torch: TileLang 前反向数值与 PyTorch 后端一致。
     test_compiled_cudnn_backward_matches_tilelang: 编译后的 cuDNN DSA 前反向与 TileLang 一致。
@@ -27,7 +27,7 @@
 
 from xtuner._testing import DeterministicDDPTestCase
 from xtuner.v1.data_proto import SequenceContext
-from xtuner.v1.model.utils import apply_legacy_reentrant_checkpointing
+from xtuner.v1.model.utils import apply_gradient_checkpointing
 from xtuner.v1.module.attention import DSAMLAConfig
 from xtuner.v1.module.attention.dsa_topk_sharing import register_dsa_topk_decoder_lifecycle_hooks
 from xtuner.v1.ops.sparse_mla import dsa_topk_indices, sparse_mla
@@ -211,13 +211,13 @@ def test_shared_layers_reuse_topk_without_cross_context_leak(self):
         assert seq_ctx.dsa_topk_cache.indices[0] is source_topk
         assert other_seq_ctx.dsa_topk_cache.indices[0] is not source_topk
 
-    def test_reentrant_checkpoint_reuses_and_releases_topk(self):
-        # 验证真实 source/shared decoder 经 reentrant checkpoint 重算后梯度有限且缓存释放。
+    def test_checkpoint_reuses_and_releases_topk(self):
+        # 验证真实 source/shared decoder 经 checkpoint 重算后梯度有限且缓存释放。
         torch.manual_seed(0)
-        source_block = apply_legacy_reentrant_checkpointing(
+        source_block = apply_gradient_checkpointing(
             _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=0))
         )
-        shared_block = apply_legacy_reentrant_checkpointing(
+        shared_block = apply_gradient_checkpointing(
             _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=1))
         )
         hidden_states = torch.randn(1, 4, 4, requires_grad=True)
diff --git a/tests/utils/test_pytree_reentrant_checkpoint.py b/tests/utils/test_pytree_reentrant_checkpoint.py
deleted file mode 100644
index 7f52d66169..0000000000
--- a/tests/utils/test_pytree_reentrant_checkpoint.py
+++ /dev/null
@@ -1,30 +0,0 @@
-"""Pytree reentrant checkpoint 的梯度行为测试。
-
-TestPytreeReentrantCheckpoint
-    test_nested_inputs_preserve_both_gradient_paths: 嵌套输入在 checkpoint 内外复用时梯度正确汇合。
-"""
-
-import torch
-from torch import nn
-from xtuner.v1.model.utils import apply_legacy_reentrant_checkpointing
-
-
-class NestedTensorBlock(nn.Module):
-    def forward(self, direct: torch.Tensor, nested: list[torch.Tensor]) -> torch.Tensor:
-        return direct * nested[0]
-
-
-class TestPytreeReentrantCheckpoint:
-    def test_nested_inputs_preserve_both_gradient_paths(self):
-        # 验证嵌套 Tensor 在 checkpoint 内外同时使用时不会重复反传旧 graph，且梯度正确相加。
-        direct_source = torch.tensor([2.0], requires_grad=True)
-        nested_source = torch.tensor([5.0], requires_grad=True)
-        direct = direct_source * 2
-        nested = nested_source * 3
-        block = apply_legacy_reentrant_checkpointing(NestedTensorBlock())
-
-        loss = block(direct, nested=[nested]).sum() + nested.square().sum()
-        loss.backward()
-
-        torch.testing.assert_close(direct_source.grad, torch.tensor([30.0]))
-        torch.testing.assert_close(nested_source.grad, torch.tensor([102.0]))
diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py
index 278295deec..7335d6d7a5 100644
--- a/xtuner/v1/config/fsdp.py
+++ b/xtuner/v1/config/fsdp.py
@@ -18,10 +18,6 @@ class FSDPConfig(BaseModel):
     recompute_ratio: Annotated[float, Parameter(help="Gradient checkpointing ratio for memory optimization")] = 1.0
     vision_recompute_ratio: Annotated[float, Parameter(help="Recompute ratio for vision modules")] = 1.0
     checkpoint_preserve_rng_state: Annotated[bool, Parameter(help="Preserve RNG state during checkpointing")] = True
-    mtp_checkpoint_use_reentrant: Annotated[
-        bool,
-        Parameter(help="Use reentrant checkpointing for MTP layers"),
-    ] = True
     # Training-time FSDP CPU offload is version-sensitive for XTuner model configs
     # that keep selected fp32 trainable parameters outside FSDP via
     # fp32_keys_pattern. The Qwen3.5-VL MoE RL path was verified to run on Torch
diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py
index 62244fd1fb..316bdfca05 100644
--- a/xtuner/v1/model/moe/moe.py
+++ b/xtuner/v1/model/moe/moe.py
@@ -47,7 +47,6 @@
 from xtuner.v1.model.utils import (
     ModelForwardExtraLogInfo,
     apply_gradient_checkpointing,
-    apply_legacy_reentrant_checkpointing,
     module_dict_repr,
 )
 from xtuner.v1.module import (
@@ -60,7 +59,6 @@
     NoAuxRouterConfig,
     RMSNorm,
 )
-from xtuner.v1.module.attention.dsa_topk_sharing import uses_dsa_topk_lifecycle
 from xtuner.v1.module.decoder_layer.dense_decoder_layer import (
     DenseDecoderLayer,
     DenseDecoderLayerMicroBatchOutput,
@@ -1170,16 +1168,7 @@ def fully_shard(
                 layer_idx=layer_idx,
                 mtp_idx=None,
             ):
-                # DSA cross-layer top-k sharing recognizes a checkpoint's original pass by grad
-                # being disabled, which only the reentrant implementation provides. Under the
-                # non-reentrant one both passes run with grad enabled, so the cache is never
-                # marked active and the shared top-k is never released. Keep those layers on the
-                # legacy path until the cache tracks the original/replay phase explicitly, which
-                # is also what lets `apply_legacy_reentrant_checkpointing` go away entirely.
-                if uses_dsa_topk_lifecycle(layer):
-                    layer = apply_legacy_reentrant_checkpointing(layer)
-                else:
-                    layer = apply_gradient_checkpointing(layer)
+                layer = apply_gradient_checkpointing(layer)
 
             self.layers[str(layer_idx)] = layer
             if layer_idx >= len(self.layers) - 1 and self.mtp_block is None:
@@ -1231,30 +1220,7 @@ 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 默认使用 reentrant 的原因：
-                    #   Case 1：最小触发条件是 compile, topk offload, MTP share weights and depth > 1.
-                    #   多个 logical depth 共用 top-k cache。reentrant 的 original
-                    #   关闭 grad、replay 开启 grad，DSA 能据此正确更新 cache 计数。
-                    #   original 不建立内部图，所以 replay 可以安全复用离散 top-k。
-                    #   non-reentrant 的两次执行都开启 grad，却仍沿用该复用策略，
-                    #   因而出现 original=COMPUTE、replay=REUSE，无法重建相同清单。
-                    #
-                    #   indexer 本身始终 no_grad。不开 compile 时，多执行/少执行一次
-                    #   indexer 不会改变 eager autograd 的保存清单；开启 compile 后，
-                    #   COMPUTE/REUSE 经过不同 graph break 和 compiled block，才可能让
-                    #   checkpoint 保存槽位错位并报 different metadata。例如 original
-                    #   保存 [A, B, C]、replay 保存 [A, X, C] 时，槽位 1 的 metadata
-                    #   不同。后续若显式记录 ORIGINAL/REPLAY phase，可再让
-                    #   non-reentrant 正确推进 cache 状态。
-                    #
-                    #   `apply_legacy_reentrant_checkpointing` 内部的 pytree 展开也是 reentrant
-                    #   专属的补丁：EP > 1、intra-layer micro-batch > 1 时 micro2 传入
-                    #   [embedding_0, embedding_1]，展开后 checkpoint 才能在 replay 前逐个
-                    #   detach，并在 backward 中把梯度交回原始 embedding graph。
-                    if self.fsdp_config.mtp_checkpoint_use_reentrant:
-                        mtp_layer = apply_legacy_reentrant_checkpointing(mtp_layer)
-                    else:
-                        mtp_layer = apply_gradient_checkpointing(mtp_layer)
+                    mtp_layer = apply_gradient_checkpointing(mtp_layer)
                 self.mtp_block.layers[mtp_idx] = mtp_layer
 
                 reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1
diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py
index 9675ac8071..a398224afd 100644
--- a/xtuner/v1/model/utils/__init__.py
+++ b/xtuner/v1/model/utils/__init__.py
@@ -1,10 +1,9 @@
-from .checkpointing import apply_gradient_checkpointing, apply_legacy_reentrant_checkpointing
+from .checkpointing import apply_gradient_checkpointing
 from .misc import ModelForwardExtraLogInfo, module_dict_repr
 
 
 __all__ = [
     "apply_gradient_checkpointing",
-    "apply_legacy_reentrant_checkpointing",
     "module_dict_repr",
     "ModelForwardExtraLogInfo",
 ]
diff --git a/xtuner/v1/model/utils/checkpointing.py b/xtuner/v1/model/utils/checkpointing.py
index 829f529027..81b8a092cb 100644
--- a/xtuner/v1/model/utils/checkpointing.py
+++ b/xtuner/v1/model/utils/checkpointing.py
@@ -4,17 +4,11 @@
 from typing import Any, Callable
 
 import torch.nn as nn
-from torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten
+from torch.utils._pytree import tree_flatten, tree_unflatten
 from torch.utils.checkpoint import checkpoint
 
 
-__all__ = [
-    "apply_gradient_checkpointing",
-    "apply_legacy_reentrant_checkpointing",
-    "checkpoint_flattened",
-    "install_checkpointing",
-    "CheckpointModule",
-]
+__all__ = ["apply_gradient_checkpointing", "checkpoint_flattened", "install_checkpointing", "CheckpointModule"]
 
 
 ContextFn = Callable[[], tuple[AbstractContextManager, AbstractContextManager]]
@@ -59,30 +53,6 @@ def checkpointed_call(original_call: Callable[..., Any], *args: Any, **kwargs: A
     return install_checkpointing(module, checkpointed_call)
 
 
-def apply_legacy_reentrant_checkpointing(module: nn.Module, *, preserve_rng_state: bool = True) -> nn.Module:
-    """Wrap ``module`` with the legacy reentrant checkpoint implementation.
-
-    Prefer :func:`apply_gradient_checkpointing`. Reentrant checkpointing runs the original pass
-    under ``no_grad`` and only the replay with grad enabled, which is required by the MTP layers of
-    DSA-based models: several logical MTP depths share one top-k cache, and only a grad-free
-    original pass lets the cache advance consistently across the two passes. This function can be
-    removed once the DSA top-k cache tracks the original/replay phase explicitly and the
-    ``mtp_checkpoint_use_reentrant`` config switch goes away.
-
-    Args:
-        module (nn.Module): Module whose forward should be recomputed during backward.
-        preserve_rng_state (bool): Restore the RNG state before recomputing. Defaults to True.
-
-    Returns:
-        nn.Module: A module that runs ``module`` under reentrant gradient checkpointing.
-    """
-
-    def checkpointed_call(original_call: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
-        return _pytree_reentrant_checkpoint(original_call, *args, preserve_rng_state=preserve_rng_state, **kwargs)
-
-    return install_checkpointing(module, checkpointed_call)
-
-
 def checkpoint_flattened(
     function: Callable[..., Any],
     *args: Any,
@@ -175,42 +145,3 @@ def original_call(*replayed_args: Any, **replayed_kwargs: Any) -> Any:
 # One checkpoint class per original class: checkpointing N layers of the same type must not build N
 # classes, and two checkpointed layers of the same type should stay `isinstance`-comparable.
 _CHECKPOINT_CLASSES: dict[type, type] = {}
-
-
-def _pytree_reentrant_checkpoint(
-    function: Callable[..., Any],
-    *args: Any,
-    preserve_rng_state: bool = True,
-    **kwargs: Any,
-) -> Any:
-    """让嵌套 Tensor 在 reentrant checkpoint 的出入口都成为 autograd 的一等公民。"""
-    # 原生 CheckpointFunction 只看得到顶层位置参数与 tensor 返回值，嵌套容器两头都会出问题。
-    #
-    # 入参侧：future_embeddings=[embedding_0, embedding_1] 对它来说只是“一个 list 参数”，
-    # 它看不到 list 里的两个 Tensor，也就不会 detach 它们。这可能会造成反向传播的错误，因为
-    # 这两个 Tensor 的梯度应该交由 CheckpointFunction.backward 的返回值交回原始 Tensor，
-    # 而不是由他们自己来传递梯度。tree_flatten 把输入摊平成
-    #   hidden, embedding_0, embedding_1
-    # 后 checkpoint 才能逐个 detach；tree_unflatten 再在 replay 前把 list 还原。
-    #
-    # 返回值侧：original 那趟在 no_grad 下执行，dict / TypedDict 里的 Tensor 因此不会被
-    # autograd.Function 注册为输出，回来时既没有 grad_fn 也不 require_grad——梯度会在此处
-    # 静默断掉（MTP 返回 TypedDict 时整个 mtp_block 的 grad 都是 None）。所以出口同样要摊平：
-    # 交给 checkpoint 的是纯 Tensor 元组，拿回来后再按 spec 还原成原本的结构。
-    flat_inputs, input_spec = tree_flatten((args, kwargs))
-    output_spec: TreeSpec | None = None
-
-    def run_function(*replayed_flat_inputs: Any) -> tuple[Any, ...]:
-        # 这里只还原参数结构，不会把 detached Tensor 重新连接到旧 graph；梯度由
-        # CheckpointFunction.backward 的返回值交回原始 Tensor。
-        nonlocal output_spec
-        replayed_args, replayed_kwargs = tree_unflatten(list(replayed_flat_inputs), input_spec)
-        flat_outputs, output_spec = tree_flatten(function(*replayed_args, **replayed_kwargs))
-        return tuple(flat_outputs)
-
-    flat_outputs = checkpoint(run_function, *flat_inputs, use_reentrant=True, preserve_rng_state=preserve_rng_state)
-    # original 与 replay 写入的 spec 相同，取任意一趟的结果即可。
-    assert output_spec is not None, "XTuner Internal Error: reentrant checkpoint did not run the function"
-    if not isinstance(flat_outputs, tuple):
-        flat_outputs = (flat_outputs,)
-    return tree_unflatten(list(flat_outputs), output_spec)

From 356868f4d6fced223d5fb658378562ca8035f3ce Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 28 Jul 2026 09:53:34 +0000
Subject: [PATCH 4/5] [Refactor] Remove the now-unused DSA lifecycle predicate

`uses_dsa_topk_lifecycle` existed for one reason: to select which decoder
layers had to stay on the reentrant checkpoint implementation, because DSA
cross-layer top-k sharing infers the checkpoint phase from
`torch.is_grad_enabled()`. The previous commit removed that implementation, so
the predicate answers a question no caller can act on any more -- a grep across
this branch and the two stacked on it finds only the definition.

The cache machinery it guarded (`_is_checkpoint_original_forward`,
`checkpoint_active`, `after_recompute_release`) stays: it is the subject of the
pending GLM-5.2 compatibility work, and the two `xfail` tests keep that gap
visible. Only the strategy-selection hook goes.

Also drops "reentrant" from the lifecycle-hook comment: non-reentrant recompute
re-invokes the decoder module and its hooks just the same, so the conclusion is
unchanged but the named mechanism no longer exists.
---
 .../v1/module/attention/dsa_topk_sharing.py   | 23 ++-----------------
 1 file changed, 2 insertions(+), 21 deletions(-)

diff --git a/xtuner/v1/module/attention/dsa_topk_sharing.py b/xtuner/v1/module/attention/dsa_topk_sharing.py
index 97c8757f1a..628675ae00 100644
--- a/xtuner/v1/module/attention/dsa_topk_sharing.py
+++ b/xtuner/v1/module/attention/dsa_topk_sharing.py
@@ -492,25 +492,6 @@ def _dsa_mtp_iteration_lifecycle_pre_hook(
         )
 
 
-def uses_dsa_topk_lifecycle(decoder_layer: torch.nn.Module) -> bool:
-    """Report whether a decoder layer participates in DSA cross-layer top-k
-    sharing.
-
-    The lifecycle distinguishes a checkpoint's original pass from its recompute pass by whether
-    grad is enabled (see ``_is_checkpoint_original_forward``), which only holds for the reentrant
-    checkpoint implementation. Callers that choose a checkpointing strategy use this to keep such
-    layers on the reentrant path; the distinction disappears once the cache tracks the
-    original/replay phase explicitly.
-
-    Args:
-        decoder_layer (torch.nn.Module): The decoder layer to inspect.
-
-    Returns:
-        bool: True if the DSA top-k lifecycle hooks are registered on the layer.
-    """
-    return getattr(decoder_layer, "_dsa_topk_decoder_lifecycle_hooks_registered", False)
-
-
 def register_dsa_topk_decoder_lifecycle_hooks(decoder_layer: torch.nn.Module) -> None:
     if getattr(decoder_layer, "_dsa_topk_decoder_lifecycle_hooks_registered", False):
         return
@@ -524,8 +505,8 @@ def register_dsa_topk_decoder_lifecycle_hooks(decoder_layer: torch.nn.Module) ->
     # recorded only pending actions and flushed them later. Remove that
     # transient state by keeping the entire residency transition at the decoder
     # boundary: the pre-hook launches H2D and the post-hook directly runs
-    # after_sparse_mla_use. Reentrant checkpoint replay invokes the decoder
-    # module and these hooks again, so main, micro-batch and MTP callers do not
+    # after_sparse_mla_use. Checkpoint recompute re-invokes the decoder module
+    # and these hooks along with it, so main, micro-batch and MTP callers do not
     # need separate lifecycle handling.
     #
     # This deliberately delays eager D2H until the decoder returns, losing its

From 13af713aea8920ae7afa186793443b2ba5c3df03 Mon Sep 17 00:00:00 2001
From: HAOCHENYE <21724054@zju.edu.cn>
Date: Tue, 11 Aug 2026 07:51:28 +0000
Subject: [PATCH 5/5] [Fix] Let a checkpointed layer's input tensors reach the
 caller's saved-tensor hooks

Inputs are flattened before being handed to `checkpoint` and reassembled
inside. Not for gradient correctness -- non-reentrant checkpointing gets
nested inputs right either way -- but for who else gets to see them.

`_CheckpointFrame.save_inputs` wraps only *top-level* tensor arguments into a
`SavedVariable`, and constructing one is what fires the ambient
`saved_tensors_hooks`. It runs just before the checkpoint installs its own
hooks, so those ambient hooks are still the caller's. That is how activation
offloading gets hold of a layer's inputs.

Converting the decoder layers to `TypedDict` I/O turned the call into
`decoder_layer(hidden_states_list, ...)`, so the only positional argument is a
list. Its tensors are stored as plain references, reach no hook, and are never
offloaded -- silently: gradients stay correct and nothing fails.

Measured on Qwen3-MoE-30BA3, `ep_size=4`, `intra_layer_micro_batch=2`, 8k,
eager, all2all, peak allocated at the last step:

| | peak allocated | peak reserved |
|---|---|---|
| offload off | 97.73 GB | 122.40 GB |
| offload on, before | 98.07 GB | 122.75 GB |
| offload on, after | 96.09 GB | 119.29 GB |

Enabling offload used to cost 0.34 GB and save nothing. Step-1 loss is
bit-identical across all three.

The regression test asserts by `data_ptr`, not by shape: a checkpointed
region's output is easily the same shape as its input, and an earlier version
of this test passed against the unfixed code for exactly that reason.
---
 tests/model/test_recompute.py | 47 +++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py
index 6fe55c7e8a..1d323cd926 100644
--- a/tests/model/test_recompute.py
+++ b/tests/model/test_recompute.py
@@ -19,12 +19,15 @@
     test_a_callable_unit_keeps_its_callers_compiled: KeptCallables 只退出自身，调用者仍编译。
     test_no_unit_withdraws_the_method_that_holds_most_compilation: 没有 unit 撤出编译占比最大的方法。
     test_attention_is_kept_by_op_identity: attention 走 op identity 而非撤出 callable。
+    test_input_tensors_reach_the_ambient_saved_tensor_hooks: 嵌套/关键字传入的输入也能进外层 hook。
 """
 
 
 
+import pytest
 import torch
 from torch import nn
+from torch.autograd.graph import saved_tensors_hooks
 
 from xtuner.v1.model.utils import apply_gradient_checkpointing
 
@@ -70,6 +73,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+class _FlexibleBlock(nn.Module):
+    """接受任意摆放的输入：位置的容器、字典、关键字参数，用来覆盖各种嵌套形状。"""
+
+    def __init__(self) -> None:
+        super().__init__()
+        # 输入 4 维、输出 6 维：输出与输入形状不同，断言才不会把输出误当成输入。
+        self.linear = nn.Linear(4, 6)
+
+    def forward(self, inputs, *, scale: float, extra: torch.Tensor | None = None) -> dict[str, torch.Tensor]:
+        tensors = list(inputs.values()) if isinstance(inputs, dict) else list(inputs)
+        if extra is not None:
+            tensors.append(extra)
+        return {"out": sum(self.linear(t) * scale for t in tensors)}
+
+
 class TestCheckpointWrapper:
     def test_wrapper_is_transparent_to_state_dict_and_attributes(self):
         # 包裹层不能出现在参数名里，否则 checkpoint 的存/取与非重算模型不兼容。
@@ -120,3 +138,32 @@ def test_module_without_a_protocol_does_not_gain_one(self):
 
         assert bool(checkpointed) is True
         assert not hasattr(type(checkpointed), "__len__")
+
+    @pytest.mark.parametrize(
+        "make_call",
+        [
+            pytest.param(lambda block, x: block([x], scale=2.0), id="nested-in-list"),
+            pytest.param(lambda block, x: block({"x": x}, scale=2.0), id="nested-in-dict"),
+            pytest.param(lambda block, x: block([], scale=2.0, extra=x), id="passed-by-keyword"),
+        ],
+    )
+    def test_input_tensors_reach_the_ambient_saved_tensor_hooks(self, make_call):
+        # 激活 offload 是靠外层 saved_tensors_hooks 拿到层输入的，而 checkpoint 只把**顶层**
+        # tensor 参数包成 SavedVariable（构造它才会触发 hook）。所以嵌套在容器里、或走关键字
+        # 传进来的 tensor 会一个 hook 都不经过——offload 静默空转，梯度却完全正确，没有任何
+        # 现象能暴露它。这里直接断言 hook 收得到。
+        packed: list[int] = []
+
+        class _Record(saved_tensors_hooks):
+            # 按 data_ptr 认张量，不按 shape：区域的输出很容易和输入同形，
+            # 按 shape 断言会把输出当成输入，测试变成恒绿。
+            def __init__(self) -> None:
+                super().__init__(lambda t: (packed.append(t.data_ptr()), t)[1], lambda t: t)
+
+        wrapped = apply_gradient_checkpointing(_FlexibleBlock())
+        x = torch.randn(2, 4, requires_grad=True)
+
+        with _Record():
+            make_call(wrapped, x)["out"].square().sum().backward()
+
+        assert x.data_ptr() in packed