用于大型Transformer的8-bit矩阵乘法介绍

原文地址：A Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using transformers, accelerate and bitsandbytes

相关博客
【自然语言处理】【大模型】用于大型Transformer的8-bit矩阵乘法介绍
【自然语言处理】【大模型】BLOOM：一个176B参数且可开放获取的多语言模型
【自然语言处理】【大模型】PaLM：基于Pathways的大语言模型
【自然语言处理】【chatGPT系列】大语言模型可以自我改进
【自然语言处理】【ChatGPT系列】WebGPT：基于人类反馈的浏览器辅助问答
【自然语言处理】【ChatGPT系列】FLAN：微调语言模型是Zero-Shot学习器
【自然语言处理】【ChatGPT系列】ChatGPT的智能来自哪里？
【自然语言处理】【ChatGPT系列】Chain of Thought：从大模型中引导出推理能力
【自然语言处理】【ChatGPT系列】InstructGPT：遵循人类反馈指令来训练语言模型
【自然语言处理】【ChatGPT系列】大模型的涌现能力

一、简介

​ 语言模型正变的越来越大，PaLM已有有540B的参数量，而OPT、GPT-3和BLOOM则大约有176B参数量。下图是近些年语言模型的尺寸。

​ 这些模型很难在常用设备上运行。例如，仅仅推理BLOOM-176B就需要8x 80GB A00 GPUs。而为了微调BLOOM-176B则需要72个GPU。PaLM则需要更多的资源。

​ 这些巨型模型需要太多GPUs才能运行，因此需要寻找方法来减少资源需求并保证模型的性能。已经有各种技术用来减小模型尺寸，例如量化、蒸馏等。在完成BLOOM-176B训练后，HuggingFace和BigScience逐步探索在少量GPU上运行大模型的方法。最终，设计出了Int8量化方法，该方法在不降低大模型性能的情况下，将显存占用降低了1至2倍。

二、机器学习中常用的数据类型

​ 浮点数在机器学习中也被称为"精度"。模型大小是有参数量及参数精度决定的，通常是float32、float16和bfloat16。

我们开始于不同浮点数的基本理解，在机器学习的背景下也被称为"精度"。模型的大小由其参数数量和精度决定，通常是float32、float16和bfloat16。下图：

​ Float32(FP32)是标准的IEEE 32-bit浮点数表示，使用这种类型可以表示范围广泛的浮点数。在FP32中，8bits被用于"指数"，23bits被用于"尾数", 1 bit则用于符号位。大多数的硬件都支持FP32操作和指令。

​ 在Float16(FP16)数据类型中，5 bits被用作"指数"，10 bits用于"尾数"。这使得FP16数的表示范围明显小于FP32，导致有上溢和下溢的风险。例如，若你做$10k\times 10k$，最终得到100k。这在FP16中是不可能的，因为其最大表示为64K。因此你会得到NaN的结果，若像神经网络那样顺序执行，所有先前的工作都会被破坏。通常，loss缩放能够一定程度上克服这个问题，但并不总是有用。

​ 因此，创建了一种新格式bfloat15(BF16)来避免这种问题。在BF16中，8bits被用于表示"指数", 7bits被用于表示"尾数"。这意味着BF16能够保留和FP32相同的动态范围，但是损失了3bits的精度。BF16可以表示巨大的数，但是精度上比FP16差。

​ 在Ampere架构中，NVIDIA也引入了TensorFloat-32(TF32)精度格式，其仅使用19 bits就合并了BF16的动态范围和FP16的精度。其目前仅在内部某些操作中使用。

​ 在机器学习的术语中FP32被称为全精度(4 bytes)，BF16和FP16则称为半精度(2 bytes)。int8(INT8)数据类型则是由8 bits表示的数，其能够存储$2^8$个不同的值([0,255]或者[-128, 127])

​ 理想情况下，训练和推理应该在FP32上进行，但是其比FP16/BF16慢两倍。因此，采用一种混合精度的方法，模型权重仍然是FP32，前向和后向传播则使用FP16/BF16，从而加快训练速度。P16/BF16被用来更新FP32权重。

​ 可以通过参数量乘以浮点数精度的大小来计算模型所占用的bytes量。例如，若模型使用bfloat16版本的BLOOM-176B模型，那么模型大小为176*10**9 x 2 bytes = 352GB！这个量级对于适配少量GPU来说相当有挑战。

​ 但是我们是否可以使用不同的数据类型以更少的存储空间来保存这些权重？一种称为量化的方法被广泛的应用于Deep Learning。

三、模型量化介绍

​ 通过实验发现，在推理中使用2 bytes的BF16/FP16精度能够几乎达到4 bytes的FP32精度相同的效果，而且模型尺寸可以减少一半。若能够进一步削减那就太棒了，但是在更低的精度上推理质量开始急剧下降。为了解决这个问题，我们引入了8 bits量化。该方法使用四分之一的精度，这样仅需要1/4的模型尺寸！但是，其不是通过丢掉另外一半bits来实现的。

