50 年前にカーニハンとリッチーの C 言語書籍の初版が発行されて以来、単精度の「float」型のサイズが 32 ビット、倍精度型のサイズが 64 ビットであることが知られていました。拡張精度を備えた 80 ビットの「long double」型もあり、これらの型は浮動小数点データ処理のほぼすべてのニーズをカバーします。しかし、近年、特に今年の LLM の台頭により、モデルのストレージとメモリのフットプリントを削減するために、開発者は浮動小数点型を可能な限り縮小し始めました。
この記事では、最も一般的な浮動小数点形式を紹介し、簡単なニューラル ネットワークを作成し、その仕組みを理解します。
「標準」32 ビット浮動小数点数
まず標準フォーマットを確認してみましょう。IEEE 754 浮動小数点演算標準は、1985 年に IEEE によって開発されました。32 個の float の一般的な数は次のようになります。
最初のビットはシンボル、次の 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ビット浮動小数点数
初期にはこの形式に対する需要はあまりなく、16 ビット浮動小数点型が IEEE 754 標準に追加されたのは 2008 年になってからでした。これには、符号ビット、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
IEEE 754 標準では 0b11111 が「無限大」用に予約されているため、ここでは 0b11110 が使用されます。同様に、可能な最小値を見つけることもできます。
ieee_754_conversion(0, 0b00001, 0b0000000000, exp_len=5, mant_len=10)
#> 0.00006104
C++ には標準の 16 ビット浮動小数点型がないため、ほとんどの開発者にとって、このような型は「未知の領域」です。
16ビット「bfloat」（BFP16）
この浮動小数点形式は、機械学習専用に設計された Google のチームによって開発されました (名前の「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.203125000000
8ビット浮動小数点(FP8)
この (比較的新しい) 形式は 2022 年に提案され、機械学習用にも作成されました。モデルが大きくなるにつれて、モデルを GPU メモリに取り込むことが困難になります。FP8 フォーマットには、E4M3 (4 ビットの指数と 3 ビットの仮数) と E5M2 (5 ビットの指数と 2 ビットの仮数) の 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
これら 2 種類の正弦波を描いてみましょう。
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 ビットの仮数を持ちます。
2 番目に考えられる 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])])
出力として、2 つのオブジェクトを取得します。1 つは実際に 4 つの数値を含む 16 ビット配列 [117,35]、もう 1 つはスケーリング係数 4.0 と 16 個の FP4 数値すべてのテンソルを含む「状態」オブジェクトです。
たとえば、最初の 4 桁の数字は「0111」(=7) で、ステータス オブジェクトでは、対応する浮動小数点値が 0.25 (0.25 4 = 1.0)であることがわかります。2 番目の数値は「0101」(=5) で、結果は 0.5 4 = 2.0 となります。3 番目の数値「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 コンピューティングは現在サポートされていません。
テスト
この記事の最後のステップとして、ニューラル ネットワーク モデルを作成してテストします。トランスフォーマー ライブラリを使用すると、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 データセットを使用します。これは 60,000 個のトレーニング イメージと 10,000 個のテスト イメージに分割されており、パラメータ 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")
ここには、モデル サイズを kb 単位で取得する「get_size_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
私たちの単純なモデルは、427 KB のニューラル ネットワーク サイズで 96% の精度を達成しました。
次に「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 分の 1 に縮小されました。ご覧のとおり、重量値は同じ範囲内にあるため、変換は簡単です。テスト実行中、精度の低下はまったくありませんでした。
4 ビット バージョンに進みます。
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%)
モデルのサイズは 427 KB から 54 KB へ 8 分の 1 に減少しましたが、精度は 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 ビット: モデルの出力です。
# > [ -9.09 -12.66 -8.42 -12.2 -0.01 -9.25 -8.29 -7.26 -8.36 -4.45]
# > Result: 4
4ビットとは以下の通りです。
# > [ -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 分の 1 に削減されます。
昨日の記事で述べたように、4 ビットですらもはや制限ではありません。GPTQ の論文では、重みを 2 ビットまたは 3 ビット (1.5 ビット!) に量子化することが記載されています。異なるレイヤーに異なる量子化を適用できる ExLlamaV2 もあります。
https://avoid.overfit.cn/post/51c2993a2f824910b241199a52d2c994
著者: ドミトリイ・エリウセフ