一直以来都是用的单机单卡训练模型,虽然很多情况下已经足够了,但总有一些情况得上分布式训练:
- 模型大到一张卡放不下;
- 单张卡batch size不敢设太大,训练速度慢;
- 当你有好几张卡,不想浪费;
- 展示一下技术。
由于还没遇到过一张显卡放不下整个模型的情况,本文的分布式训练仅限数据并行。主要从数据并行的原理和一些简单的实践例子进行说明。
文章目录
原理介绍
与每个step一个batch数据相比,数据并行是指每个step用更多的数据(多个batch)进行计算------即多个batch的数据并行进行前向计算。既然是并行,那么就涉及到多张卡一起计算。单卡和多卡训练过程如下图^1^所示,主要有三个过程:
- 各卡分别计算损失和梯度,即图中红线部分;
- 所以梯度整合到主device,即图中蓝线部分;
- 主device进行参数更新,并将新模型拷贝到其他device上,即图中绿线部分。
左图是单GPU训练;右图是多GPU训练的一个变体:(1)计算损失和梯度,(2)所有梯度聚合在一个GPU上,(3)发生参数更新,并将参数重新广播给所有GPU
如果不使用数据并行,在显存足够的情况下,我们可以将batch_size设大,这和数据并行的区别在哪呢?如果只将batch_size设大,计算还是在一张卡上完成,速度相对来说是不如将数据均分后放在不同卡上并行计算的。当然,考虑到卡之间的通信问题,要发挥多卡并行的力量需要进行一定权衡。
torch中主要有两种数据并行方式:DP和DDP。
DataParallel
DP是较简单的一种数据并行方式,直接将模型复制到多个GPU上并行计算,每个GPU计算batch中的一部分数据,各自完成前向和反向后,将梯度汇总到主GPU上。其基本流程:
- 加载模型、数据至内存;
- 创建DP模型;
- DP模型的forward过程:
- 一个batch的数据均分到不同device上;
- 为每个device复制一份模型;
- 至此,每个device上有模型和一份数据,并行进行前向传播;
- 收集各个device上的输出;
- 每个device上的模型反向传播后,收集梯度到主device上,更新主device上的模型,将模型广播到其他device上;
- 3-4循环。
在DP中,只有一个主进程,主进程下有多个线程,每个线程管理一个device的训练。因此,DP中内存中只存在一份数据,各个线程间是共享这份数据的。DP和Parameter Server的方式很像。
小样
python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
# 假设我们有一个简单的数据集类
class SimpleDataset(Dataset):
def __init__(self, data, target):
self.data = data
self.target = target
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx], self.target[idx]
# 假设我们有一个简单的神经网络模型
class SimpleModel(nn.Module):
def __init__(self, input_dim):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(input_dim, 1)
def forward(self, x):
return torch.sigmoid(self.fc(x))
# 假设我们有一些数据
n_sample = 100
n_dim = 10
batch_size = 10
X = torch.randn(n_sample, n_dim)
Y = torch.randint(0, 2, (n_sample, )).float()
dataset = SimpleDataset(X, Y)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# ===== 注意:刚创建的模型是在 cpu 上的 ===== #
device_ids = [0, 1, 2]
model = SimpleModel(n_dim).to(device_ids[0])
model = nn.DataParallel(model, device_ids=device_ids)
optimizer = optim.SGD(model.parameters(), lr=0.01)
for epoch in range(10):
for batch_idx, (inputs, targets) in enumerate(data_loader):
inputs, targets = inputs.to('cuda'), targets.to('cuda')
outputs = model(inputs)
loss = nn.BCELoss()(outputs, targets.unsqueeze(1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item()}')
其中最重要的一行便是:
python
model = nn.DataParallel(model, device_ids=device_ids)
注意,模型的参数和缓冲区都要放在device_ids[0]
上。在执行forward
函数时,模型会被复制到各个GPU上,对模型的属性进行更新并不会产生效果,因为前向完后各个卡上的模型就被销毁了。只有在device_ids[0]
上对模型的参数或者buffer进行的更新才会生效! ^2^
DistributedDataParallel
DDP,顾名思义,即分布式的数据并行,每个进程独立进行训练,每个进程会加载完整的数据,但是读取不重叠的数据。DDP执行流程^3^:
-
准备阶段
- 环境初始化
- 在各张卡上初始化进程并建立进程间通信,对应代码:
init_process_group
。
- 在各张卡上初始化进程并建立进程间通信,对应代码:
- 模型广播
- 将模型parameter、buffer广播到各节点,对应代码:
model = DDP(model).to(local_rank)
。
- 将模型parameter、buffer广播到各节点,对应代码:
- 创建管理器reducer,给每个参数注册梯度平均hook。
- 环境初始化
-
准备数据
- 加载数据集,创建适用于分布式场景的数据采样器,以防不同节点使用的数据不重叠。
-
训练阶段
- 前向传播
- 同步各进程状态(parameter和buffer);
- 当DDP参数
find_unused_parameter
为true
时,其会在forward
结束时,启动一个回溯,标记未用到的参数,提前将这些设置为ready
。
- 计算梯度
- reducer外面:各进程各自开始反向计算梯度;
- reducer外面:当某个参数的梯度计算好了,其之前注册的grad hook就会触发,在reducer里把这个参数的状态标记为
ready
; - reducer里面:当某个bucket的所有参数都是
ready
时,reducer开始对这个bucket的所有参数开始一个异步的all-reduce梯度平均操作; - reducer里面:当所有bucket的梯度平均都结束后,reducer把得到的平均梯度正式写入到parameter.grad里。
- 优化器应用梯度更新参数。
- 前向传播
小样
python
import argparse
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 1. 基础模块 ###
class SimpleModel(nn.Module):
def __init__(self, input_dim):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(input_dim, 1)
cnt = torch.tensor(0)
self.register_buffer('cnt', cnt)
def forward(self, x):
self.cnt += 1
# print("In forward: ", self.cnt, "Rank: ", self.fc.weight.device)
return torch.sigmoid(self.fc(x))
class SimpleDataset(Dataset):
def __init__(self, data, target):
self.data = data
self.target = target
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx], self.target[idx]
# 2. 初始化我们的模型、数据、各种配置 ####
## DDP:从外部得到local_rank参数。从外面得到local_rank参数,在调用DDP的时候,其会自动给出这个参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank
## DDP:DDP backend初始化
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')
## 假设我们有一些数据
n_sample = 100
n_dim = 10
batch_size = 25
X = torch.randn(n_sample, n_dim) # 100个样本,每个样本有10个特征
Y = torch.randint(0, 2, (n_sample, )).float()
dataset = SimpleDataset(X, Y)
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
data_loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)
## 构造模型
model = SimpleModel(n_dim).to(local_rank)
## DDP: Load模型要在构造DDP模型之前,且只需要在master上加载就行了。
ckpt_path = None
if dist.get_rank() == 0 and ckpt_path is not None:
model.load_state_dict(torch.load(ckpt_path))
## DDP: 构造DDP model ------------------ 必须在 init_process_group 之后才可以调用 DDP
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
## DDP: 要在构造DDP model之后,才能用model初始化optimizer。
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_func = nn.BCELoss().to(local_rank)
# 3. 网络训练 ###
model.train()
num_epoch = 100
iterator = tqdm(range(100))
for epoch in iterator:
# DDP:设置sampler的epoch,
# DistributedSampler需要这个来指定shuffle方式,
# 通过维持各个进程之间的相同随机数种子使不同进程能获得同样的shuffle效果。
data_loader.sampler.set_epoch(epoch)
# 后面这部分,则与原来完全一致了。
for data, label in data_loader:
data, label = data.to(local_rank), label.to(local_rank)
optimizer.zero_grad()
prediction = model(data)
loss = loss_func(prediction, label.unsqueeze(1))
loss.backward()
iterator.desc = "loss = %0.3f" % loss
optimizer.step()
# DDP:
# 1. save模型的时候,和DP模式一样,有一个需要注意的点:保存的是model.module而不是model。
# 因为model其实是DDP model,参数是被`model=DDP(model)`包起来的。
# 2. 只需要在进程0上保存一次就行了,避免多次保存重复的东西。
if dist.get_rank() == 0 and epoch == num_epoch - 1:
torch.save(model.module.state_dict(), "%d.ckpt" % epoch)
结合上面的代码,一个简化版的DDP流程:
- 读取DDP相关的配置,其中最关键的就是:
local_rank
; - DDP后端初始化:
dist.init_process_group
; - 创建DDP模型,以及数据加载器。注意要为加载器创建分布式采样器(
DistributedSampler
); - 训练。
DDP的通常启动方式:
CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node 2 ddp.py
一些概念
以上过程中涉及到一些陌生的概念,其实走一遍DDP的过程就会很好理解:每个进程是一个独立的训练流程,不同进程之间共享同一份数据。为了避免不同进程使用重复的数据训练,以及训练后同步梯度,进程间需要同步。因此,其中一个重点就是每个进程序号,或者说使用的GPU的序号。
node
:节点,可以是物理主机,也可以是容器;rank
和local_rank
:都表示进程在整个分布式任务中的编号。rank
是进程在全局的编号,local_rank
是进程在所在节点上的编号。显然,如果只有一个节点,那么二者是相等的。在启动脚本中的--nproc_per_node
即指定一个节点上有多少进程;world_size
:即整个分布式任务中进程的数量。
DDP与DP的区别
- DP是单进程多线程的,只能在单机上工作;DDP是多进程的,可以在多级多卡上工作。DP通常比DDP慢,主要原因有:1)DP是单进程的,受到GIL的限制;2)DP每个step都需要拷贝模型,以及划分数据和收集输出;
- DDP可以与模型并行相结合;
- DP的通信成本随着卡数线性增长,DDP支持Ring-AllReduce,通信成本是固定的。
本文利用pytorch进行数据并行训练进行了一个粗浅的介绍。包括DP和DDP的基本原理,以及简单的例子。实际在分布式过程中涉及到的东西还是挺多的,比如DP/DDP中梯度的回收是如何进行的,DDP中数据采样的细节,DDP中的数据同步操作等。更多的还是要基于真实的需求出发才能真的体会得到。