​ 量化基本上是从一种数据类型"舍入"为另一种数据类型来完成的。例如，若一个数量类型范围0…9，另一个范围则是0…4。那么第一个数据类型中的"4"将会被舍入为第二种数据类型中的"2"。然而，若第一种数据类型中的"3"，其会位于第二种数据类型的1和2之间，然后通常会被舍入为"2"。也就是说第一种数据类型中的"4"和"3"都会对应第二种数据类型中的"2"。这表明量化是可能带来信息丢失的噪音过程，一种有损压缩。

​ 有两种常见的8-bit量化技术：zero-point量化和absolute maximum(absmax)量化。zero-point量化和absmax量化会将浮点数值映射至更加紧凑的int8(1 byte)值。这些方法首先会将输入按照量化常数进行缩放，从而实现规范化。

​ 举例来说，在zero-point量化中，若范围是$[-1.0\dots 1.0]$并希望量化至范围$[-127\dots 127]$。那么应该按照因子127进行缩放，然后四舍五入至8-bit精度。为了还原原始值，需要将int8的值除以量化因子127。例如，0.3被缩放为$0.3\times 127=38.1$，然后四舍五入为38。若要恢复，则$38/127=0.2992$。在这个例子中量化误差为0.008。随着这些微小的误差在模型各个层中传播，会逐步积累和增长并导致性能下降。

​ 再来看看absmax量化的细节。为了在absmax量化中完成fp16和int8的映射，需要先除以张量中的绝对最大值(令整个张量介于-1至1之间)，然后在乘以目标数据类型的总范围。例如，在一个向量上应用absmax量化，该向量为
$$v=\left[\begin{array}{cccccccc}1.2& -0.5& -4.3& 1.2& -3.1& 0.8& 2.4& 5.4\end{array}\right]$$
从向量中选择最大值，即5.4。而int8的范围为[-127,127]，所以量化过程为$v/5.4\times 127=v\times \frac{127}{5.4}\approx v\times 23.5$，即整个向量乘以缩放因子23.5。最终得到的量化后向量为
$$\left[\begin{array}{cccccccc}28& -12& -101& 28& -73& 19& 56& 127\end{array}\right]$$

​ 为了还原原始值，可以使用全精度的int8数除以量化因子23.5。但是由于四舍五入的原因，会丢失一些精度。

​ 这些技巧能够以多种方式组合。例如，当涉及矩阵乘法时，ow-wise或者vector-wise量化可以使得结果更加准确。以矩阵乘法$A\times B=C$为例，相对于使用每个张量的绝对最大值来规范张量，vector-vise量化则会寻找矩阵A每行的绝对最大值和矩阵B每列的绝对最大值。然后通过除以这些绝对最大值向量来规范化矩阵A和B。然后执行$A\times B$来得到C。为了最终返回FP16精度的值，通过计算A和B绝对最大值向量的外积来反规范化。

​ 这些技术虽然能够量化模型，但是在较大模型上会带来性能下降。Hugging Face Transformers和Accelerate库集成了一种称为LLM.int8()的8-bit量化算法，能够在176B参数量模型上使用且不降低模型效果。

四、LLM.int8()简介

​ 理解Transformer中与规模相关的涌现特性对于理解为什么传统量化方式在大模型中失败至关重要。性能的下降是由异常特征值导致的，会在后面解释这一情况。LLM.int8()算法本质上可以由三个步骤来完成矩阵乘法：

1. 对输入的hidden states逐列的提取异常值(即大于某个阈值的值)；
2. 分别对FP16中的异常值和INT8中的非异常值执行矩阵乘法；
3. 对非异常的结果进行反量化，并将两者结果合并来获得最终的FP16结果；

三个步骤如下图所示：

1. 异常值特征

​ 在整个分布之外的值，称为异常值。异常值检测被广泛使用，而拥有特征分布的先验知识有助于异常值检测任务。

​ 具体来说，我们观察到经典的量化算法在超过6B参数量的transformer模型上失效了。虽然在较小的模型上也能观测到较大的异常值特征。但是，我们观察到一个参数量的阈值，transformer中的异常值会系统性的出现在每个层中。

​ 由于8-bit精度的局限性，因此仅使用几个特别大的值来量化向量将导致非常差的结果。此外，transformer架构的内在特征就是将所有的元素连接在一起，这将导致错误跨越多层传播并被加剧。因此，开发出了混合精度分解来实现这种极端异常值的量化。

2. MatMul内部

​ 一旦得到hidden state，使用自定义阈值来抽取异常值并分解矩阵为上述两部分。我们发现使用6作为阈值进行抽取可以完整的恢复推理性能。异常值部分以fp16实现，所以是经典的矩阵乘法；而8-bit则是通过vector-wise量化将模型权重和hidden state量化至8-bit的精度。即hidden-state使用row-wise量化，模型权重使用column-wise量化。经过这个步骤后，再将结果反量化并以半精度返回。

3. 零退化意味着什么

​ 如何评估性能下降？8-bit模型到底损失了多少性能？这里在8-bit模型和native模型上运行了常见的基准，分别针对OPT-175B和BLOOM-176B。

• 对于OPT-175B

