一个温和的介绍:使用transformers、accelerate和bitsandbytes进行规模化的8位矩阵乘法的转换器

一个温和的介绍:使用transformers、accelerate和bitsandbytes进行规模化的8位矩阵乘法转换器

介绍

语言模型的规模不断扩大。在撰写本文时,PaLM具有540B的参数,OPT、GPT-3和BLOOM约有176B的参数,而且我们趋向于使用更大的模型。下图显示了一些最近语言模型的规模。

因此,这些模型很难在易于获取的设备上运行。例如,仅对BLOOM-176B进行推理,您需要拥有8个80GB的A100 GPU(每个约15000美元)。要对BLOOM-176B进行微调,您需要72个这样的GPU!像PaLM这样更大的模型需要更多的资源。

由于这些庞大的模型需要如此多的GPU运行,我们需要找到减少这些要求的方法,同时保持模型的性能。已经开发了各种技术来尝试缩小模型的大小,您可能听说过量化和蒸馏等等。

在完成BLOOM-176B的训练后,我们在HuggingFace和BigScience的帮助下寻找了使这个大模型在较少的GPU上更易于运行的方法。通过我们的BigScience社区,我们了解到了关于Int8推理的研究,该研究不会降低大模型的预测性能,并将大模型的内存占用减少了2倍。我们很快开始合作进行这项研究,并最终将其完全集成到Hugging Face的transformers中。通过本博客文章,我们为所有Hugging Face模型提供了LLM.int8()集成,并在下面详细解释了这一点。如果您想了解更多关于我们的研究的信息,可以阅读我们的论文《LLM.int8(): 用于大规模Transformer的8位矩阵乘法》。

本文着重概述量化技术的高级概述,概述了将其纳入transformers库的困难,并制定了此合作伙伴关系的长期目标。

在这里,您将了解到什么使一个大模型使用了这么多内存?是什么使BLOOM使用了350GB?我们将逐渐介绍一些基本前提。

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

我们从对不同浮点数据类型的基本理解开始,这些数据类型在机器学习的上下文中也被称为“精度”。

模型的大小由其参数的数量及其精度确定,通常为float32、float16或bfloat16之一(下图来自:https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/)。

Float32(FP32)代表标准化的IEEE 32位浮点表示。使用这种数据类型可以表示广泛的浮点数。在FP32中,8位用于“指数”,23位用于“尾数”,1位用于数的符号。此外,大多数硬件支持FP32操作和指令。

在浮点16(FP16)数据类型中,5位用于指数,10位用于尾数。这使得FP16数的可表示范围远低于FP32。这使FP16数面临溢出(尝试表示一个非常大的数字)和下溢(表示一个非常小的数字)的风险。

例如,如果执行10k * 10k,最终得到100M,这在FP16中是不可能表示的,因为最大可能的数字是64k。因此,您将得到NaN(不是一个数字)的结果,如果您有像神经网络中的顺序计算一样的连续计算,所有先前的工作都会被破坏。通常,使用损失缩放来克服此问题,但它并不总是有效。

为了避免这些限制,创建了一种新的格式,bfloat16(BF16)。在BF16中,8位用于指数(与FP32相同),7位用于小数部分。

这意味着在BF16中,我们可以保持与FP32相同的动态范围。但是与FP16相比,我们失去了3位精度。现在对于大数值没有任何问题,但是精度比FP16差。

在安培架构中,NVIDIA还引入了TensorFloat-32(TF32)精度格式,将BF16的动态范围和FP16的精度结合起来,只使用19位。目前它只在某些操作中内部使用。

在机器学习术语中,FP32被称为全精度(4字节),而BF16和FP16被称为半精度(2字节)。此外,int8(INT8)数据类型由一个8位表示组成,可以存储2^8个不同的值(在[0, 255]或[-128, 127]之间的有符号整数)。

虽然理想情况下,训练和推断应该使用FP32进行,但FP32的速度比FP16/BF16慢两倍,因此采用混合精度方法,在FP32中保持权重作为精确的“主权重”参考,而在前向和后向传递中使用FP16/BF16来增强训练速度。然后使用FP16/BF16梯度更新FP32主权重。

在训练过程中,主权重始终以FP32存储,但在实践中,半精度权重在推断过程中通常提供与FP32相似的质量–只有当模型接收到多个梯度更新时才需要精确的参考。这意味着我们可以使用半精度权重并使用一半的GPU来实现相同的结果。

