单机多卡训练

单机多卡训练,第1张

  • 启动方式的修改

单机单卡的启动

python run.py --model bert 

单机多卡的启动,2是卡的个数

python -m torch.distributed.launch --nproc_per_node 2 run.py --model bert

  • 卡的设置方式修改

上面改成分布式启动后,会自动传 local_rank 参数给程序,我们需要解析收到的 local_rank参数并进行设置

parser.add_argument("--local_rank", type=int) # 新增
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 修改前

self.device =  args.local_rank  #修改后
torch.cuda.set_device(args.local_rank)  # 新增
torch.distributed.init_process_group(backend="nccl", init_method='env://')  # 新增
  
  

  • 加载程序

修改前

model = x.Model(config).to(config.device)

修改后

model = x.Model(config).cuda()
model = torch.nn.parallel.DistributedDataParallel(model,device_ids=[torch.cuda.current_device()])

读取数据的时候是否增加 DistributedSampler 方式并不影响,设置并行会把数据分成 n 份,每个卡上加载单份数据进行训练。n 个卡之间如何交互,模型如何累计?需要研究一下。

原理:

  1. pytorch 把模型同步到各个 GPU
  2. 数据分成 n 份后,依次放入各个 GPU
  3. 每个GPU分别进行前项计算过程;
  4. 前向过程计算完后,pytorch再从各个GPU中收集计算后的结果,然后再按照次序将其拼接起来,计算loss。
    整个过程其实就是 同步模型参数→分别前向计算→计算损失→梯度反传

 

欢迎分享,转载请注明来源:内存溢出

原文地址: https://outofmemory.cn/langs/789173.html

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2022-05-05
下一篇 2022-05-05

发表评论

登录后才能评论

评论列表(0条)

保存