`GradientCheckpointingLayer` checkpoints FSDP2 hooks; `use_reentrant=True` silently drops gradients when a layer runs more than once per bac

快速结论：当 Transformers 的 GradientCheckpointingLayer 与 FSDP2 的 fully_shard 组合，且梯度检查点使用 use_reentrant=True 时，同一个 decoder layer 在单次 backward 前被多次 forward，该层梯度可能只保留最后一次 forward 的贡献，且不会报错。优先确认训练流程是否“多次 forward 后一次 backward”，例如部分 RL/DPO 训练器。

适用环境：Issue 中确认 transformers 5.14.1（main 提交 cdb75de 上 GradientCheckpointingLayer.__call__ 未变），torch 2.16.0a0（main），并在 torch 2.11.0、2.15.0 和 2.16 nightly 上出现同类问题；Python 3.12、CPU、gloo、2 个 rank 即可复现，无需 GPU。

最快修复方案：暂无确认的一步修复方案。Issue 讨论中验证有效的方向是让检查点只包裹子类 forward、不包裹 super().__call__，使模块 hook（包括 FSDP2 hook）位于 checkpoint 外层；该补丁尚未合入，属于待评审方案。

注意事项：该补丁仅在 Issue 讨论中针对 cdb75de 验证，覆盖了部分模型测试，不代表已发布版本可用；use_reentrant=False 在复现中未显示同样的梯度丢失，但 FSDP2 官方测试主要覆盖非 reentrant 检查点，不能据此认为任意组合都安全。回退方案（遇到 fully_shard 层且 use_reentrant=True 时抛错）只是讨论中的备选，未被验证。

问题场景

用户使用 Transformers 的 LlamaForCausalLM（以及其他继承 GradientCheckpointingLayer 的模型层），开启 gradient_checkpointing_enable()，并对每个 decoder layer 以及整体模型调用 fully_shard 进行 FSDP2 分片。在单次 backward() 之前对同一层做多次 forward（例如 RL/DPO 训练器每个 backward 前跑多个 forward），梯度与未分片参考模型不一致，但训练不报错，optimizer.step() 照常执行。

报错原文

`GradientCheckpointingLayer` checkpoints FSDP2 hooks; `use_reentrant=True` silently drops gradients when a layer runs more than once per backward

gradient_checkpointing_enable(use_reentrant=True) + fully_shard(layer)     forwards_per_backward=1 grad_rel_error=0.000
...
with 2 forwards before one backward, the gradient error is 0.476 vs. the unsharded model, and 0.0 with 1 forward

原因分析

最可能的原因是 checkpoint 嵌套顺序与 FSDP2 的期望相反。GradientCheckpointingLayer.__call__ 对 super().__call__ 做 checkpoint，等于把 nn.Module.__call__ 连同模块的 forward hooks 一起包进检查点，形成 checkpoint(fsdp_layer, x)；而 FSDP2 期望 FSDP 在外、checkpoint 在内，例如 fully_shard(checkpoint_wrapper(layer))。

在 use_reentrant=True 下，原始 forward 在 no_grad 中执行，FSDP2 只能在 recompute 阶段注册 backward hook，并且每次 backward 中每层只有第一次 recompute 能成功注册。Issue 中的状态记录显示：第二次 forward 的 recompute 正常注册 hook 并归约梯度，而第一次 forward 的 recompute 时状态已为 PRE_BACKWARD，_fsdp_state._pre_forward 提前返回，梯度被丢弃。因此该层梯度只反映最后一次 forward，且没有任何异常提示。

环境排查

确认 transformers 版本，以及 GradientCheckpointingLayer.__call__ 是否仍是 cdb75de 中的实现（对 super().__call__ 做 checkpoint）。
确认 PyTorch 版本；Issue 在 torch 2.11.0、2.15.0、2.16 nightly/2.16.0a0 上均观察到问题。
确认梯度检查点参数：是否显式传入 use_reentrant=True，还是使用默认值。
确认 FSDP2 使用方式：是否对 layer 调用 fully_shard(layer)，以及是否再对整体模型 fully_shard(model)。
确认训练循环是否存在一个 backward 前多次 forward；这是触发梯度静默丢失的关键条件之一。
单进程 gloo 即可复现，CPU 环境即可排查，无需 GPU。

解决步骤

先复现并量化问题：用 Issue 中的脚本对比分片模型与未分片模型的梯度相对误差，分别设置每个 backward 前 1 次和 2 次 forward。1 次 forward 误差应接近 0，2 次 forward 若误差明显非 0，说明命中该问题。
临时规避：将 gradient_checkpointing_kwargs 中的 use_reentrant 改为 False，或在当前训练流程中避免同一层在单次 backward 前多次 forward。前者在 Issue 复现中未出现同样的梯度丢失，但属于可优先尝试的规避手段，并非官方修复。
调整 FSDP2 嵌套顺序（可优先尝试）：不要让 Transformers 的 layer 级 checkpoint 包住 FSDP hook，改为 FSDP 在外、checkpoint 在内，例如在分片时使用 fully_shard(checkpoint_wrapper(layer, checkpoint_impl=...)) 的结构。Issue 中该结构在 1 次和 2 次 forward 下梯度误差均为 0。
若需要从框架侧修复：参考 Issue 讨论中的补丁，使 GradientCheckpointingLayer 只对子类 forward 做 checkpoint，而不是对 super().__call__ 做 checkpoint，从而让模块 hook 在 checkpoint 外层执行（hooks -> checkpoint(forward) -> hooks）。该补丁包含 __init_subclass__ 包装 forward、用 _gc_running 标志防止 super().forward 被重复 checkpoint 等改动，属于待维护者评审的方案，合入前不建议直接用于生产。
如果无法调整嵌套顺序，可参考讨论中的回退策略：当 use_reentrant=True 且层已被 fully_shard 时显式报错，避免静默丢梯度。该策略仅为讨论备选，未被验证。

验证方法

用 Issue 的复现脚本，对 use_reentrant=True + fully_shard(layer) 分别跑 forwards_per_backward=1 和 2，检查 grad_rel_error 是否都接近 0。修复前 2 次 forward 的误差约为 0.476，修复后讨论中的补丁报告所有行均为 0.0。也可增加 forward hook 计数断言，确认每次 forward 在 grad 开启时 hook 只执行一次。