pytorch FSDP分布式训练minist案例

参考:
https://pytorch.ac.cn/tutorials/intermediate/FSDP_tutorial.html

http://admin.guyuehome.com/39789/1000

FSDP（Fully Sharded Data Parallel）—类似实现zero系列

FSDP是PyTorch
1.11版本引入的一种新的数据并行策略。它的主要目标是解决在多GPU和多节点环境下训练大型模型的内存效率问题。FSDP通过将模型的参数、梯度和优化器状态分片到多个GPU上，从而允许在有限的内存资源下训练更大的模型。

主要特点：

参数分片：将模型的参数分片到多个GPU上，每个GPU只保存一部分参数。

梯度分片：梯度也会被分片，每个GPU只计算和存储与其参数分片对应的梯度。

优化器状态分片：优化器的状态也会被分片，每个GPU只维护与其参数分片对应的优化器状态。

通信优化：通过减少通信量和优化通信模式，提高训练效率。

运行版本：
torch 2.4.0

代码：
python DDP_mnist.py

<