pytorch模型训练之fp16、apm、多GPU模型、梯度检查点（gradient checkpointing）显存优化等

本章内容分四个部分讲，fp16、apm以及pytorch的多gpu训练模式、gradient checkpointing显存优化。本节内容基于
pytorch==1.2.0，
transformers==3.0.2
python==3.6
pytorch 1.6+的版本有自带amp模式，这里暂时先不讨论，后期再做补充。

一、fp16和fp32

先说说fp16和fp32，当前的深度学习框架大都采用的都是fp32来进行权重参数的存储，比如Python float的类型为双精度浮点数fp64，PyTorch Tensor的默认类型为单精度浮点数fp32。随着模型越来越大，加速训练模型的需求就产生了。在深度学习模型中使用fp32主要存在几个问题，第一模型尺寸大，训练的时候对显卡的显存要求高；第二模型训练速度慢；第三模型推理速度慢。其解决方案就是使用低精度计算对模型进行优化。

• 推理过程中，模型优化目前比较成熟的方案就是fp16量化和int8量化，NVIDIA TensorRT等框架就可以支持，这里不再赘述。训练方面的方案就是混合精度训练，它的基本思想很简单: 精度减半(fp32→ fp16) ，训练时间减半。与单精度浮点数float32（32bit，4个字节）相比，半精度浮点数float16仅有16bit，2个字节组成。可以很明显的看到，使用fp16可以解决或者缓解上面fp32的两个问题：显存占用更少：通用的模型fp16占用的内存只需原来的一半，训练的时候可以使用更大的batchsize。计算速度更快：有论文指出半精度的计算吞吐量可以是单精度的 2-8 倍。

• 训练过程中，如果直接使用半精度进行计算会导致的两个问题的处理：舍入误差(Rounding Error)和溢出错误(Grad Overflow / Underflow)。
  舍入误差：float16的最大舍入误差约为 (~2 ^-10)，比float32的最大舍入误差(~2 ^-23) 要大不少。 对足够小的浮点数执行的任何操作都会将该值四舍五入到零，在反向传播中很多甚至大多数梯度更新值都非常小，但不为零。 在反向传播中舍入误差累积可以把这些数字变成0或者nan， 这会导致不准确的梯度更新，影响网络的收敛。
  溢出错误：由于float16的有效的动态范围约为 (5.96×10^-8 ~ 6.55×10^4)，比单精度的float32(1.4x10^-45 ~ 1.7x10^38)要狭窄很多，精度下降(小数点后16相比较小数点后8位要精确的多)会导致得到的值大于或者小于fp16的有效动态范围，也就是上溢出或者下溢出。在深度学习中，由于激活函数的的梯度往往要比权重梯度小，更易出现下溢出的情况。2018年ICLR论文Mixed Precision Training中提到，简单的在每个地方使用FP16会损失掉梯度更新小于2^-24的值——大约占他们的示例网络所有梯度更新的5%。

解决方案就是使用混合精度训练（Mixed Precision）和损失缩放(Loss Scaling)：

1. 混合精度训练： 混合精度训练是一种通过在FP16上执行尽可能多的操作来大幅度减少神经网络训练时间的技术，在像线性层或是卷积操作上，FP16运算较快，但像Reduction运算又需要 FP32的动态范围。通过混合精度训练的方式，便可以在部分运算操作使用FP16，另一部分则使用 FP32，混合精度功能会尝试为每个运算使用相匹配的数据类型，在内存中用FP16做储存和乘法从而加速计算，用FP32做累加避免舍入误差。这样在权重更新的时候就不会出现舍入误差导致更新失败，混合精度训练的策略有效地缓解了舍入误差的问题。
2. 损失缩放： 即使用了混合精度训练，还是会存在无法收敛的情况，原因是激活梯度的值太小，造成了下溢出。损失缩放是指在执行反向传播之前，将损失函数的输出乘以某个标量数(论文建议从8开始)。 乘性增加的损失值产生乘性增加的梯度更新值，提升许多梯度更新值到超过FP16的安全阈值2^-24。 只要确保在应用梯度更新之前撤消缩放，并且不要选择一个太大的缩放以至于产生inf权重更新(上溢出) ，从而导致网络向相反的方向发散。

二、基于apex的apm（Apex混合精度加速）

