[Refactor] Switch gradient checkpointing to the non-reentrant implementation (1/3)

.dev_scripts/repro_fsdp_checkpoint_prefetch.py

@@ -0,0 +1,125 @@
+"""Check FSDP2 checkpoint composition and backward-prefetch bookkeeping.
+
+Run on two GPUs, for example::
+
+    torchrun --standalone --nproc-per-node=2 \
+        .dev_scripts/repro_fsdp_checkpoint_prefetch.py --mode reentrant
+    torchrun --standalone --nproc-per-node=2 \
+        .dev_scripts/repro_fsdp_checkpoint_prefetch.py --mode non-reentrant
+
+The model intentionally has no MoE, compile, explicit forward prefetch, or large
+activations. It isolates the default FSDP2 backward-prefetch bookkeeping from
+the Qwen3.5-VL async-RL case; an outer checkpoint wrapper should keep the
+logical post-forward order equal to the number of layers in both modes.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import fully_shard
+
+from xtuner.v1.model.utils import apply_gradient_checkpointing
+
+
+class Block(nn.Module):
+    def __init__(self, hidden_size: int) -> None:
+        super().__init__()
+        self.up = nn.Linear(hidden_size, hidden_size * 2, bias=False)
+        self.down = nn.Linear(hidden_size * 2, hidden_size, bias=False)
+        self.grad_modes: list[bool] = []
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        self.grad_modes.append(torch.is_grad_enabled())
+        return hidden_states + self.down(torch.nn.functional.silu(self.up(hidden_states)))
+
+
+class Model(nn.Module):
+    def __init__(self, num_layers: int, hidden_size: int) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([Block(hidden_size) for _ in range(num_layers)])
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return hidden_states
+
+
+@dataclass
+class Observation:
+    max_pending_groups: int = 0
+    max_pending_bytes: int = 0
+    max_post_forward_order: int = 0
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mode", choices=("reentrant", "non-reentrant"), required=True)
+    parser.add_argument("--num-layers", type=int, default=6)
+    parser.add_argument("--hidden-size", type=int, default=512)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group("nccl")
+
+    torch.manual_seed(0)
+    model = Model(args.num_layers, args.hidden_size).cuda()
+    mesh = init_device_mesh("cuda", (dist.get_world_size(),))
+    use_reentrant = args.mode == "reentrant"
+
+    for index, layer in enumerate(model.layers):
+        layer = apply_gradient_checkpointing(layer, use_reentrant=use_reentrant)
+        model.layers[index] = layer
+        fully_shard(layer, mesh=mesh, reshard_after_forward=True)
+    fully_shard(model, mesh=mesh, reshard_after_forward=True)
+
+    observation = Observation()
+
+    def observe_fsdp_state(_module: nn.Module, _inputs: tuple[object, ...], _output: object) -> None:
+        # Synchronize only in this diagnostic script so async all-gathers have a
+        # stable lifetime at the point where we account for their storage.
+        torch.cuda.synchronize()
+        param_groups = [fully_shard.state(layer)._fsdp_param_group for layer in model.layers]
+        pending = [group._all_gather_result for group in param_groups]
+        pending = [result for result in pending if result is not None]
+        pending_bytes = sum(result.all_gather_output.nbytes for result in pending)
+        comm_ctx = param_groups[0].comm_ctx
+        observation.max_pending_groups = max(observation.max_pending_groups, len(pending))
+        observation.max_pending_bytes = max(observation.max_pending_bytes, pending_bytes)
+        observation.max_post_forward_order = max(
+            observation.max_post_forward_order,
+            len(comm_ctx.post_forward_order),
+        )
+
+    # Register after fully_shard so this observer runs after FSDP's post-forward
+    # hook. It measures real FSDP state without replacing any production method.
+    handles = [layer.register_forward_hook(observe_fsdp_state) for layer in model.layers]
+    inputs = torch.randn(2, 8, args.hidden_size, device="cuda", requires_grad=True)
+    model(inputs).float().square().mean().backward()
+    torch.cuda.synchronize()
+
+    if dist.get_rank() == 0:
+        grad_modes = [layer.grad_modes for layer in model.layers]
+        print(f"mode={args.mode}")
+        print(f"grad_modes={grad_modes}")
+        print(f"max_post_forward_order={observation.max_post_forward_order}")
+        print(f"max_pending_groups={observation.max_pending_groups}")
+        print(f"max_pending_bytes={observation.max_pending_bytes}")
+
+    for handle in handles:
+        handle.remove()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

tests/engine/test_moe_train_engine_float8.py

@@ -20,7 +20,7 @@
 from xtuner.v1.utils.device import get_device
 from xtuner.v1.model.base import ModelItem
 from xtuner.v1.loss.ce_loss import CELossConfig
-from xtuner.v1.model.moe.moe import BalancingLossConfig
+from xtuner.v1.model.moe.moe import MOE_BLOCK_FORWARD, BalancingLossConfig
 
 
 
@@ -35,11 +35,9 @@ class TestMoEEngineFloat8(DeterministicDDPTestCase):
         "device,ep_size,hsdp_sharding_size,sim_tol,rtol",
         [
             ("cuda", 1, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.01),
-            # ep8 is a smoke/trend coverage for the FSDP shard-mesh-size-1 FP8 path.
-            # It shares the ep1 reference below, but is not expected to align step-by-step
-            # because EP changes routing/collective order and accumulates FP8 numeric drift.
-            # Observed 10-step loss:
-            # [2.4714, 2.4714, 1.8044, 1.5210, 0.9570, 0.6952, 0.4370, 0.3123, 0.1714, 0.1100]
+            # EP8 covers checkpoint replay across layer-varying routed-token shapes while MoEBlock
+            # remains fullgraph-compiled. It shares the EP1 reference below, but EP changes routing
+            # and collective order, so the two loss curves need not align step-by-step.
             ("cuda", 8, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.15),
         ],
     )
@@ -66,6 +64,10 @@ def test_tile_wise_fp8(self, device, ep_size, hsdp_sharding_size, sim_tol, rtol)
             optim_cfg=optim_cfg,
             fsdp_cfg=fsdp_cfg,
         )
+        if ep_size > 1:
+            # Regression contract: checkpoint replay must remain correct while the EP expert block
+            # keeps its strict full-graph compile boundary.
+            self.assertEqual(engine.model.compile_cfg.get(MOE_BLOCK_FORWARD), {"fullgraph": True})
         engine.from_hf(hf_path=QWEN3_MOE_PATH)
 
         loss_cfg = CELossConfig()

tests/model/test_fsdp_checkpoint.py

@@ -0,0 +1,62 @@
+import torch
+
+from xtuner._testing import DeterministicDDPTestCase
+from xtuner.v1.config import FSDPConfig
+from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.model.dense.qwen3 import Qwen3DenseConfig
+from xtuner.v1.module.attention import MHAConfig
+
+
+class TestFSDPCheckpoint(DeterministicDDPTestCase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    def test_reentrant_checkpoint_keeps_fsdp_outside_recompute(self):
+        self.create_pg("cuda")
+        config = Qwen3DenseConfig(
+            vocab_size=64,
+            max_position_embeddings=64,
+            eos_token_id=2,
+            bos_token_id=1,
+            num_hidden_layers=2,
+            hidden_size=32,
+            intermediate_size=64,
+            rms_norm_eps=1e-6,
+            hidden_act="silu",
+            attention=MHAConfig(
+                num_attention_heads=4,
+                num_key_value_heads=2,
+                head_dim=8,
+                qk_norm=True,
+            ),
+            compile_cfg=False,
+        )
+        model = config.build().cuda()
+        grad_modes: list[bool] = []
+        original_layer = model.layers["0"]
+        original_layer.register_forward_pre_hook(lambda _module, _inputs: grad_modes.append(torch.is_grad_enabled()))
+
+        model.fully_shard(
+            FSDPConfig(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.bfloat16,
+                torch_compile=False,
+            )
+        )
+        checkpoint_calls = 0
+
+        def record_checkpoint_call(_module, _inputs, _output):
+            nonlocal checkpoint_calls
+            checkpoint_calls += 1
+
+        model.layers["0"].register_forward_hook(record_checkpoint_call)
+        input_ids = torch.randint(0, config.vocab_size, (1, 8), device="cuda")
+        output = model(SequenceContext.from_input_ids((input_ids,)))
+        assert output.logits is not None
+        output.logits.sum().backward()
+
+        # The original layer and its lifecycle hooks must be replayed, while
+        # the outer FSDP/checkpoint boundary is one logical forward only.
+        assert grad_modes == [False, True]
+        assert checkpoint_calls == 1

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,

tests/model/test_qwen3_5_dense.py

@@ -99,7 +99,7 @@ def test_decoder_layer_bitwise_parity(self, device, layer_idx):
             loss_hf.backward()
 
             x_xt = base.clone().requires_grad_(True)
-            o_xt = xt_layer(x_xt, position_embeddings=(cos, sin), seq_ctx=seq_ctx)
+            o_xt = xt_layer(x_xt, position_embeddings=(cos, sin), seq_ctx=seq_ctx)["hidden_states"]
             loss_xt = F.cross_entropy(F.linear(model.norm(o_xt), model.lm_head.weight).reshape(-1, cfg.vocab_size), labels)
             loss_xt.backward()