单机多卡训练

• 启动方式的修改

单机单卡的启动

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