AMP(自动混合精度）的关键词有两个：自动，混合精度。
自动：Tensor的dtype类型会自动变化，框架按需自动调整tensor的dtype,当然有些地方还需手动干预。
混合精度：采用不止一种精度的Tensor，torch.FloatTensor和torch.HalfTensor

首先一个默认的原始训练流程为：

import torch
model = torch.nn.Linear(D_in, D_out)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for img, label in dataloader:
	out = model(img)
	loss = LOSS(out, label)
	loss.backward()
	optimizer.step()
	optimizer.zero_grad()

这种状态下，我测试的模型的单卡显存占用为：8466MB，速度为18s/40iter 如果想使用半精度的训练，可以直接将流程换为：

import torch
model = torch.nn.Linear(D_in, D_out).half()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for img, label in dataloader:
	out = model(img.half())
	loss = LOSS(out, label)
	loss.backward()
	optimizer.step()
	optimizer.zero_grad()

即将模型和输入的数据转为半精度即可。这样将达到最高的速度和最小的模型体积。采用该方法，我的测试模型单卡显存占用为：4978MB，速度为34s/40iter，对比这两种情况，显存占用的确是大大降低了，之前模型训练需要的显存位8466*4 = 33864MB，采用半精度训练后显存占用降低为：4978*4 = 11948MB。

接下来是混合精度的实现，这里主要用到Apex的amp工具。代码修改为：

import torch
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")

for img, label in dataloader:
	out = model(img)
	loss = LOSS(out, label)
	# loss.backward() 	
        with amp.scale_loss(loss, optimizer) as scaled_loss:
    	     scaled_loss.backward()
	optimizer.step()
	optimizer.zero_grad()

实际流程为：调用 amp.initialize 按照预定的 opt_level 对 model 和 optimizer 进行设置。在计算 loss 时使用 amp.scale_loss 进行回传。需要注意以下几点：

1. 在调用 amp.initialize 之前，模型需要放在 GPU 上，也就是需要调用 cuda() 或者 to()。
2. 在调用 amp.initialize 之前，模型不能调用任何分布式设置函数。
3. 此时输入数据不需要再转换为半精度

下表展示了不同 opt_level 设置的差异：

设置编号	00	01	02	03
cast_model_type	torch.float32	None	torch.float16	torch.float16
patch_torch_functions	False	True	False	False
keep_batchnorm_fp32	None	None	True	False
master_weights	False	None	True	False
loss_scale	1.0	“dynamic”	“dynamic”	1.0

概括起来：00相当于原始的单精度训练。01在大部分计算时采用半精度，但是所有的模型参数依然保持单精度，对于少数单精度较好的计算（如softmax）依然保持单精度。02相比于01，将模型参数也变为半精度。03基本等于最开始实验的全半精度的运算。值得一提的是，不论在优化过程中，模型是否采用半精度，保存下来的模型均为单精度模型，能够保证模型在其他应用中的正常使用。这也是Apex的一大卖点。

实际对比使用中01和02设置单卡分别使用的显存量为01:5402MiB，02:5426MiB，基本不分上下。时间和半精度持平，为34s/40iter。

三、多GPU模式

pytorch提供了两种多gpu训练的方式，方案一：利用nn.DataParallel实现，实现简单，不涉及多进程，另一种是用采用分布式并行训练torch.nn.parallel.DistributedDataParallel和torch.utils.data.distributed.DistributedSampler结合多进程实现。第二种方式效率更高，但是实现起来稍难，第二种方式同时支持多节点分布式实现。方案二的效率要比方案一高，即使是在单运算节点上。下面详细说明下两者的区别。

方案一：nn.DataParallel 实现

核心在于使用nn.DataParallel将模型wrap一下，代码其他地方不需要做任何更改:

model = nn.DataParallel(model) ## 只需要这一句话即可

为方便说明，我们假设模型输入为(32, 768)，这里的 32 表示batch_size，模型输出为(32, 768)，使用 4 个GPU训练。nn.DataParallel起到的作用是将这 32 个样本拆成 4 份，发送给 4 个GPU 分别做 forward，然后生成 4 个大小为(8, 768)的输出，然后再将这 4 个输出都收集到cuda:0上并合并成(32, 768)。

可以看出，nn.DataParallel没有改变模型的输入输出，因此其他部分的代码不需要做任何更改，非常方便。但弊端是，后续的loss计算只会在cuda:0上进行，没法并行，因此会导致负载不均衡的问题。

方案二：分布式数据并行(distributed data parallel)

分布式数据并行(distributed data parallel)，是通过多进程实现的，相比与方案一要复杂很多。可以从以下几个方面理解：

1. 从一开始就会启动多个进程（进程数等于 GPU 数），每个进程独享一个 GPU，每个进程都会独立地执行代码。这意味着每个进程都独立地初始化模型、训练，当然，在每次迭代过程中会通过进程间通信共享梯度，整合梯度，然后独立地更新参数。
2. 每个进程都会初始化一份训练数据集，当然它们会使用数据集中的不同记录做训练，这相当于同样的模型喂进去不同的数据做训练，也就是所谓的数据并行。这是通过torch.utils.data.distributed.DistributedSampler函数实现的，不过逻辑上也不难想到，只要做一下数据partition，不同进程拿到不同的parition就可以了，官方有一个简单的demo，感兴趣的可以看一下代码实现：Distributed Training
3. 进程通过local_rank变量来标识自己，local_rank为0的为master，其他是slave。这个变量是torch.distributed包帮我们创建的，使用方法如下：

import argparse  # 必须引入 argparse 包
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()

必须以如下方式运行代码：

python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 train.py ## 单机 2 GPU                                   

这样的话，torch.distributed.launch就以命令行参数的方式将args.local_rank变量注入到每个进程中，每个进程得到的变量值都不相同。比如使用 4 个GPU的话，则 4 个进程获得的args.local_rank值分别为0、1、2、3。

上述命令行参数nproc_per_node表示每个节点（节点就代表机器）需要创建多少个进程(使用几个GPU就创建几个)；nnodes表示使用几个节点，因为我们是做单机多核训练，所以设为1。

4. 因为每个进程都会初始化一份模型，为保证模型初始化过程中生成的随机权重相同，需要设置随机种子。

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

使用方法通过如下示意代码展示：

from torch.utils.data.distributed import DistributedSampler  # 负责分布式dataloader创建，也就是实现上面提到的partition。

# 负责创建 args.local_rank 变量，并接受 torch.distributed.launch 注入的值
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()

# 每个进程根据自己的local_rank设置应该使用的GPU
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)

