参考:
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
<