PyTorch 17. GPU并发

多GPU运算的分发并行机制

torch.nn.DataParallel

功能：包装模型，实现分发并行机制

torch.nn.DataParallel(module, device_ids = None, output_device=None, dim=0)

主要参数：
module: 需要包装分发的模型
device_ids: 可分发的gpu，默认分发到所有可见可用gpu
output_device:结果输出设备

结合之前博客的多gpu训练和单gpu测试，进行查看，具体可见
这篇博客
缺点：

1. 在每个训练批次（batch）中，因为模型的权重都在一个进程上先算出来，然后再把他们分发到每个GPU上，所以网络通信就成为了一个瓶颈，而GPU使用率也通常很低。
2. 除此之外，nn.DataParallel需要所有的GPU都在一个节点上，且并不支持Apex的混合精度训练。

使用torch.distributed加速并行训练：

DataParallel: 单进程控制多GPU
DistributedDataParallel: 多进程控制多GPU，一起训练模型

和单进程训练不同的是，多进程训练需要注意一下事项：

1. 在喂数据的时候，一个batch被分到了多个进程，每个进程在取数据的时候要确保拿到的是不同的数据（DistributedSampler）
2. 要告诉每个进程自己是谁，使用哪块GPU（args.local_rank）
3. 在做BN的时候注意同步数据。

使用方式

在多进程的启动方面，我们无需自己手写multiprocess进行一系列复杂的CPU、GPU分配任务，PyTorch为我们提供了一个很方便的启动器torch.distributed.launch用于启动文件，所以我们运行训练代码的方式就变成这样：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

其中的 --nproc_per_node 参数用于指定为当前主机创建的进程数，由于我们是单机多卡，所以这里的node数量为1，这里只设置所使用的GPU数量即可。

初始化

在启动器为我们启动python脚本后，在执行过程中，启动器会将当前进行的index通过参数传递给python，我们可以这样获得当前进程的index：即通过参数local_rank来告诉我们当前进程使用的是哪个GPU，用于我们在每个进程中指定不同的device:

def parse():
	parser = argparse.ArgumentParser()
	parser.add_argument('--local_rank', type=int, default=0, help='node rank for distributed training')
	args = parser.parse_args()
	return args

def main():
	args = parse()
	torch.cuda.set_device(args.local_rank)
	torch.distributed.init_process_group(
		'nccl',
		init_method='env://'
	)
	device = torch.device(f'cuda:{args.local_rank}')

其中torch.distributed.init_process_group用于初始化GPU通信方式（NCLL）和参数的获取方式（env代表通过环境变量）。使用init_process_group设置GPU之间通信使用的后端和端口，通过NCCL实现GPU通信

Dataloader

在我们初始化data_loader的时候需要使用到torch.utils.data.distributed.DistributedSampler这个特性：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

这样就能给每个进程一个不同的sampler，告诉每个进程自己分别取哪些数据

模型的初始化

和nn.DataParallel的方式一样，

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

使用DistributedDataParallel包装模型, 它能帮助我们为不同GPU上求得的提取进行all reduce(即汇总不同GPU计算所得的梯度，并同步计算结果)。all reduce 后不同GPU中模型的梯度均为all reduce之前各GPU梯度的均值。

同步BN

现有的标准BN仅能在单卡上进行批归一化，跨卡同步BN可以对全局的样本进行归一化，apex中可以使用

多机多卡DDP（DistributedDataParallel）

进程组的相关概念
Group：进程组，大部分情况下DDP的各个进程是在同一个进程组下
world_size: 总的进程数量（原则上一个process占用一个GPU是较优的）
rank: 当前进程的序号，用于程序间通讯 ，rank=0的主机为master节点
local_rank: 当前进程对应的GPU号

DDP的基本用法（代码编写流程）

1. 使用torch.distributed.init_process_group初始化进程组
2. 使用torch.nn.parallel.DistributedDataParallel创建分布式模型
3. 使用torch.utils.data.distributed.DistributedSampler创建DataLoader
4. 调整其他必要的地方
5. 使用torch.distributed.launch / torch.multiprocessing 或slurm开始训练

使用apex加速（混合精度训练，并行训练，同步BN）：

APEX是来自英伟达的一个深度学习加速库。由英伟达开源，完美支持PyTorch框架，用于改变数据格式来减小模型显存占用的工具。其中amp（Automatic Mixed Precision），将模型的大部分操作都用Float16数据类型测试，一些特别操作仍然使用Float32。并且用户仅仅通过三行代码即可完美将自己的训练代码迁移到该模型。

apex的使用

Amp:Automatic Mixed Precision

apex.amp是一种通过仅更改脚本的3行来启用混合精度训练的工具，通过向amp.initialize提供不同的flags，用户可以轻松地实验不同的纯精度和混合精度训练模式

Distributed Training

apex.parallel.DistributedDataParallel是一个模块包装器，类似于torch.nn.parallel.DistributedDataParallel，它支持方便的多进程分布式训练，针对NVIDIA的NCCL通信库进行了优化。

Synchronized Batch Normalization

apex.parallel.SyncBatchNorm扩展了torch.nn.modules.batchnorm.__BatchNorm以支持同步BN。它在多进程训练期间减少了跨进程的统计数据。同步BN已用于每个GPU上只能容纳一个小的本地minibatch的情况。

import apex
sync_bn_model = apex.parallel.convert_syncbn_model(model)

Checkpointing

为了正确保存和加载读者的amp训练，amp引入了amp.state_dict()，其中包含所有loss_scalers及其相应的未跳过的步骤，以及amp.load_state_dict()以恢复属性。

# Initialization
opt_level = 'O1'
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)

# Train your model
...
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
...

# Save checkpoint
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'amp': amp.state_dict()
}
torch.save(checkpoint, 'amp_checkpoint.pt')
...

# Restore
model = ...
optimizer = ...
checkpoint = torch.load('amp_checkpoint.pt')

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
amp.load_state_dict(checkpoint['amp'])

# Continue training

注意

1. 保存和加载模型参数时要多留意保存和加载amp.state_dict()
2. model和optimizer都要使用amp.initialize()进行包装

用户指定数据格式

amp.initialize(net, opt, opt_level="O1")

其中的opt-level参数是用户指定采用何种数据格式做训练的

1. O0：采用纯FP32训练
2. O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16还是FP32进行计算
3. O2：混合精度训练，不存在黑白名单，除了BN，几乎都用FP16计算
4. O3：纯FP16训练

溢出问题
因为Float16保存数据位数变少了，能保存数据的上限和下限的绝对值也小了。在处理求和的操作，如sigmoid，softmax等，会导致数据溢出，得到错误结果，针对这些操作，我们想使用float32作为数据格式，我们仅需在模型定义中，在构造函数__init__()中加入以下待：

from apex import amp
class xxxNet(Module):
	def __init__(using_map=False)
		...
		...
		if using_amp:
		     amp.register_float_function(torch, 'sigmoid')
		     amp.register_float_function(torch, 'softmax')

和register_float_function相似的注册函数有：

1. amp.register_half_function（module,function_name）
2. amp.register_float_function (module, function_name)
3. amp.register_promote_function (module, function_name)