PyTorch 分布式训练模式 FSDP 设计分析

在 AI 大模型训练场景中,数据是海量的,模型也是超大的,对于训练大模型会带来很大挑战,比如对算力的需求,对处理大模型的工程复杂度,等等。PyTorch 给出了一种实现方式——FSDP(Fully Sharded Data Parallel),它提供了易用的 API,可以非常方便地解决大模型分布式训练的难题。 FSDP 是在 DDP(DistributedDataParallel)的基础上提出的,首先我们了解一下 PyTorch 的 DDP(DistributedDataParallel) 训练模式的一些特点: 在 DDP 中,核心的能力还是训练数据并行(Data Parallel)。以多机多卡方式为例,每个 process/worker 都会持有模型的一个副本(Replica),通过使每个 process/worker 处理一个 batch 的数据试下并行处理,最后使用 all-reduce 操作对多个不同 process/worker 计算得到的梯度进行累加求和;接着,再将优化器状态、梯度通过跨多个 process/worker 进行复制,使得每个 process/worker 上的模型参数都得到同步更新。也就是说,在 DDP 中并没有实现对模型参数的分片管理,即模型并行(Model Parallel)。 在 FSDP 中实现了模型的分片管理能