• 对于BLOOM-176B

​ 可以看到这些模型的性能下降为0，因为这些指标的绝对差值小于标准误差。

4. 比native模型更快？

​ LLM.int8()方法的主要目标在不降低性能的情况下，使得大模型更容易被使用。但是，如果该方法非常的慢则就不实用了。所以，我们对多个模型的生成速度进行了基准测试。实验发现使用LLM.int8()的BLOOM-176B要比fp16版本慢15%至23%，这是一个可以接受的范围。但是较小的模型下降会更多。开发人员正在逐步优化这个问题。

五、Transformers集成

1. 使用

​ 本文重点描述的模块是Linear8bitLt，你可以直接从bitsandbytes库中引入。其来自于经典的torch.nn模块，并使用下面的代码来轻易的使用和部署。

​ 下面是一个使用bitsandbytes将一个小模型转换为int8类型。

• 正确的引入

import torch
import torch.nn as nn

from bitsandbytes.nn import Linear8bitLt

• 先定义一个fp16的模型

fp16_model = nn.Sequential(nn.Linear(64, 64),
                           nn.Linear(64, 64))

• 假设该模型已经完成训练，保存模型

torch.save(fp16_model.state_dict(), "model.pt")

• 现在再定义一个int8模型

int8_model = nn.Sequential(
    Linear8bitLt(64, 64, has_fp16_weights=False),
    Linear8bitLt(64, 64, has_fp16_weights=False)
)

添加参数has_fp16_weights很重要。默认值为True，其被用于Int8/FP16混合精度训练。然而，这里关注的是推理，所以将其设置为False。

• 现在将fp16的模型加载至int8模型中

int8_model.load_state_dict(torch.load("model.pt"))
# print(int8_model[0].weight)
int8_model = int8_model.to("cuda:0") # 执行该代码时会进行量化
# print(int8_model[0].weight)

通过输出print(int8_model[0].weight)可以看到模型被量化为Int8类型，那么怎么还原为FP16权重呢？

(int8_model[0].weight.CB * int8_model[0].weight.SCB) / 127

• 使用int8模型进行推理

input_ = torch.randn((1,64), dtype=torch.float16)
hidden_states = int8_model(input_.to(torch.device('cuda:0')))

2. 你只需要accelerate

​ 当使用大模型时，acceleate库包含了有用的程序。init_empty_weights方法特别有用，因为任何模型(无论大小)都可以作为上下文管理器使用此方法进行初始化，而无需为模型权重分配任何内存。

import torch.nn as nn
from accelerate import init_empty_weights

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(100000, 100000) for _ in range(1000)])

这个初始化的模型会被放置至Pytorch的元设备上，其是一种不用分配存储空间来表示shape和dtype的潜在机制。

起初，该函数在.from_pretrained函数中被调用，并将所有参数重写为torch.nn.Parameter。但是，这不符合我们的需求，因为希望在Linear8bitLt模块中保留Int8Params类。因此我们将

module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))

修改为

param_cls = type(module._parameters[name])
kwargs = module._parameters[name].__dict__
module._parameters[name] = param_cls(module._parameters[name].to(torch.device("meta")), **kwargs)

通过这个修改，我们可以通过自定义函数在没有任何内存消耗的情况下，利用这个上下文管理器将所有的nn.Linear替换为bnb.nn.Linear8bitLt。

def replace_8bit_linear(model, threshold=6.0, module_to_not_convert="lm_head"):
    for name, module in model.named_children():
        if len(list(module.children())) > 0:
            # 递归
            replace_8bit_linear(module, threshold, module_to_not_convert)

        if isinstance(module, nn.Linear) and name != module_to_not_convert:
            with init_empty_weights():
                model._modules[name] = bnb.nn.Linear8bitLt(
                    module.in_features,
                    module.out_features,
                    module.bias is not None,
                    has_fp16_weights=False,
                    threshold=threshold,
                )
    return model

该函数会递归的将元设备上的所有nn.Linear替换为Linear8bitLt模块。属性has_fp16_weights必须被设置为False，以便加载int8权重和量化信息。

3. 如何在transformers中使用

from transformers import AutoTokenizer, AutoModelForCausalLM

def inference(payload, model, tokenizer):
    input_ids = tokenizer(payload, return_tensors="pt").input_ids.to(model.device)
    print(f"输入:\n    {payload}")
    logits = model.generate(input_ids, num_beams=1, max_new_tokens=128)
    print(f"生成:\n    {tokenizer.decode(logits[0].tolist()[len(input_ids[0]):])}")

model_name = "bigscience/bloomz-7b1-mt"
payload = "一个传奇的开端，一个不灭的神话，这不仅仅是一部电影，而是作为一个走进新时代的标签，永远彪炳史册。你认为这句话的立场是赞扬、中立还是批评？"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
model_native = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
# 比较推理结果
inference(payload, model_8bit, tokenizer)
inference(payload, model_native, tokenizer)
# 计算显存节约程度
mem_fp16 = model_native.get_memory_footprint()
mem_int8 = model_8bit.get_memory_footprint()
print(mem_fp16/mem_int8)