PyTorch中DistributedDataParallel使用笔记

1. 基本概念

在使用DistributedDataParallel时有一些概率必须掌握

多机多卡 含义
world_size 代表有几台机器,可以理解为几台服务器
rank 第几台机器,即第几个服务器
local_rank 某台机器中的第几块GPU
单机多卡 含义
world_size 代表机器一共有几块GPU
rank 第几块GPU
local_rank 第几块GPU,与rank相同

2. 使用方法

2.1. 修改主函数

在运行的时候,DistributedDataParallel会往你的程序中加入一个参数local_rank,所以要现在你的代码中解析这个参数,如:

python 复制代码
parser.add_argument("--local_rank", type=int, default=1, help="number of cpu threads to use during batch generation")

2.2. 初始化

python 复制代码
torch.distributed.init_process_group(backend="nccl")

os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2"  # 有几块GPU写多少

2.3. 设定device

python 复制代码
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
global device
device = torch.device("cuda", local_rank)

我没用arg.local_rank,新定义了一个local_rank变量,是因为我更信任distributed.get_rank()这个函数

这里用torch.device来写,并且加了global,是因为后面模型和数据都要用到这个device,不会出错

2.4. 模型加载到多gpu

python 复制代码
model.to(device)  # 这句不能少,最好不要用model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)  # 这句加载到多GPU上

2.5. 数据加载到gpu

python 复制代码
数据.to(device)

2.6. 启动

python 复制代码
torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:12345 train_cylinder_asym.py

参考文献

Pytorch并行计算(二): DistributedDataParallel介绍_dist.barrier_harry_tea的博客-CSDN博客

DistributedDataParallel多GPU分布式训练全过程总结 跟着做90%成功_BRiAq的博客-CSDN博客

相关推荐
多巴胺与内啡肽.1 小时前
深度学习--自然语言处理统计语言与神经语言模型
深度学习·语言模型·自然语言处理
深度之眼1 小时前
2025时间序列都有哪些创新点可做——总结篇
人工智能·深度学习·机器学习·时间序列
Sherlock Ma2 小时前
PDFMathTranslate:基于LLM的PDF文档翻译及双语对照的工具【使用教程】
人工智能·pytorch·语言模型·pdf·大模型·机器翻译·deepseek
不吃香菜?3 小时前
PyTorch 实现食物图像分类实战:从数据处理到模型训练
人工智能·深度学习
Light604 小时前
智启未来:深度解析Python Transformers库及其应用场景
开发语言·python·深度学习·自然语言处理·预训练模型·transformers库 |·|应用场景
数据智能老司机4 小时前
构建具备自主性的人工智能系统——在生成式人工智能系统中构建信任
深度学习·llm·aigc
谦行5 小时前
工欲善其事,必先利其器—— PyTorch 深度学习基础操作
pytorch·深度学习·ai编程
xwz小王子6 小时前
Nature Communications 面向形状可编程磁性软材料的数据驱动设计方法—基于随机设计探索与神经网络的协同优化框架
深度学习
生信碱移7 小时前
大语言模型时代,单细胞注释也需要集思广益(mLLMCelltype)
人工智能·经验分享·深度学习·语言模型·自然语言处理·数据挖掘·数据可视化
硅谷秋水8 小时前
通过模仿学习实现机器人灵巧操作:综述(上)
人工智能·深度学习·机器学习·计算机视觉·语言模型·机器人