要计算模型的大小(以字节为单位),需要将参数数量乘以所选择精度的字节数。例如,如果我们使用BLOOM-176B模型的bfloat16版本,则有176*10**9 x 2字节 = 352GB!正如前面所讨论的,这对于几个GPU来说是一个相当大的挑战。

但是,如果我们可以使用不同的数据类型来存储这些权重以节省内存呢?一种称为量化的方法已被广泛应用于深度学习中。

模型量化简介

经验上,我们发现与4字节FP32精度相比,使用2字节BF16/FP16半精度几乎可以得到相同的推断结果,从而将模型大小减半。如果能进一步减少,那将是令人惊奇的,但是推断质量结果在更低的精度下开始急剧下降。

为了解决这个问题,我们引入了8位量化。这种方法使用四分之一的精度,因此只需要模型大小的四分之一!但这不仅仅是丢弃另外一半的位数。

量化实质上是通过将一种数据类型“舍入”为另一种数据类型来完成的。例如,如果一个数据类型的范围是0到9,而另一个数据类型的范围是0到4,那么第一种数据类型中的值“4”将舍入为第二种数据类型中的“2”。然而,如果我们在第一种数据类型中有值“3”,它位于第二种数据类型的1和2之间,那么通常会舍入为“2”。这表明量化是一个噪声过程,可能会导致信息丢失,一种有损压缩。

最常见的两种8位量化技术是零点量化和绝对最大值(absmax)量化。零点量化和absmax量化将浮点值映射为更紧凑的int8(1字节)值。首先,这些方法通过将输入归一化并将其乘以量化常数来对其进行缩放。

例如,在零点量化中,如果我的范围是-1.0…1.0,并且我想量化到-127…127的范围,我希望乘以因子127,然后将其舍入为8位精度。要恢复原始值,您需要将int8值除以相同的量化因子127。例如,值0.3将缩放为0.3*127 = 38.1。通过舍入,我们得到38的值。如果我们将其反转,我们得到38/127=0.2992 – 在这个示例中,我们有0.008的量化误差。这些看似微小的错误往往会累积并随着它们通过模型的层级传播而增加,导致性能下降。

(图像来源:这篇博文)

现在让我们来看看absmax量化的细节。要计算fp16数和其对应的int8数在absmax量化中的映射,首先你必须将其除以张量的绝对最大值,然后乘以数据类型的总范围。

例如,假设你想在一个包含[1.2, -0.5, -4.3, 1.2, -3.1, 0.8, 2.4, 5.4]的向量中应用absmax量化。你提取其中的绝对最大值,这种情况下为5.4。Int8的范围为[-127, 127],因此我们将127除以5.4,得到缩放因子23.5。因此,将原始向量乘以它得到量化向量[28, -12, -101, 28, -73, 19, 56, 127]

要恢复原始值,只需用量化因子对int8数进行全精度除法运算,但由于上面的结果是“四舍五入”的,会丢失一些精度。

对于无符号int8,我们会减去最小值并按绝对最大值进行缩放。这与零点量化的做法接近。它类似于最小-最大缩放,但后者以保持值的比例,使值“0”始终用整数表示,不会有任何量化误差。

这些技巧可以以多种方式组合,例如,按行或按向量进行量化,当涉及矩阵乘法时可以获得更准确的结果。观察矩阵乘法A*B=C,与每个张量单独归一化的常规量化不同,按向量进行量化会找到A的每一行和B的每一列的绝对最大值。然后我们通过这些向量对A和B进行归一化。然后我们将A*B相乘得到C。最后,为了恢复FP16值,我们通过计算A和B的绝对最大值向量的外积来进行反归一化。关于这种技术的更多细节可以在LLM.int8()论文或Tim的博客中关于量化和新兴特性的博文中找到。

尽管这些基本技术使我们能够量化深度学习模型,但对于更大的模型,它们通常会导致准确性下降。我们将LLM.int8()实现集成到Hugging Face Transformers和Accelerate库中,它是第一种在176B参数的大模型(如BLOOM)上不会降低性能的技术。

LLM.int8()的简要概述:大型语言模型的零降级矩阵乘法

在LLM.int8()中,我们已经证明了理解transformer的尺度相关的新兴特性对于理解为什么传统量化在大型模型上失败至关重要。我们证明了性能下降是由异常特征引起的,我们将在下一节中解释。LLM.int8()算法本身可以解释如下。

实质上,LLM.int8()试图通过三个步骤完成矩阵乘法计算:

  1. 从输入的隐藏状态中按列提取异常值(即大于某个阈值的值)。
  2. 在异常值中使用FP16进行矩阵乘法,而在非异常值中使用int8。
  3. 对非异常值结果进行反量化,然后将异常值和非异常值结果相加,得到完整结果的FP16。

