背景概述 Fully Sharded Data Parallel（FSDP）作为一种高效的分布式训练策略，通过将模型参数分片存储于多张 GPU 上，显著降低了单卡显存占用，支持超大模型的训练。然而，在结合反向重计算（Checkpointing）技术以进一步节省显存时，部分场景下仍可能出现显存占用异常升高的问题。问题现象在使用 FSDP 对多模态模型进行分片并启用反向重计算机制后，观察到反向传播阶段的显存占用约为前向传播的两倍，显存峰值显著上升，影响训练稳定性与资源利用率。通过设备内存快照工具分析，发现反向计算过程中存在两处显存未释放的异常堆积点。根因分析 1. 显存快照定位通过 PyTorch 提供的设备内存快照功能，定位到反向计算阶段存在两个持续增长的显存占用点。 ●第一处为 reduce 操作，用于在 FSDP 中聚合各卡梯度，属于正常行为。 ●第二处为 allgather 操作，用于在重计算过程中收集分片参数以恢复前向计算状态，其申请的显存未被及时释放，构成异常。 2. 调用栈分析深入分析调用栈发现，allgather 操作在反向重计算中被重复触发。当 fully_shard 被嵌套于重计算逻辑内部时，触发了 pre_forward 与 pre_backward 钩子函数，导致参数被多次 unshard，但未在后续及时 reshard，造成显存累积。 3. 核心问题当 fully_shard 被包裹在 checkpoint_wrapper（反向重计算）内部时，FSDP 的默认行为会在每次反向传播前自动执行 unshard 操作以获取完整参数，但未在计算完成后主动触发 reshard，导致分片参数未及时释放，显存持续累积。解决方案 1. 显式启用 reshard_after_backward 在调用 fully_shard 后，显式设置 set_reshard_after_backward(True)，确保在反向传播完成后立即对分片参数执行 reshard 操作，释放临时占用的显存。 2. 框架级优化已在最新版本中合入针对 FSDP 与重计算协同场景的优化补丁，修复了在 checkpoint_wrapper 内调用 fully_shard 时重复触发 allgather 及显存未释放的问题。该优化通过控制 pre_forward 与 pre_backward 的执行逻辑，避免不必要的参数展开与内存残留。相关代码合入请参考：feat: torch fullyshard patch, optimize interaction between fullyshard and checkpoint_wrapper 建议与总结 1. FSDP 原理简述 FSDP（Fully Sharded Data Parallel）通过分片模型参数降低每张卡的显存占用来支持超大模型训练。每张卡持有 1/DP 卡数的参数，在需要参数计算时通过 allgather 收集这一层计算所需参数，每张卡参数需要的显存峰值为完整参数/层数。 2. 内存快照使用参考 PyTorch 内存快照工具的使用方式请参考：内存快照 - Ascend Extension for PyTorch 3. 最佳实践建议 ●在同时对一个 block 应用重计算和 fully_shard 时，建议先使能重计算，再进行 fully_shard，确保 fully_shard 包裹在 checkpoint_wrapper 外部。 ●使用设备内存快照工具定期检查显存使用情况，及时发现异常堆积点。 ●对于视觉模块等特殊结构，可考虑单独进行 fully_shard 或增加重计算逻辑，进一步优化显存占用。