# 初始化分布式环境，主要用来帮助进程间通信
torch.distributed.init_process_group(backend='nccl')

# 固定随机种子
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# 初始化模型
model = Net()
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 只 master 进程做 logging，否则输出会很乱
if args.local_rank == 0:
    tb_writer = SummaryWriter(comment='ddp-training')

# 分布式数据集
train_sampler = DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)  # 注意这里的batch_size是每个GPU上的batch_size

# 分布式模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)

详细代码参考：ddp_train.py

ddp有用的技巧：torch.distributed.barrier

在读huggingface/transformers中的源码，比如examples/run_ner.py会看到一下代码：

# Load pretrained model and tokenizer
if args.local_rank not in [-1, 0]:
    torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab                          

args.model_type = args.model_type.lower()
config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                      num_labels=num_labels,
                                      cache_dir=args.cache_dir if args.cache_dir else None)
tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
                                            do_lower_case=args.do_lower_case,
                                            cache_dir=args.cache_dir if args.cache_dir else None)
model = model_class.from_pretrained(args.model_name_or_path,
                                    from_tf=bool(".ckpt" in args.model_name_or_path),
                                    config=config,
                                    cache_dir=args.cache_dir if args.cache_dir else None)

if args.local_rank == 0:
    torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

上述代码要实现预训练模型的下载和读入内存，如果4个进程都分别下载一遍显然是不合理的，那如何才能实现只让一个进程下载呢？这个时候就可以使用barrier函数。当slave进程（local_rank!=0）运行到第一个if时就被barrier住了，只能等着，但master进程可以往下运行完成模型的下载和读入内存，但在第二个if语句时遇到barrier，那会不会被barrier住呢？答案是不会，因为master进程和slave进程集合在一起了(barrier)，barrier会被解除，这样大家都往下执行。当然这时大家执行的进度不同，master进程已经执行过模型读入，所以从第二个if往下执行，而slave进程尚未执行模型读入，只会从第一个if往下执行。

可以看到barrier类似一个路障，进程会被拦住，直到所有进程都集合齐了才放行。适合这样的场景：只一个进程下载，其他进程可以使用下载好的文件；只一个进程预处理数据，其他进程使用预处理且cache好的数据等。

模型保存

模型的保存与加载，与单GPU的方式有所不同。这里通通将参数以cpu的方式save进存储, 因为如果是保存的GPU上参数，pth文件中会记录参数属于的GPU号，则加载时会加载到相应的GPU上，这样就会导致如果你GPU数目不够时会在加载模型时报错，像下面这样：

RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

模型保存都是一致的，不过时刻记住DDP中你有多个进程在同时跑，所以会保存多个模型到存储上，如果使用共享存储就要注意文件名的问题，当然一般只在rank0进程上保存参数即可，因为所有进程的模型参数是同步的。

torch.save(model.module.cpu().state_dict(), "model.pth")      

模型的加载

model=torch.load("model.pth")

以下是huggingface/transformers代码中用到的模型保存代码：

## 只保存 rank0 进程上模型参数
if torch.distributed.get_rank() == 0:
    model_to_save = model.module if hasattr(model, "module") else model  # Take care of distributed/parallel training
    model_to_save.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

