[Refactor] Unify SAC checkpointing and GLM DSA dataflow

[jayhenry]

Summary

• Unify activation checkpointing around one PyTree-aware reentrant boundary so structured micro-batch inputs and outputs stay on the real autograd and offload path.
• Pass GLM DSA top-k IDs explicitly through decoder and MTP outputs, removing the mutable SequenceContext cache lifecycle.
• Reuse frozen source-indexer results from checkpoint-local FIFO frames during replay.
• Keep FSDP wrappers outside recompute and dense checkpoint boundaries outside compiled graphs.

Regression coverage

• compile + top-k offload + shared-weight MTP with depth > 1; validates pinned-memory offload and restore, and that the source indexer is not recomputed
• EP > 1 + intra-layer micro-batch 2
• FSDP checkpointing, torch.compile, FP8, activation offload, and nested PyTree inputs and outputs

Validation

• Post-rebase on latest upstream/main: 16 passed
  • tests/model/test_recompute.py
  • tests/model/test_glm52_mtp_checkpoint_repro.py
  • tests/model/test_fsdp_checkpoint.py
• Qwen3.5 full 40-layer stack5 regression: step losses and grad norms matched upstream; steady sequence TGS +0.36%; peak allocated memory -0.14 GB.
• GLM AdamW production-shaped and Muon EP2/micro2 training regressions completed.

[YanhuiDua]

@claude review

[github-actions]

Claude:

Summary

本 PR 把 activation checkpointing 统一到一个 PyTree-aware 的 reentrant 边界（apply_activation_checkpointing + reuse_during_recompute），并把 GLM-5.2 的 DSA top-k 从 SequenceContext 上的可变缓存改为经 decoder / MTP 输出显式穿参，同时将 decoder 与 MTP 的返回值从位置元组改为 TypedDict，GLM-5.2 相关代码收敛到新的 xtuner/v1/model/moe/glm52/ 包。整体方向是把隐式生命周期换成显式数据流，抽象更清晰；但 offload 策略在基类与 GLM 子类各写了一份，另有一处配置项失效、一处 offload 匹配语义放宽、一处 vision compile 降级需要确认。

Main Flowchart after this PR

flowchart TD
    A["MoE._forward / _micro_batch_forward"] --> B["_decoder_stack /<br/>_micro_batch_decoder_stack<br/>(新增, 11 参数 + 可变 output dict)"]
    B --> C{"模型类型"}
    C -->|"通用 MoE"| D["MoE._call_decoder_layer<br/>(offload 窗口 + block_idx 规则)"]
    C -->|"GLM-5.2"| E["Glm52MoE._call_decoder_layer<br/>(整段 override: DSA IDs + 另一套 block_idx 规则)"]
    D --> F["_saved_tensors_offload_ctx<br/>(storage-ptr 匹配, reserve_pin_memory)"]
    E --> F
    F --> G["decoder_layer(...) -> TypedDict"]
    G --> H["GLM52AttnOutputs.dsa_topk_ids<br/>显式传给下一 consumer 层"]
    H --> I["MTPBlock._call_decoder_layer<br/>GLM52MTPBlock 无条件透传 dsa_topk_ids"]

    style D fill:#ffe0b2,stroke:#e65100
    style E fill:#ffcdd2,stroke:#b71c1c
    style F fill:#ffcdd2,stroke:#b71c1c
    style I fill:#ffe0b2,stroke:#e65100

Loading

核心原理实现与单测

• Checkpoint 统一边界：apply_activation_checkpointing 固定 reentrant 并用 PyTree 桥接展平输入 / 还原输出，reuse_during_recompute 以 ContextVar + 按 callable 的 FIFO frame 支撑 replay 复用。tests/model/test_recompute.py 全部走 public API，覆盖了 grad 模式序列 [False, True]、全 detach 输入仍产参数梯度、冻结模块不建立 replay 边、嵌套 / 关键字输入可被外层 saved_tensors_hooks 捕获（即 offload 不会静默空转）、replay 输出结构变化报错、以及跨两次 checkpoint 的 FIFO 隔离——核心行为有真实代码路径覆盖，未 mock 项目内模块。
• GLM DSA 显式数据流：source 层算一次 IDs，consumer 层由 Glm52MoE._call_decoder_layer 显式传入；tests/module/attention/test_dsa_mla.py、tests/model/test_glm52_moe.py::TestGlm52ExplicitDsaDataflow、tests/model/test_glm52_mtp_checkpoint_repro.py（indexer 调用计数 1 / 2、pinned D2H 真实命中并在 backward 恢复）覆盖了 source 不重算与 top-k offload。
• 覆盖缺口（未达 Warning，仅提示）：_call_decoder_layer 的 offload block-index 计算无 CPU 可跑单测；index_share_for_mtp_iteration=False 无任何测试。

抽象与信息隐藏评估

• Warning — xtuner/v1/model/moe/moe.py _call_decoder_layer 与 xtuner/v1/model/moe/glm52/glm52.py 的 override：Seam 切在整个 _call_decoder_layer 上，导致同一条 activation-offload 规则（窗口张量 + block_idx 编号）在基类和 GLM 子类各存一份且公式不同，基类策略变更时 GLM 会静默偏移。
• Warning — xtuner/v1/model/moe/moe.py _decoder_stack / _micro_batch_decoder_stack：抽出的私有方法带 11 个关键字参数并把 output: dict 当可变入出参写入，没有隐藏规则，只是把循环搬家。

其他 Issues

• Warning — xtuner/v1/model/moe/glm52/glm52.py：删除 index_share_for_mtp_iteration 的唯一运行时使用点后，该配置项仍保留在 config、from_hf 校验和 to_hf_config 回写中，但设为 False 既不改变行为也不报错，变成静默失效的配置。
• Warning — xtuner/v1/model/moe/moe.py _saved_tensors_offload_ctx：offload 匹配从 data_ptr() 放宽为 untyped_storage().data_ptr()，会把共享 storage 的视图一并纳入窗口，且单 micro-batch 路径新增常驻 pinned 缓冲并改变 block_idx 编号，而现有基准未覆盖 XTUNER_ACTIVATION_OFFLOAD=1 的单 micro-batch 路径。
• Warning — xtuner/v1/model/compose/qwen3_vl/modeling_vision.py 与 xtuner/v1/model/compose/intern_s1/modeling_vision.py：compile 目标改为 checkpoint wrapper 的 forward 且降为 fullgraph=False，但内层 vision layer 并没有注释所称的"独立 full-graph 配置"，fullgraph 保证被静默移除且本 PR 无 VL 回归验证。

Verdict

REQUEST_CHANGES

[claude]

Claude: [正确性] 匹配条件从原来的 x.data_ptr() == hidden_states.data_ptr() 放宽为 storage 指针，会把共享 storage 的非零偏移视图一并纳入窗口，同一 storage 可能被重复 pack/D2H。另外单 micro-batch 路径新增了 reserve_pin_memory=True 且 block_idx 由绝对层号改为相对号，改变了 OffloadManager 的 key 与 prefetch 顺序，而本 PR 基准未覆盖 XTUNER_ACTIVATION_OFFLOAD=1 的单 micro-batch 路径。