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相关的改动。

API接口

FSDP2的接口不是向后兼容的，相应的接口名称也从FSDP变成了fully_shard，以一个model为例，不再像FSDP1一样可以直接通过一个接口包裹整个模型，而转变成逐层进行模型的包裹。

for module in model.modules();
  if isinstance(module, TransformerBlock):
    fully_shard(module, **fsdp_kwargs)
fully_shard(model, **fsdp_kwargs)

原因是在FSDP2中，unit变成了最小的处理单位，unit可以是一个layer，一个nnModule，也可以是一整个模型，分块能更合理地去管理每个分块参数的生命周期。如图所示，每个颜色代表一个unit的参数。

Image Alt Text

相比之下，FSDP1的处理方式为：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
# defnition of model
......
model = FSDP(model)

分布式与DTensor

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.

DTensor是pytorch的一种分布式数据结构，是torch.Tensor的子类。DTensor通过Devicemesh描述了张量在哪个设备上，以什么划分方式分布。Placement参数描述了每个参数上的储存方式。FSDP2使用了DTensor作为参数的分布式表达，每个rank只持有完整Tensor的局部shard。 DTensor相关介绍见专栏文章 PyTorch DTensor 功能及原理源码解读 DTensor取代的是FSDP1 使用的FlatParameter参数切割方式，会把所有要分块的参数拼接成一个巨大的Tensor，根据要分块的大小对参数进行padding，之后按大小分配给各个rank，同时，梯度也会随参数切分而切分。

To shard the FlatParameter, FSDP divides it into equal-sized chunks, where the number of chunks equals the sharding factor, and assigns one chunk per rank.

Image Alt Text

FSDP1 FlatParam

FSDP的FlatParam分块导致某些Tensor可能会被划分到不同的rank，而某些Tensor完整的保留在某些rank中。基于DTensor的切分方式通常是沿Dim0切分，FSDP2采用这种切分方式保证每个rank都获得每个param的一部分参数，并且只需要对部份Tensor进行padding。

Image Alt Text

经过FSDP2的init_device_mesh，生成分布式的分块信息。DTensor在进行sharding操作时会保留这些分块信息，FSDP2中，在FSDPParam中进行维护。在进行通信操作，如All-reduce或reduce-scatter时，DTensor中储存的分块信息会被读取，以完成参数的shard和unshard操作。打印fsdp前后的参数信息，可以得知sharding前grad属于<class ’torch.Tensor’>，而sharding后转变成了<class ’torch.distributed.tensor.Tensor’> Stream、Event控制

Improving memory management system that achieves lower and deterministic GPU memory by avoiding recordStream (doc) and does so without any CPU synchronization. 通过避免recordStream，无需CPU同步降低了GPU显存占用

recordStream是torch用来管理CUDA张量生命周期的方法，因为GPU的操作是异步的，比如下发一个计算或者通信算子时，CPU操作会立即返回，而实际的数据计算或者通信仍然在GPU上进行，那么此时如果尝试读写这块正在操作的buffer，就有可能导致一些不可预知的错误。

CUDA tensor的生命周期是由GPU上内存的分配合释放管理的，recordStream会记录tensor被哪些流绑定，如果要释放一个tensor，就必须保证依赖该tensor的流全部执行完毕，而不是仅与操作它的流相关。

FSDP1的局限性： FSDP1的prefetch实现了在计算当前层参数时，预取下一层参数，以实现通信/计算的overlap，但是通常计算流和通信流是两条stream，且不同stream间是相互并行的。因此，存在一种情况：即如果某次计算由于耗时较久，久到连续下发的几个all-gather操作全部完成了，计算还没有完成，此时显存就会被通信后的参数不断占据，甚至会出现oom。

FSDP1对这个问题的解法也十分粗暴，即，设置了一个rate_limiter,即速率控制器，在每个TransformerBlock后阻塞CPU，rate_limiter中的限制默认为2，保证只有当上一层结束时，才会对下一层进行all-gather操作。源码：

class _FreeEventQueue:
    """
    This tracks all pending frees corresponding to inflight all-gathers. The
    queueing pattern is iterative enqueues with a single dequeue per iteration
    once the limit ``_max_num_inflight_all_gathers`` is reached.
    """

    def __init__(self) -> None:
        self._queue: collections.deque[torch.Event] = collections.deque()
        self._max_num_inflight_all_gathers = 2  # empirically chosen

关于FSDP1还有一个讨论，通过recordStream控制，假设存在两个流streamA和streamB，当streamA持有某个Tensor时，直到这条steamA 完全结束时，相应的Tensor才会被释放，因此，即使这个Tensor生命周期理论已经结束了，但streamA没有退出，此时如果在streamB上尝试malloc一块内存，cachingAllocator会判断此时没有空闲内存，而去重新分配空间。这就导致了1. 内存存在峰值 2. 由于反复malloc，性能变差。社区讨论的内存峰值如图：

Image Alt Text

基本原理如图所示：当在尝试del一个正在使用的tensor后，如果使用过该tensor的流还未结束，那么此时的malloc操作无法复用内存，如图上的malloc C，如果在流结束之后，如malloc E，才可以顺利复用A的内存。

Image Alt Text

改进的点如下图所示，同步不要再发生在CPU上了，而通过在stream后插入标记来记录操作是否完成，通信流会等待依赖的计算流，而计算流也会等待依赖的通信流。理解一下这张图，以i为例，计算layer i时，预取了i+1，i+1的通信会产生一个event i，计算流会去wait这个event，直到结束后才会去进行layer i+1的计算。

Image Alt Text