PyTorch FSDP2 对比 FSDP1的升级
概述 FSDP2 与FSDP1 保持了相同的并行思路，将参数、梯度等分块，在需要时all-gather聚集到一起，用完即刻丢弃。通过forward 和 backward hook自动处理相应的shard和unshard操作，节省显存，同时通过prefetch进一步掩盖通信和计算的耗时，提升性能。但是FSDP2和FSDP1也存在一些diff，甚至接口也不再向后兼容。 FSDP2梳理及更新点 官方的描述是这样的： Comparing with FSDP1, FSDP2 has following advantages: Representing sharded parameters as DTensor sharded on dim-i, allowing for easy manipulation of individual parameters, communication-free sharded state dicts, and a simpler meta-device initialization flow. Improving memory management system that achieves lower and deterministic GPU memory by avoiding recordStream (doc) and does so without any CPU synchronization. Offering a tensor subclass extension point to customize the all-gather, e.g. for float8 all-gather for float8 linears (doc), and NF4 for QLoRA (doc) Mixing frozen and non-frozen parameters can in the same communication group without using extra memory. 主要梳理下DTensor以及recordStream相关的改动。 ...