混合精度训练说明

关注

原创 2024-12-22 22:10:54 1.1k 阅读 AI 写同款· GEO 优化 ›

【NLP】使用混合精度技术加速大型语言模型由于大型语言模型 (LLM) 的计算要求和内存占用量大，因此训练和使用它们的成本很高。本文将探讨如何利用低精度格式将训练和推理速度提高 3 倍，同时不影响模型精度。尽管我们主要关注大型语言模型示例，但这些技术中的大多数都是通用的，也适用于其他深度学习架构。阅读详情

什么是混合精度训练？混合精度训练有什么用？
这里总结一下。
本文总结自kapathy的build gpt2

通常在训练过程中，model里面的数据默认都是torch.float32类型，
也就是用32bit的float型数据表示变量。
比如特征提取中提取的特征，描述子，网络的参数等，显示是torch.float32.

32bit的float带来的影响是比如拿它和int8比，占用内存会比较多，计算量会比较大。
所以训练时也就会比较慢，我们可以考虑丢弃一些精度来达到性能的提升。

看一下GPU支持的类型，比如nv a100吧，

在这里插入图片描述
可以看到实际上GPU是支持到float64的，但实际上模型训练时不需要那么高的精度，也会导致训练速度变慢，所以默认一般是float32.

精度下调的性能提升

float32最大可达到19.5TFLOPS，意味着可达到每秒19.5 trillion次float操作。
那么如果我们降低一些精度，比如降到TF32（见上图），每秒就可以达到156TFLOPS的性能，差不多达到了8倍性能提升。
如果进一步降低精度，到BFLOAT16, 就可以达到16倍性能提升。
注：右边*的性能指的是用稀疏化。

int8一般用于推理而不是训练。因为int8是均匀的空间，而训练时activation, weight都是正态分布的, 训练时要用float。

另一方面，精度降低之后占的内存少，易于搬运，这涉及到memory bandwidth和模型的memory. 可以参考图中的GPU memory bandwidth.
解释一下memory bandwidth, 一般情况下data需要搬到GPU上然后运算，运算完再搬回去（涉及到GPU内存），但是受限于这个bandwidth, 明明有多余算力，但是data还没搬进来，就需要等，导致利用度不够高。然而如果你降低了精度，数据占的内存就会变小，一次就可以搬运更多，那么每次参与计算的就会更多，在这个搬运上也会提升性能。

小结一下，适当降低精度会让计算量减少从而提升性能，另一方面，会让数据占的内存减少，每次搬运的数据更多，从而在有限的memory bandwidth上达到性能提升。

图上有个名词叫tensor core, 现在介绍一下什么是tensor core.
它是a100中的一个instruction, 它做的事情是4x4的矩阵乘法。
在这里插入图片描述

矩阵的乘法会broke up成这些4x4的矩阵乘。
比如在transformer中很多linear layer就需要矩阵乘法，特别是最后一层的classify layer是一个大矩阵乘。
矩阵乘就通过tensor core来加速。

不同精度的数据

TF32
在这里插入图片描述
通过这个图可以看到TF32和FP32表示的范围是一样的，最左边的sign是符号位，中间8位range表示数字可表达的范围，它们是一样的，区别就是TF32舍弃了一些小数点位。整个只有19bit, 而不是32bit。
这些是在硬件上完成的，pytorch代码上是不可见的。

为什么叫混合精度呢，是因为你的input还是fp32, output还是fp32, 但是在内部计算上，后面的bit被舍弃了，为提升性能降低了一些精度。所以结果会类似是一个近似结果，但你会几乎看不出差别。
虽然说明书上写的用TF32会有8x性能提升，但实际上那只是矩阵乘的时候用了TF32, 其他部分仍然是FP32, 另外还受限于memory bandwidth, 所以实际上大概率是达不到的。

但是注意一点，这是a100支持的TF32，有的GPU可能不支持。

说了这么多，到底怎么用TF32训练呢。
只需要一行code. 用到了torch.set_float32_matmul_precision, 有兴趣的可以查看官方文档。它有一些参数，“highest”， ”high“，”medium“等，其中"high"就是TF32. 默认是"highest", 也就是float32.

torch.set_float32_matmul_precision('high')

#model = GPT(..)
#model.to(device)
#training process

BFLOAT16
每个float只有16bit.
它和FP16有什么区别呢，还是看上面的图，它的range和float32是一样的，比FP16表示的范围要宽，只不过精度进一步被cut.
另外当你用FP16训练时，由于它表示的范围比FP32要小，所以你还要做gradient scaling操作。用BF16不需要做gradient scaling.

具体怎么用BF16呢，你可以参考torch.autocast, 但你不需要考虑gradient scaling的部分。
官方文档上有说，autocasting时，不要在model或input处call half() or bfloat16()。
你只能在forward和计算loss处用BF16.
具体如下, 只需要加一句torch.autocast：

model = GPT(..)
model.to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(iterations):
     x, y = train_loader.next_batch()
     x, y = x.to(device), y.to(device)
     optimizer.zero_grad()
     #用BF16
     with torch.autocast(device_type=device, dtype=torch.bfloat16):
         logits, loss = model(x, y)
     loss.backward()
     optimizer.step()

但是你看transformer里面的embedding table的weight, 仍然是float32, 具体哪些模块能cast到BF16，你可以参考官方文档里面的CUDA Ops that can autocast to float16. 需要矩阵乘的时候会cast, 很多操作仍然会保持float32不变，比如norm。