四、gradient checkpointing显存优化

当我们采用了分布式训练以及混合精度训练都不能降低显存大小的时候（比如多语言large模型，光词表就有几十万），现有的 GPU 资源无法训练一个设备装不下的模型。下面我们介绍一种改善这个问题的技术：梯度检查点（gradient checkpointing）

简单的说，梯度检查点的工作原理是在反向传播时重新计算深度神经网络的中间值（而通常情况是在前向传播时存储的)。这个策略是用时间（重新计算这些值两次的时间成本）来换空间（提前存储这些值的内存成本）。

神经网络如何使用内存

神经网络使用的总内存基本上是两个部分的总和。

• 第一部分是模型使用的静态内存。尽管 PyTorch 模型中内置了一些固定开销，但总的来说几乎完全由模型权重决定。而如今，在生产中使用的现代深度学习模型的总参数在100万到10亿之间。作为参考，一个带 16GB GPU 内存的 NVIDIA T4 的实际限制大约在1-1.5亿个参数之间。
• 第二部分是模型的计算图所占用的动态内存。在训练模式下，每次通过神经网络的前向传播都为网络中的每个神经元计算一个激活值，这个值随后被存储在所谓的计算图中。必须为批次中的每个单个训练样本存储一个值，因此数量会迅速的累积起来。总成本取决于模型大小和批处理大小，并设置适用于您的GPU内存的最大批处理大小的限制。一开始存储激活的原因是，在反向传播期间计算梯度时需要用到激活。

梯度检查点是如何起作用的

大型模型在静态和动态方面都很耗资源。首先，它们很难适配 GPU，而且哪怕你把它们放到了设备上，也很难训练，因为批处理大小被迫限制的太小而无法收敛。

梯度检查点（gradient checkpointing）的工作原理是从计算图中省略一些激活值（由前向传播产生，其中这里的”一些“是指可以只省略模型中的部分激活值，折中时间和空间，陈天奇在它的论文中Training Deep Nets with Sublinear Memory Cost使用了如下动图的方法，即前向传播的时候存一个节点释放一个节点，空的那个等需要用的时候再backword的时候重新计算）。这减少了计算图使用的内存，降低了总体内存压力（并允许在处理过程中使用更大的批次大小）。

PyTorch 通过torch.utils.checkpoint.checkpoint和torch.utils.checkpoint.checkpoint_sequential提供梯度检查点，根据官方文档的 notes，它实现了以下功能，在前向传播时，PyTorch 将保存模型中的每个函数的输入元组。在反向传播过程中，对于每个函数，输入元组和函数的组合以实时的方式重新计算，插入到每个需要它的函数的梯度公式中，然后丢弃（显存中只保存输入数据和函数）。网络计算开销大致相当于每个样本通过模型前向传播开销的两倍。

梯度检查点首次发表在2016年的论文《Training Deep Nets With Sublinear Memory Cost》中。论文声称提出的梯度检查点算法将模型的动态内存开销从 �(�) （n为模型中的层数）降低到 �(�) ，并通过实验展示了将 ImageNet 的一个变种从 48GB 压缩到了 7GB 内存占用。

五、chunk_size_applying（按多个小批量和低维度计算 FFN 部）

在Transformer的每个残差块中，自我注意层通常后面跟着2个前馈层。前馈层的中间嵌入尺寸通常大于模型的隐藏尺寸(例如，bert-base-uncased)。对于 input size 为 [bs, seq_length] 的输入，存储中间前馈嵌入所需的内存 [bs, seq_length, intermediate_size] 可以占内存使用的很大一部分。Reformer: The Efficient Transformer的作者注意到，由于计算是独立于sequence_length维度的，它在数学上等价于计算两个前馈层的输出嵌入 [bs, hidden_size]_0，[bs, hidden_size]_1， …， [bs, hidden_size]_n 配置。然后将它们连接到 [bs, seq_length, hidden_size] 中。其中n = seq_length（当然如果设置chunk_size=64的话，就会64个维度64个维度的算，最后再拼起来），这会增加计算时间，减少内存使用，但会产生数学上等价的结果。不过，在默认操作中不会特意设置这两个值（在源代码中默认为 0 和 1），所以会直接等效于正常的 forward 过程。

---

---

---

pytorch模型训练之fp16、apm、多GPU模型、梯度检查点（gradient checkpointing）显存优化等 - 知乎

参考文献：

pytorch 使用amp进行混合精度训练

Apex混合精度加速

飞狗：PyTorch的自动混合精度（AMP）（含有torch1.5以下和以上的两种amp用法）

pytorch单机多卡的训练方式(github)

使用梯度检查点(gradient checkpointing)训练比内存还大的pytorch模型