这些步骤可以用以下动画总结:

异常特征的重要性

一般来说,超出一些数字全局分布范围的值被称为异常值。异常值检测已经广泛应用并在当前的文献中进行了研究,对于特征分布的先验知识有助于异常值检测任务。更具体地说,我们观察到,规模上的经典量化在>6B参数的基于transformer的模型上失败。虽然较小模型中也存在大量异常特征,但我们观察到这些异常特征与transformer的每一层中都存在高度系统化的模式。有关这些现象的更多细节,请参阅LLM.int8()论文和新兴特性博文。

正如前面提到的,8位精度非常受限,因此对具有几个较大值的向量进行量化可能会产生非常错误的结果。此外,由于基于transformer的架构内置了将所有元素连接在一起的特性,这些错误在传播到多个层时往往会不断累积。因此,混合精度分解已经被开发出来,以便在存在极端离群值的情况下实现高效的量化。下面将对此进行讨论。

在MatMul中

计算隐藏状态后,我们使用自定义阈值提取离群值,并按照上述说明将矩阵分解为两部分。我们发现,以这种方式提取所有绝对值大于等于6的离群值可以恢复完整的推理性能。离群值部分使用fp16完成,因此它是一个经典的矩阵乘法,而8位矩阵乘法则通过将权重和隐藏状态进行向量化量化(即,对隐藏状态进行行级量化,对权重矩阵进行列级量化)来完成。在这一步之后,结果被反量化并以半精度返回,以便将其添加到第一个矩阵乘法中。

0降级意味着什么?

我们如何正确评估该方法的性能降级?使用8位模型时,我们在生成方面失去了多少质量?

我们使用lm-eval-harness运行了几个常见的基准测试,并报告了结果。

对于OPT-175B:

对于BLOOM-176:

我们确实观察到这些模型的性能没有降级,因为度量指标的绝对差异都低于标准误差(除了BLOOM-int8在lambada上略优于原生模型)。要进行更详细的性能评估,与最先进的方法进行比较,请参阅论文!

它比原生模型更快吗?

LLM.int8()的主要目的是在不降低性能的情况下使大型模型更易于使用。但是,如果该方法非常慢,它将变得不太有用。因此,我们对多个模型的生成速度进行了基准测试。我们发现,使用LLM.int8()的BLOOM-176B版本比fp16版本慢15%至23%——这仍然是可以接受的。对于较小的模型(如T5-3B和T5-11B),我们发现减速更大。我们努力提高这些小模型的速度。在一天内,我们将T5-3B的每个令牌的推理时间从312毫秒提高到173毫秒,并将T5-11B的推理时间从45毫秒降低到25毫秒。此外,已经识别出了问题,并且在即将发布的版本中,LLM.int8()对于小模型来说可能会更快。目前,当前的数字如下表所示。

这3个模型分别是BLOOM-176B、T5-11B和T5-3B。

Hugging Face transformers集成细节

接下来,让我们讨论Hugging Face transformers集成的具体细节。让我们看一下如何使用以及在设置过程中可能遇到的常见问题。

用法

负责整个博客文章中描述的魔法的模块名为Linear8bitLt,您可以轻松从bitsandbytes库中导入它。它派生自经典的torch.nn模块,可以在您的架构中轻松使用和部署,使用下面描述的代码。

下面是一个逐步示例,用于说明以下用例:假设您希望使用bitsandbytes将一个小型模型转换为int8。

  1. 首先,我们需要以下正确的导入!
import torch
import torch.nn as nn

import bitsandbytes as bnb
from bnb.nn import Linear8bitLt
  1. 然后,您可以定义自己的模型。请注意,您可以将任何精度的检查点或模型转换为8位(FP16、BF16或FP32),但是,当前,我们的Int8模块要求模型的输入为FP16才能正常工作。因此,在这里我们将我们的模型视为fp16模型。
fp16_model = nn.Sequential(
    nn.Linear(64, 64),
    nn.Linear(64, 64)
)
  1. 假设你已经在你最喜欢的数据集和任务上训练好了你的模型!现在是时候保存模型了:
[... 训练模型 ...]
torch.save(fp16_model.state_dict(), "model.pt")
  1. 现在你的state_dict已经保存好了,让我们来定义一个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精度下训练。然而,我们对内存效率推断感兴趣,因此需要使用has_fp16_weights=False

  1. 现在是时候加载你的8位模型了!
int8_model.load_state_dict(torch.load("model.pt"))
int8_model = int8_model.to(0) # 量化发生在这里

