DataParallel（DP）：Parameter Server模式，一张卡位reducer，实现也超级简单，一行代码。 有个不能接受的缺陷是：DataParallel是基于Parameter server的算法，所有的loss都在主卡上计算，负载不均衡的问题比较严重，有时在模型较大的时候（比如bert-large），主卡占满了，其他的卡一半都占不到，及其浪费资源。 示例代码：
# coding=utf-8 import torch import torch.nn as nn from torch.autograd import Variable from torch.utils.data import Dataset, DataLoader class RandomDataset(Dataset): def __init__(self, size, length): self.len = length self.data = torch.randn(length, size) def __getitem__(self, index): return self.data[index] def __len__(self): return self.len class Model(nn.Module): def __init__(self, input_size, output_size): super(Model, self).__init__() self.fc = nn.Linear(input_size, output_size) def forward(self, input): output = self.fc(input) return output input_size = 5 output_size = 2 batch_size = 30 data_size = 30 dataset = RandomDataset(input_size, data_size) rand_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True) model = Model(input_size, output_size) if torch.cuda.is_available(): model.cuda() if torch.cuda.device_count() > 1: model = nn.DataParallel(model) # 关键代码 for data in rand_loader: if torch.cuda.is_available(): input_var = Variable(data.cuda()) else: input_var = Variable(data) output = model(input_var)是的，你没有看错，这个函数是为了分布式训练设计的。但是，即使在单机多卡上，官方也建议使用新的DistributedDataParallel，采用all-reduce算法。
（1）初始化后端
torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend='nccl', init_method='env://')（2）模型并行化 这里也很简单，使用DistributedDataParallel函数warp一下就可以：
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)（3）数据并行 这里需要注意，如果指定了sampler，则shuffle=False，其中DataLoader的num_worker是每一个卡独立设置。
dataset = RandomDataset(input_size, data_size) sampler = torch.utils.data.distributed.DistributedSampler(dataset) rand_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False, sampler=sampler)（4）启动脚本
python -m torch.distributed.launch --nproc_per_node=8 train_face_troch.py完整代码示例：
# coding=utf-8 import torch import torch.distributed import torch.nn as nn from torch.autograd import Variable from torch.utils.data import Dataset, DataLoader import apex import argparse class RandomDataset(Dataset): def __init__(self, size, length): self.len = length self.data = torch.randn(length, size) self.label = torch.mean(self.data, dim=-1) def __getitem__(self, index): return self.data[index], self.label[index] def __len__(self): return self.len class Model(nn.Module): def __init__(self, input_size, output_size): super(Model, self).__init__() self.fc = nn.Linear(input_size, output_size) def forward(self, input): output = self.fc(input) return output def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('--local_rank', default=0, type=int) args = parser.parse_args() return args input_size = 5 output_size = 2 batch_size = 30 data_size = 30 args = parse_args() local_rank = args.local_rank torch.cuda.set_device(local_rank) # 设定cuda的默认GPU，每个rank不同 torch.distributed.init_process_group(backend='nccl', init_method='env://') dataset = RandomDataset(input_size, data_size) sampler = torch.utils.data.distributed.DistributedSampler(dataset) rand_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False, sampler=sampler) model = Model(input_size, output_size) if torch.cuda.is_available(): model.cuda() model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank) optimizer = torch.optim.Adam(model.parameters()) criterion = torch.nn.CrossEntropyLoss() # if torch.cuda.device_count() > 1: # model = nn.DataParallel(model) for data, label in rand_loader: data = data.cuda() label = label.cuda() output = model(data) loss = criterion(output, label) optimizer.zero_grad() loss.backward() optimizer.step()大规模数据训练时，混合精度训练时必须的，这速度谁用谁知道。 参考： https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed
这里主要需要改两个地方，一个是amp.initialize这个函数，一个是backward。
# coding=utf-8 import torch import torch.distributed import torch.nn as nn from torch.autograd import Variable from torch.utils.data import Dataset, DataLoader import apex import argparse class RandomDataset(Dataset): def __init__(self, size, length): self.len = length self.data = torch.randn(length, size) self.label = torch.mean(self.data, dim=-1) def __getitem__(self, index): return self.data[index], self.label[index] def __len__(self): return self.len class Model(nn.Module): def __init__(self, input_size, output_size): super(Model, self).__init__() self.fc = nn.Linear(input_size, output_size) def forward(self, input): output = self.fc(input) return output def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('--local_rank', default=0, type=int) args = parser.parse_args() return args input_size = 5 output_size = 2 batch_size = 30 data_size = 30 args = parse_args() local_rank = args.local_rank # 初始化 torch.cuda.set_device(local_rank) torch.distributed.init_process_group(backend='nccl', init_method='env://') dataset = RandomDataset(input_size, data_size) sampler = torch.utils.data.distributed.DistributedSampler(dataset) rand_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False, sampler=sampler) model = Model(input_size, output_size) model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank) optimizer = torch.optim.Adam(model.parameters()) model, optimizer = amp.initialize(model, optimizer, opt_level='O1') # 这里是字母O criterion = torch.nn.CrossEntropyLoss() if torch.cuda.is_available(): model.cuda() if torch.cuda.device_count() > 1: model = nn.DataParallel(model) for data, label in rand_loader: data = data.cuda() label = label.cuda() output = model(data) loss = criterion(output, label) optimizer.zero_grad() #loss.backward() with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() optimizer.step()