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
<