请注意,在第二行的调用中,当模型设置在GPU上时,量化步骤完成。如果在调用.to函数之前打印int8_model[0].weight,你将得到:

int8_model[0].weight
Parameter containing:
tensor([[ 0.0031, -0.0438,  0.0494,  ..., -0.0046, -0.0410,  0.0436],
        [-0.1013,  0.0394,  0.0787,  ...,  0.0986,  0.0595,  0.0162],
        [-0.0859, -0.1227, -0.1209,  ...,  0.1158,  0.0186, -0.0530],
        ...,
        [ 0.0804,  0.0725,  0.0638,  ..., -0.0487, -0.0524, -0.1076],
        [-0.0200, -0.0406,  0.0663,  ...,  0.0123,  0.0551, -0.0121],
        [-0.0041,  0.0865, -0.0013,  ..., -0.0427, -0.0764,  0.1189]],
       dtype=torch.float16)

而如果在第二行调用之后打印它,你将得到:

int8_model[0].weight
Parameter containing:
tensor([[   3,  -47,   54,  ...,   -5,  -44,   47],
        [-104,   40,   81,  ...,  101,   61,   17],
        [ -89, -127, -125,  ...,  120,   19,  -55],
        ...,
        [  82,   74,   65,  ...,  -49,  -53, -109],
        [ -21,  -42,   68,  ...,   13,   57,  -12],
        [  -4,   88,   -1,  ...,  -43,  -78,  121]],
        device='cuda:0', dtype=torch.int8, requires_grad=True)

权重值被“截断”,正如我们在前面的章节中所解释的量化。此外,这些值似乎分布在[-127, 127]之间。你可能还想知道如何获取FP16权重,以便在fp16中执行离群值的矩阵乘法?你可以简单地执行以下操作:

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

你将得到:

tensor([[ 0.0028, -0.0459,  0.0522,  ..., -0.0049, -0.0428,  0.0462],
        [-0.0960,  0.0391,  0.0782,  ...,  0.0994,  0.0593,  0.0167],
        [-0.0822, -0.1240, -0.1207,  ...,  0.1181,  0.0185, -0.0541],
        ...,
        [ 0.0757,  0.0723,  0.0628,  ..., -0.0482, -0.0516, -0.1072],
        [-0.0194, -0.0410,  0.0657,  ...,  0.0128,  0.0554, -0.0118],
        [-0.0037,  0.0859, -0.0010,  ..., -0.0423, -0.0759,  0.1190]],
       device='cuda:0')

这足够接近原始的FP16值(上面的输出两次)了!

  1. 现在,您可以通过确保输入位于正确的GPU上并且为FP16来安全地推断使用您的模型:
input_ = torch.randn((1, 64), dtype=torch.float16)
hidden_states = int8_model(input_.to(torch.device('cuda', 0)))

请查看完整的最小代码示例!

值得一提的是,您应该注意这些模块与nn.Linear模块略有不同,因为它们的参数来自bnb.nn.Int8Params类,而不是nn.Parameter类。稍后您将看到,这在我们的旅程中带来了额外的障碍!

现在是时候了解如何将其整合到transformers库中了!

accelerate就是你所需要的

当使用庞大的模型时,accelerate库提供了许多有用的工具。尤其有用的是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)]) # 这将占用~0 RAM!

初始化的模型将放置在PyTorch的meta设备上,这是一种在不为存储分配内存的情况下表示形状和数据类型的底层机制。多酷啊!

最初,此函数在.from_pretrained函数中调用,并将所有参数覆盖为torch.nn.Parameter。但是,这不符合我们的要求,因为我们希望保留Int8Params类用于上述Linear8bitLt模块。我们在以下PR中修复了该问题:

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层替换为在meta设备上初始化的Linear8bitLt模块。属性has_fp16_weights必须设置为False,以便直接将权重加载为int8以及量化统计信息。

我们还会忽略某些模块(这里是lm_head),因为我们希望保留其原始精度以获得更精确和稳定的结果。

但事情还没有结束!上述函数在init_empty_weights上下文管理器下执行,这意味着新模型仍然位于meta设备上。对于在此上下文管理器下初始化的模型,accelerate将手动加载每个模块的参数并将其移动到正确的设备上。在bitsandbytes中,设置Linear8bitLt模块的设备是一个关键步骤(如果您感兴趣,可以在此处检查代码片段),正如我们在玩具脚本中所见。

