16,8和4位浮点数是如何工作的（附代码）
本文约5700字，建议阅读10分钟
本文将介绍最流行的浮点格式，创建一个简单的神经网络，并了解它是如何工作的。
50年前Kernighan、Ritchie和他们的C语言书的第一版开始，人们就知道单精度“float”类型有32位大小，双精度类型有64位大小。还有一种具有扩展精度的80位“长双精度”类型，这些类型几乎涵盖了浮点数据处理的所有需求。但是在最近几年，尤其是今年LLM的兴起，为了减小模型的存储和内存占用，开发人员开始尽可能地缩小浮点类型。
在本文中，我们将介绍最流行的浮点格式，创建一个简单的神经网络，并了解它是如何工作的。
“标准”32位浮点数
我们先回顾一下标准格式。IEEE 754浮点运算标准由IEEE于1985年制定。32浮点型的典型数字是这样的:
第一个比特（bit）是一个符号，接下来的8个比特代表一个指数，最后一个比特代表尾数。最终值的计算公式为:
我们创建一个辅助函数以二进制形式打印浮点值:
import struct
def print_float32(val: float): """ Print Float32 in a binary form """ m = struct.unpack('I', struct.pack('f', val))[0] return format(m, 'b').zfill(32)
print_float32(0.15625)
# > 00111110001000000000000000000000再创建一个逆向转换函数，这将在后面有用:
def ieee_754_conversion(sign, exponent_raw, mantissa, exp_len=8, mant_len=23): """ Convert binary data into the floating point value """ sign_mult = -1 if sign == 1 else 1 exponent = exponent_raw - (2 ** (exp_len - 1) - 1) mant_mult = 1 for b in range(mant_len - 1, -1, -1): if mantissa & (2 ** b): mant_mult += 1 / (2 ** (mant_len - b))
return sign_mult * (2 ** exponent) * mant_mult
ieee_754_conversion(0b0, 0b01111100, 0b01000000000000000000000)
#> 0.15625作为开发，你肯定知道浮点类型的准确性是有限的，比如这个：
val = 3.14 print(f"{val:.20f}")
# > 3.14000000000000012434在一般情况下，这不是一个大问题，但是我们拥有的比特位越少，得到的精度就越低。
16位浮点数
早期对这种格式的需求并不大，直到2008年才将16位浮点类型添加到IEEE 754标准中。它有一个符号位，5个指数位和10位尾数(分数):
他的转换逻辑与32位浮点数相同，但精度较低。以二进制形式打印一个16位浮点数:
import numpy as np
def print_float16(val: float): """ Print Float16 in a binary form """ m = struct.unpack('H', struct.pack('e', np.float16(val)))[0] return format(m, 'b').zfill(16)
print_float16(3.14)
# > 0100001001001000使用之前使用的方法，我们可以进行反向转换:
ieee_754_conversion(0, 0b10000, 0b1001001000, exp_len=5, mant_len=10)
# > 3.140625我们还可以找到Float16中可以表示的最大值:
ieee_754_conversion(0, 0b11110, 0b1111111111, exp_len=5, mant_len=10)
#> 65504.0这里使用0b11110，是因为在IEEE 754标准中，0b11111是为“无穷大”保留的。同理还可以找到可能的最小值:
ieee_754_conversion(0, 0b00001, 0b0000000000, exp_len=5, mant_len=10)
#> 0.00006104对于大多数开发人员来说，像这样的类型是一种“未知的领域”，因为c++中也没有标准的16位浮点类型。
16位" bfloat " (BFP16)
这种浮点格式是由谷歌团队开发的，它是专门为机器学习设计的(名字中的“B”也代表“大脑”)。该类型是对“标准”16位浮点数的修改:指数被扩大到8位，因此“bfloat16”的动态范围实际上与float-32相同。但尾数的大小被减少到7位:
让我们做一个和之前类似的计算:
ieee_754_conversion(0, 0b10000000, 0b1001001, exp_len=8, mant_len=7)
#> 3.140625可以看到由于指数较大，bfloat16格式具有更宽的范围:
ieee_754_conversion(0, 0b11111110, 0b1111111, exp_len=8, mant_len=7)
#> 3.3895313892515355e+38这比前面示例中的65504.0要好得多，但是正如前面提到的：因为尾数中的位数更少，所以bfloat16的精度更低，可以在Tensorflow中测试这两种类型:
import tensorflow as tf
print(f"{tf.constant(1.2, dtype=tf.float16).numpy().item():.12f}")
# > 1.200195312500
print(f"{tf.constant(1.2, dtype=tf.bfloat16).numpy().item():.12f}")
# > 1.2031250000008位浮点(FP8)
这种(相对较新的)格式是在2022年提出的，它也是为机器学习而创建的——因为模型变得更大，将它们放入GPU内存是一个挑战。FP8格式有两种变体:E4M3(4位指数和3位尾数)和E5M2(5位指数和2位尾数)。
让我们来获取两种格式的最大可能值:
ieee_754_conversion(0, 0b1111, 0b110, exp_len=4, mant_len=3)
# > 448.0
ieee_754_conversion(0, 0b11110, 0b11, exp_len=5, mant_len=2)
# > 57344.0也可以在Tensorflow中使用FP8:
import tensorflow as tf from tensorflow.python.framework import dtypes
a_fp8 = tf.constant(3.14, dtype=dtypes.float8_e4m3fn) print(a_fp8)
# > 3.25
a_fp8 = tf.constant(3.14, dtype=dtypes.float8_e5m2) print(a_fp8)
# > 3.0让我们在这两种类型中画一个正弦波:
import numpy as np import tensorflow as tf from tensorflow.python.framework import dtypes import matplotlib.pyplot as plt
length = np.pi * 4 resolution = 200 xvals = np.arange(0, length, length / resolution) wave = np.sin(xvals) wave_fp8_1 = tf.cast(wave, dtypes.float8_e4m3fn) wave_fp8_2 = tf.cast(wave, dtypes.float8_e5m2)
plt.rcParams["figure.figsize"] = (14, 5) plt.plot(xvals, wave_fp8_1.numpy()) plt.plot(xvals, wave_fp8_2.numpy()) plt.show()可以看到，有一些差别，但是都还不错。
可以明显看到一些精度的损失，但这个图像看起来仍然像正弦波!
4位浮点类型
现在让我们来看看最“疯狂”的东西——4位浮点值!4位浮点数(FP4)是遵循IEEE标准的最小可能值，具有1位符号，2位指数和1位尾数:
第二种可能的4位实现是所谓的NormalFloat (NF4)数据类型。NF4值针对保存正态分布变量进行了优化。其他数据类型很难做到这一点，但所有可能的NF4值都可以很容易地打印在一个列表中:
[-1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0]FP4和NF4类型都在bitsandbytes库中有相应的实现。作为一个例子，让我们将[1.0,2.0,3.0,4.0]数组转换为FP4:
from bitsandbytes import functional as bf
def print_uint(val: int, n_digits=8) -> str: """ Convert 42 => '00101010' """ return format(val, 'b').zfill(n_digits)
device = torch.device("cuda") x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device) x_4bit, qstate = bf.quantize_fp4(x, blocksize=64)
print(x_4bit) # > tensor([[117], [35]], dtype=torch.uint8)
print_uint(x_4bit[0].item()) # > 01110101 print_uint(x_4bit[1].item()) # > 00100011
print(qstate) # > (tensor([4.]), # > 'fp4', # > tensor([ 0.0000, 0.0052, 0.6667, 1.0000, 0.3333, 0.5000, 0.1667, 0.2500, # > 0.0000, -0.0052, -0.6667, -1.0000, -0.3333, -0.5000, -0.1667, -0.2500])])作为输出，我们得到两个对象:一个16位数组[117,35]，实际上包含我们的4个数字，和一个“状态”对象，包含缩放因子4.0和所有16个FP4数字的张量。
例如，第一个4位数字为“0111”(=7)，在状态对象中我们可以看到对应的浮点值为0.25;0.25 *4 = 1.0。第二个数字是“0101”(=5)，结果是0.5*4 = 2.0。对于第三个数字，“0010”是2,0.666*4 = 2.666，接近但不等于3.0。对于4位值显然有了一些精度损失。对于最后一个值，“0011”是3,1000 *4 = 4.0。
逆向转换不需要手动操作bitsandbytes可以帮我们自动完成
x = bf.dequantize_fp4(x_4bit, qstate) print(x)
# > tensor([1.000, 2.000, 2.666, 4.000])4位格式也有一个有限的动态范围。例如，数组[1.0,2.0,3.0,64.0]将被转换为[0.333,0.333,0.333,64.0]。但对于规范化的数据，还是可以接受的。作为一个例子，让我们画一个FP4格式的正弦波:
import matplotlib.pyplot as plt import numpy as np from bitsandbytes import functional as bf
length = np.pi * 4 resolution = 256 xvals = np.arange(0, length, length / resolution) wave = np.sin(xvals)
x_4bit, qstate = bf.quantize_fp4(torch.tensor(wave, dtype=torch.float32, device=device), blocksize=64) dq = bf.dequantize_fp4(x_4bit, qstate)
plt.rcParams["figure.figsize"] = (14, 5) plt.title('FP8 Sine Wave') plt.plot(xvals, wave) plt.plot(xvals, dq.cpu().numpy()) plt.show()可以看到精度的损失:
特别说明，在写这篇文章的时候，4位类型NF4只适用于CUDA；目前还不支持CPU计算。
测试
作为本文的最后一步，我们创建一个神经网络模型并对其进行测试。使用transformers库，通过将load_in_4-bit参数设置为True，就可以以4位加载预训练模型。但这并不能让我们理解它是如何工作的。所以我们将创建一个小型神经网络，训练它并以4位精度使用它。
首先，让我们创建一个神经网络模型:
import torch import torch.nn as nn import torch.optim as optim from typing import Any
class NetNormal(nn.Module): def __init__(self): super().__init__() self.flatten = nn.Flatten() self.model = nn.Sequential( nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10) )
def forward(self, x): x = self.flatten(x) x = self.model(x) return F.log_softmax(x, dim=1)我们使用MNIST数据集，数据集分为6万张训练图像和1万张测试图像；可以使用参数train=True|False在DataLoader中指定选择。
from torchvision import datasets, transforms
train_loader = torch.utils.data.DataLoader( datasets.MNIST("data", train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ])), batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader( datasets.MNIST("data", train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ])), batch_size=batch_size, shuffle=True)训练过程以“正常”方式进行，使用默认精度:
device = torch.device("cuda")
batch_size = 64 epochs = 4 log_interval = 500
def train(model: nn.Module, train_loader: torch.utils.data.DataLoader, optimizer: Any, epoch: int): """ Train the model """ model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step()
if batch_idx % log_interval == 0: print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}]\tLoss: {loss.item():.5f}')
def test(model: nn.Module, test_loader: torch.utils.data.DataLoader): """ Test the model """ model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) t_start = time.monotonic() output = model(data) test_loss += F.nll_loss(output, target, reduction='sum').item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset) t_diff = time.monotonic() - t_start
print(f"Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset)}%)\n")
def get_size_kb(model: nn.Module): """ Get model size in kilobytes """ size_model = 0 for param in model.parameters(): if param.data.is_floating_point(): size_model += param.numel() * torch.finfo(param.data.dtype).bits else: size_model += param.numel() * torch.iinfo(param.data.dtype).bits print(f"Model size: {size_model / (8*1024)} KB")
# Train model = NetNormal().to(device) optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5) for epoch in range(1, epochs + 1): train(model, train_loader, optimizer, epoch) test(model, test_loader)
get_size(model)
# Save torch.save(model.state_dict(), "mnist_model.pt")这里还有一个“get_size_kb”方法来获取以kb为单位的模型大小。
训练过程是这样的:
Train Epoch: 1 [0/60000] Loss: 2.31558 Train Epoch: 1 [32000/60000] Loss: 0.53704 Test set: Average loss: 0.2684, Accuracy: 9225/10000 (92.25%)
Train Epoch: 2 [0/60000] Loss: 0.19791 Train Epoch: 2 [32000/60000] Loss: 0.17268 Test set: Average loss: 0.1998, Accuracy: 9401/10000 (94.01%)
Train Epoch: 3 [0/60000] Loss: 0.30570 Train Epoch: 3 [32000/60000] Loss: 0.33042 Test set: Average loss: 0.1614, Accuracy: 9530/10000 (95.3%)
Train Epoch: 4 [0/60000] Loss: 0.20046 Train Epoch: 4 [32000/60000] Loss: 0.19178 Test set: Average loss: 0.1376, Accuracy: 9601/10000 (96.01%)
Model size: 427.2890625 KB我们的简单模型准确率达到96%，神经网络大小为427 KB。
下面我们将“Linear”层替换为“Linear8bitLt”。
from bitsandbytes.nn import Linear8bitLt
class Net8Bit(nn.Module): def __init__(self): super().__init__() self.flatten = nn.Flatten() self.model = nn.Sequential( Linear8bitLt(784, 128, has_fp16_weights=False), nn.ReLU(), Linear8bitLt(128, 64, has_fp16_weights=False), nn.ReLU(), Linear8bitLt(64, 10, has_fp16_weights=False) )
def forward(self, x): x = self.flatten(x) x = self.model(x) return F.log_softmax(x, dim=1)
device = torch.device("cuda")
# Load model = Net8Bit() model.load_state_dict(torch.load("mnist_model.pt")) get_size_kb(model) print(model.model[0].weight)
# Convert model = model.to(device)
get_size_kb(model) print(model.model[0].weight)
# Run test(model, test_loader)结果如下：
Model size: 427.2890625 KB Parameter(Int8Params([[ 0.0071, 0.0059, 0.0146, ..., 0.0111, -0.0041, 0.0025], ..., [-0.0131, -0.0093, -0.0016, ..., -0.0156, 0.0042, 0.0296]]))
Model size: 107.4140625 KB Parameter(Int8Params([[ 9, 7, 19, ..., 14, -5, 3], ..., [-21, -15, -3, ..., -25, 7, 47]], device='cuda:0', dtype=torch.int8))
Test set: Average loss: 0.1347, Accuracy: 9600/10000 (96.0%)原始模型以标准浮点格式加载;它的大小是一样的，权重看起来像[0.0071,0.0059，…]。模型尺寸缩小了4倍。正如我们所看到的，权重值在相同的范围内，因此转换很容易-在测试运行期间，根本没有准确性损失!
再继续4bit的版本：
from bitsandbytes.nn import LinearFP4, LinearNF4
class Net4Bit(nn.Module): def __init__(self): super().__init__() self.flatten = nn.Flatten() self.model = nn.Sequential( LinearFP4(784, 128), nn.ReLU(), LinearFP4(128, 64), nn.ReLU(), LinearFP4(64, 10) )
def forward(self, x): x = self.flatten(x) x = self.model(x) return F.log_softmax(x, dim=1)
# Load model = Net4Bit() model.load_state_dict(torch.load("mnist_model.pt")) get_model_size(model) print(model.model[2].weight)
# Convert model = model.to(device)
get_model_size(model) print(model.model[2].weight)
# Run test(model, test_loader)输出如下所示:
Model size: 427.2890625 KB Parameter(Params4bit([[ 0.0916, -0.0453, 0.0891, ..., 0.0430, -0.1094, -0.0751], ..., [-0.0079, -0.1021, -0.0094, ..., -0.0124, 0.0889, 0.0048]]))
Model size: 54.1015625 KB Parameter(Params4bit([[ 95], [ 81], [109], ..., [ 34], [ 46], [ 33]], device='cuda:0', dtype=torch.uint8))
Test set: Average loss: 0.1414, Accuracy: 9579/10000 (95.79%)模型大小减少了8倍，从427 KB减少到54 KB，但准确率下降1%。这怎么可能?至少对这个模型来说，答案很简单：
权重或多或少是均匀分布的，并且精度损失不是太大。
神经网络使用Softmax作为输出，最大值的索引决定了实际结果。所以对于寻找最大索引，值本身并不重要。例如，当其他值为0.1或0.2时，该值为0.8或0.9没有任何区别!
我们从测试数据集加载数字并检查模型输出。
dataset = datasets.MNIST('data', train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]))
np.set_printoptions(precision=3, suppress=True) # No scientific notation
data_in = dataset[4][0] for x in range(28): for y in range(28): print(f"{data_in[0][x][y]: .1f}", end=" ") print()打印输出显示了我们想要预测的数字:
让我们看看“标准”模型将返回什么：
# Suppress scientific notation np.set_printoptions(precision=2, suppress=True)
# Predict with torch.no_grad(): output = model(data_in.to(device)) print(output[0].cpu().numpy()) ind = output.argmax(dim=1, keepdim=True)[0].cpu().item() print("Result:", ind)
# > [ -8.27 -13.89 -6.89 -11.13 -0.03 -8.09 -7.46 -7.6 -6.43 -3.77] # > Result: 4最大元素位于第5个位置(numpy数组中的元素从0开始编号)，对应于数字“4”。
这是 8-bit:模型的输出：
# > [ -9.09 -12.66 -8.42 -12.2 -0.01 -9.25 -8.29 -7.26 -8.36 -4.45] # > Result: 44-bit的如下:
# > [ -8.56 -12.12 -7.52 -12.1 -0.01 -8.94 -7.84 -7.41 -7.31 -4.45] # > Result: 4可以看到实际输出值略有不同，但最大索引是保持不变的。
总结
在本文中，我们测试了16位、8位和4位浮点数的不同方案，创建了一个神经网络，并能够以8位和4位精度运行它。通过将精度从标准浮点数降低到4位浮点数，内存占用减少了8倍，但精度损失最小。
就像我们昨天的文章中提到的，即使是4位也已经不是极限了;在GPTQ论文中，提到了将权重量化为2甚至3元(1.5比特!)。还有可以对不同层应用不同量化的ExLlamaV2。