https://zhouyeyu.github.io/tags/fsdp2/ Recent content in FSDP2 on YuZhouye's Blog Hugo -- 0.155.0 en-us Thu, 18 Dec 2025 21:37:02 +0800 https://zhouyeyu.github.io/posts/pytorch-fsdp2-%E5%AF%B9%E6%AF%94-fsdp1%E7%9A%84%E5%8D%87%E7%BA%A7/ Thu, 18 Dec 2025 21:37:02 +0800 https://zhouyeyu.github.io/posts/pytorch-fsdp2-%E5%AF%B9%E6%AF%94-fsdp1%E7%9A%84%E5%8D%87%E7%BA%A7/ <h2 id="概述">概述</h2> <p>FSDP2 与FSDP1 保持了相同的并行思路，将参数、梯度等分块，在需要时all-gather聚集到一起，用完即刻丢弃。通过forward 和 backward hook自动处理相应的shard和unshard操作，节省显存，同时通过prefetch进一步掩盖通信和计算的耗时，提升性能。但是FSDP2和FSDP1也存在一些diff，甚至接口也不再向后兼容。</p> <h2 id="fsdp2梳理及更新点">FSDP2梳理及更新点</h2> <p>官方的描述是这样的：</p> <blockquote> <p>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相关的改动。</p>