GPT ネオックスの llm.int ()

これは、レイヤーを llm.int8 () nn.Linear 線形レイヤーに変換するユーティリティ関数を実装します。

LLM.int8 () の論文では、大規模な言語モデルでパフォーマンスを低下させることなく、外れ値を処理する際に int8 量子化を使用してメモリ使用量を削減できることが示されています。重みと入力をスケーリングされた8ビット整数に変換し、行列乗算を行ってint32の結果を生成し、それをfloat16に変換して再スケーリングします。その結果、大規模な言語モデルでは、一部の特徴によって極端な値 (外れ値) がモデルの出力の大部分を占める可能性があることが示されています。これらの機能は8ビットの整数空間に制限されるため、モデルのパフォーマンスが低下します。解決策として、これらの外れ値(指定されたしきい値を超える)を選択し、その乗算をfloat16空間で個別に計算します。外れ値の割合は約 0.01% なので、メモリ使用量は増加せず、モデルのパフォーマンス低下を防ぎます

GPT-NoEx レイヤーを変換するコードは model.py で定義されています。

int8 量子化による GPT-Neox の使用例を以下に示します。

33

bitsandbytes パッケージをインポート

34try:
35    from bitsandbytes.nn import Linear8bitLt, Int8Params
36except ImportError:
37    raise ImportError('''Please install `bitsandbytes` with `pip install bitsandbytes -U`''')
38
39import torch
40from torch import nn

nn.Linear レイヤーを llm.int8 () リニアレイヤーに変換します

  • linear_module nn.Linear 変換するレイヤーです
  • device モデルのデバイスです
  • threshold 外れ値の検出に使用する閾値です
43def make_llm_int8_linear(linear_module: nn.Linear, device: torch.device, threshold: float = 6.0):

53    assert isinstance(linear_module, nn.Linear)

空の Linear8BitLT モジュールを作成します。

56    int8_lin = Linear8bitLt(
57        linear_module.in_features,
58        linear_module.out_features,
59        linear_module.bias is not None,
60        has_fp16_weights=False,
61        threshold=threshold,
62    )

ウェイトをクオンタイズ

65    int8_lin._parameters['weight'] = Int8Params(linear_module.weight.data.cpu(),
66                                                requires_grad=False,
67                                                has_fp16_weights=False).to(device)

float16 空間にバイアスを設定します。

70    if linear_module.bias is not None:
71        int8_lin._parameters['bias'] = nn.Parameter(linear_module.bias.data,
72                                                    requires_grad=False)

75    return int8_lin