Uh oh!
There was an error while loading. Please reload this page.
[Refactor] Unify SAC checkpointing and GLM DSA dataflow - #2039
Conversation
d78f607 to
81e0010CompareFlatten structured inputs and outputs at one checkpoint boundary so nested MTP micro-batch tensors receive gradients. Keep FSDP outside replay and remove the non-reentrant MTP switch.
Move GLM-specific decoder and MTP adapters under the model package, thread DSA IDs through keyed outputs, and remove SequenceContext cache lifecycle state. Share one saved-tensor offload window for activation and DSA ID storage.
Keep reusable no-grad outputs in checkpoint-call-local FIFO frames, freeze DSA indexers through their model config, and preserve the original int32 storage across shared layers and offload. Add real regressions for source call counts and pinned-memory DSA offload.
Only inject the detached-input grad entry when the checkpointed module still has trainable parameters. This preserves MTP parameter gradients while keeping frozen vision outputs detached, avoiding a replay with no differentiable outputs.
b2047db to
5c15855CompareYanhuiDua
commented
Sep 2, 2026
@claude review |
Claude: Summary本 PR 把 activation checkpointing 统一到一个 PyTree-aware 的 reentrant 边界（ Main Flowchart after this PRflowchart 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
核心原理实现与单测
抽象与信息隐藏评估
其他 Issues
VerdictREQUEST_CHANGES |
| activation_offload = int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1 | ||
| dsa_topk_offload = int(os.getenv("XTUNER_DSA_TOPK_OFFLOAD", "0")) == 1 | ||
| offload_tensors: list[torch.Tensor] = [] | ||
| if activation_offload and layer_idx >= self.config.first_k_dense_replace: | ||
| offload_tensors = list(hidden_states) if is_micro_batch else [cast(torch.Tensor, hidden_states)] | ||
| if dsa_topk_offload and layer_idx in self._dsa_topk_last_consumers and dsa_topk_ids is not None: | ||
| offload_tensors.extend(dsa_topk_ids if isinstance(dsa_topk_ids, list) else [dsa_topk_ids]) | ||
| # The offload context expects a dense zero-based block index across every | ||
| # active activation or DSA-ID window, including dense GLM layers. | ||
| offload_block_idx = sum( | ||
| (activation_offload and previous_idx >= self.config.first_k_dense_replace) | ||
| or ( | ||
| dsa_topk_offload | ||
| and previous_idx in self._dsa_topk_last_consumers | ||
| and self._dsa_topk_source_layers[previous_idx] != previous_idx | ||
| ) | ||
| for previous_idx in (int(idx) for idx in self.layers) | ||
| if previous_idx < layer_idx | ||
| ) |
There was a problem hiding this comment.
Claude: [设计] activation-offload 规则出现两份实现。
Problem：基类 MoE._call_decoder_layer 已把「窗口张量选择 + block_idx 编号」写死在方法体内，GLM 为了插入 DSA IDs 只能整段 override，于是同一条规则有了两个不同公式：基类是 sum(first_k_dense_replace <= idx < layer_idx)，这里则把 DSA 末位消费层也计入密集编号。做 Deletion test：删掉这个 override 会同时丢失 DSA 数据流和 offload 规则，说明 Seam 切在了整个 _call_decoder_layer 上——位置偏大。后果是基类调整 offload 策略时，GLM 侧不会报错，只会静默错位到另一套编号。
Solution：基类保留唯一的 _call_decoder_layer（窗口构造 + block 计数），只把变化点做成小钩子，GLM 仅覆盖钩子：
# MoEdef_extra_offload_tensors(self, layer_idx, hidden_states) ->list[torch.Tensor]:
return []
def_layer_forward_kwargs(self, layer_idx, previous_layer_results) ->dict:
return {}Benefits：offload 规则单点化（Locality），子类需要了解的 Interface 从「整个调用 + 窗口 + 编号」缩小到两个钩子（Leverage），block_idx 编号也能在基类上直接写 CPU 单测。
| GLM52MoEDecoderLayerOutput | GLM52MoEDecoderLayerMicroBatchOutput | None, | ||
| previous_layer_results, | ||
| ) | ||
| dsa_topk_ids = None if previous_results is None else previous_results["dsa_topk_ids"] |
There was a problem hiding this comment.
Claude: [兼容性] index_share_for_mtp_iteration 的唯一运行时使用点已被删除，这里无条件透传上一 depth 的 IDs。但该字段仍保留在 config、from_hf 校验与 to_hf_config 回写中，设为 False 既不生效也不报错。建议按该配置决定是否透传，或在配置层直接拒绝 False。
| """Build one policy-neutral saved-tensor offload window. | ||
| The decoder-stack caller decides which tensors belong to the current | ||
| window and advances ``block_idx`` only when the list is non-empty. | ||
| """ | ||
| if not tensors: | ||
| return contextlib.nullcontext() |
There was a problem hiding this comment.
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 路径。
| if self.compile_cfg: | ||
| layer.forward = torch.compile(layer.forward, fullgraph=True) | ||
| # The PyTree checkpoint adapter is an intentional graph break; model compute | ||
| # inside it keeps the independently configured full-graph compilation. | ||
| compiled_forward = torch.compile(type(layer).forward, fullgraph=False) | ||
| layer.forward = compiled_forward.__get__(layer, type(layer)) |
There was a problem hiding this comment.
Claude: [性能] 此处 type(layer) 已是 CheckpointWrapper，而内层 Qwen3VLVisionLayer.forward 并没有注释所说的「独立 full-graph 配置」（该文件只 import 了 maybe_compile，未装饰任何 vision forward）。结果是 fullgraph 保证被静默移除，vision 侧 graph break 不再报错，且本 PR 无 VL 回归验证。建议把 fullgraph=True 保留在被包裹模块自身的 forward 上，并补一条 VL 验证。
Summary
Regression coverage
Validation