本文，我们将了解如何基于 PyTorch 最新的
工作原理 📝
(图源: 链接)
上述工作流概述了 FSDP 的幕后流程。我们先来了解一下 DDP 是如何工作的，然后再看 FSDP 是如何改进它的。在 DDP 中，每个工作进程 (加速器 / GPU) 都会保留一份模型的所有参数、梯度和优化器状态的副本。每个工作进程会获取不同的数据，这些数据会经过前向传播，计算损失，然后再反向传播以生成梯度。接着，执行 all-reduce 操作，此时每个工作进程从其余工作进程获取梯度并取平均。这样一轮下来，每个工作进程上的梯度都是相同的，且都是全局梯度，接着优化器再用这些梯度来更新模型参数。我们可以看到，每个 GPU 上都保留完整副本会消耗大量的显存，这限制了该方法所能支持的 batch size 以及模型尺寸。
FSDP 通过让各数据并行工作进程分片存储优化器状态、梯度和模型参数来解决这个问题。进一步地，还可以通过将这些张量卸载到 CPU 内存来支持那些 GPU 显存容纳不下的大模型。在具体运行时，与 DDP 类似，FSDP 的每个工作进程获取不同的数据。在前向传播过程中，如果启用了 CPU 卸载，则首先将本地分片的参数搬到 GPU/加速器。然后，每个工作进程对给定的 FSDP 包装模块/层执行 all-gather 操作以获取所需的参数，执行计算，然后释放/清空其他工作进程的参数分片。在对所有 FSDP 模块全部执行该操作后就是计算损失，然后是后向传播。在后向传播期间，再次执行 all-gather 操作以获取给定 FSDP 模块所需的所有参数，执行计算以获得局部梯度，然后再次释放其他工作进程的分片。最后，使用 reduce-scatter 操作对局部梯度进行平均并将相应分片给对应的工作进程，该操作使得每个工作进程都可以更新其本地分片的参数。如果启用了 CPU 卸载的话，梯度会传给 CPU，以便直接在 CPU 上更新参数。
如欲深入了解 PyTorch FSDP 工作原理以及相关实验及其结果，请参阅 [7,8,9]。
问题
如果在 accelerate 中使用 PyTorch FSDP 时遇到任何问题，请提交至 accelerate。
但如果你的问题是跟 PyTorch FSDP 配置和部署有关的 - 你需要提交相应的问题至 PyTorch。
参考文献
[2] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
[3] DeepSpeed: Extreme-scale model training for everyone - Microsoft Research
[4] Megatron-LM: Training Multi-Billion Parameter Language Models Using
Model Parallelism
[5] Introducing GPipe, an Open Source Library for Efficiently Training Large-scale Neural Network Models
[6] Which hardware do you need to train a 176B parameters model?
[7] Introducing PyTorch Fully Sharded Data Parallel (FSDP) API | PyTorch
[10] Fit More and Train Faster With ZeRO via DeepSpeed and FairScale