【code】多GPU多卡并行如何操作

最新推荐文章于 2025-07-13 21:22:13 发布

原创最新推荐文章于 2025-07-13 21:22:13 发布 · 3.5k 阅读

5 ·

CC 4.0 BY-SA版权

本文详细解析了PyTorch中实现多GPU并行训练的正确方法，包括如何初始化模型和优化器，以及在多GPU环境下进行参数更新的具体步骤。同时，介绍了主GPU的设定方式和torch.cuda模块的高级用法。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

最初以为多卡并行仅仅只是一句把模型model以及optimize放在nn.DataParallel上就可以完成：

import torch.nn as nn

【1.返回模型】

self.bert_model = nn.DataParallel(self.bert_model, device_ids=args.gpu_ids, output_device=torch.cuda.current_device())

或【2.返回计算结果】

nn.parallel.data_parallel(module, inputs, device_ids=args.gpu_ids, output_device=torch.cuda.current_device(), dim=0, module_kwargs=None)

但并非是这样。如果单单使用这行命令会发现原本正确的程序开始报错。

正确的使用方法为：

model:

model = nn.DataParallel(model, device_ids=args.bert_gpu_ids) #, output_device=torch.cuda.current_device())

optimizer:

optimizer = torch.optim.Adam([param for _, param in param_optimizer],lr=args.lr)

optimizer = nn.DataParallel(optimizer, device_ids=args.bert_gpu_ids) #, output_device=torch.cuda.current_device())

for param_group in optimizer.module.param_groups:
param_group['lr'] = cur_lr

optimizer.module.step()

NOTE :

主 GPU 默认情况下是 0 号 GPU，也可以通过 torch.cuda.set_device(id) 来手动更改默认 GPU。
提供的多 GPU 并行列表中需要包含有主 GPU。
torch.cuda.set_device(args.device):设置当前设备。官网不鼓励使用此函数来设置。在大多数情况下，最好使用CUDA_VISIBLE_DEVICES环境变量。

https://pytorch-cn.readthedocs.io/zh/latest/package_references/torch-cuda/

eg:
os.environ['CUDA_VISIBLE_DEVICES'] = "5,6,7"

net = torch.nn.DataParallel(net, device_ids=[0,1,2])

参考博客：https://blog.youkuaiyun.com/qq_19598705/article/details/80396325