在调用量化步骤两次时,出现了失败的情况。我们不得不实现accelerateset_module_tensor_to_device函数(称为set_module_8bit_tensor_to_device),以确保我们不会调用它两次。让我们在下面的部分详细讨论这个问题!

使用accelerate设置设备时要非常小心

accelerate库中,我们进行了非常微妙的平衡!一旦加载了模型并将其设置在正确的设备上,有时仍然需要调用set_module_tensor_to_device来在所有设备上分发带有钩子的模型。这是在acceleratedispatch_model函数内完成的,其中可能涉及多次调用.to,而我们希望避免这样做。我们需要进行2个拉取请求才能实现我们想要的效果!最初的拉取请求在这里打破了一些测试,但是这个拉取请求成功地修复了一切!

总结

因此,最终的解决方案是:

  1. 使用正确的模块在设备上初始化模型
  2. 逐个将参数设置在正确的GPU设备上,并确保您从不执行此过程两次!
  3. 在正确的位置添加新的关键字参数,并添加一些良好的文档
  4. 添加非常详细的测试!详细信息请查看我们这里的测试。这听起来可能很容易,但是我们一起经历了许多艰难的调试会话,通常涉及CUDA内核!

说了这么多,这个集成冒险非常有趣;从对不同库进行深入研究和进行一些“手术”,到对齐一切并使其正常工作!

现在是时候看看如何从这个集成中受益以及如何成功地在transformers中使用它了!

如何在transformers中使用它

硬件要求

在CPU上不支持8位张量核心。bitsandbytes可以在支持8位张量核心的硬件上运行,这些硬件包括图灵和安培GPU(RTX 20s、RTX 30s、A40-A100、T4+)。例如,Google Colab GPU通常是NVIDIA T4 GPU,他们的最新一代GPU确实支持8位张量核心。我们的演示是基于Google Colab的,所以请查看下面的演示!

安装

只需使用下面的命令安装最新版本的库(确保您使用的是python>=3.8),并运行下面的命令进行尝试

pip install accelerate
pip install bitsandbytes
pip install git+https://github.com/huggingface/transformers.git

示例演示 – 在Google Colab上运行T5 11B

查看在BLOOM-3B模型上运行8位模型的Google Colab演示!

这是运行T5-11B的演示。T5-11B模型检查点是FP32格式,占用42GB内存,在Google Colab上无法容纳。使用我们的8位模块,它只使用11GB内存,轻松放置:

或者这个BLOOM-3B的演示:

改进的范围

在我们的观点中,这种方法极大地改善了访问非常大的模型的能力。没有性能下降,它使得具有较少计算能力的用户能够访问以前无法访问的模型。我们发现了几个可以在将来改进的领域,以使这种方法对大型模型变得更好!

更快的小型模型推理速度

正如我们在基准测试部分所看到的,我们可以将小型模型(<=6B参数)的运行时速度提高近2倍。然而,尽管大型模型如BLOOM-176B的推理速度非常稳健,但对于小型模型仍然有改进空间。我们已经确定了问题,并可能实现与fp16相同的性能,或者获得小的加速。在接下来的几周内,您将看到这些变化被整合进来。

支持Kepler GPU(GTX 1080等)

虽然我们支持过去四年中的所有GPU,但一些旧的GPU(如GTX 1080)仍然被广泛使用。虽然这些GPU没有Int8张量核心,但它们具有Int8向量单元(一种“弱”张量核心)。因此,这些GPU也可以体验到Int8加速。然而,为了实现快速推理,需要一个完全不同的软件堆栈。虽然我们计划整合对Kepler GPU的支持,以使LLM.int8()功能更广泛地可用,但由于其复杂性,这需要一些时间来实现。

在Hub上保存8位状态字典

将8位状态字典直接加载到8位模型中目前是不可能的,因为模型计算的统计数据(记住weight.CBweight.SCB)目前未存储或考虑在状态字典内,并且Linear8bitLt模块尚不支持此功能。我们认为保存并将其推送到Hub可能有助于提高可访问性。

CPU支持

正如在本文开头所述,CPU设备不支持8位核心。然而,我们能否克服这个问题?在CPU上运行这个模块也会显著提高可用性和可访问性。

在其他模态上进行扩展

目前,语言模型占据了非常大的模型。将这种方法应用于非常大的视觉、音频和多模态模型可能是未来几年更好地提高可访问性的有趣事情。

致谢

特别感谢以下人士对文章的可读性改进以及在transformers的集成过程中做出的贡献(按字母顺序排列):JustHeuristic(Yozh)、Michael Benayoun、Stas Bekman、Steven Liu、Sylvain Gugger、Tim Dettmers。