🕌

fluxのモデルflux1-dev-bnb-nf4とQLORAとnf4

2024/09/17に公開

モデル

生成 AI

データ圧縮

tech

はじめに

flux用の軽量モデルとしてflux1-dev-bnb-nf4というものが存在する。bnb-nf4というのは管理上のデータタイプであるが、bnbはbitsandbytesライブラリをさすので実際の名前はnf4になる。ではnf4とはなんであろうか。

一言で言うと従来のfp16と呼ばれるモデルのおおよそ1/4、fp8からはだいたい半分にモデルのサイズを減らしてVRAMが少ないGPUでも動かせる様にする技術でディスクにも優しい技術である。その代償としてAIとしての性能を犠牲にしているけど容量を減らしたほど性能が落ちない様にしてある。しかも同じ4bitならfp4やINT4を使うより精度が良い。

（説明終わり）

QLORA

nf4は"QLORA: Efficient Finetuning of Quantized LLMs"と言う論文に出てくる。QLoRAはnf4を含む複数のテクニックを併用しデータ量を減らしつつ精度を保つアルゴリズムになるようだ。

特徴として

nf4：正規分布した新しい4bitの浮動小数点数
二重量子化：量子化定数を量子化する二重量子化で精度を保ちながらデータ量を減らす
ページドオプティマイザ：GPUのメモリだけではなくコンピュータ上のメモリを利用し最適化する（必要なVRAMを減らすが、速度は遅くなる）

――となっている。ページドオプティマイザはnVidiaのGPUに依存した機能のため二重量子化とnf4がキモになる様である。

nf4とは

nf4データタイプ

nf4の値は-1.0～1.0の間に収まる。つまり-1.0より下と1.0より上の値はnf4に格納出来ない。この範囲にデータを収めるためにデータを扱う場合は正規分布や量子化を利用する。nf4は16種類しか数字が扱えないため以下の数値しか取らない。

    data = [
        -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,
    ]

なお筆者の解釈は正しく無い可能性があるので正確な情報が欲しい場合は英語の論文を読んだ方がよい。

このテクニックはモデルのサイズ削減がメインであり、計算は最終的にfp32やbf16で行われるため速度が速くなるわけではない。内部も4bitで演算できれば4倍速にできるが実際にはfp16に比べてやや遅くなっている。しかしメモリ削減の効果は劇的で本来動作しないはずのGPUでもある程度のモデルまでなら動かせてしまう。

nf4は一段目はLLM.INT8()を利用したデータ量削減のテクニックを使っているのでそれについて調べてみる。

LLM.INT8

浮動小数点数から固定小数点変換図

これは昔ながらの固定小数点演算のテクニックの一種である。

例えば、テストの点数が

333, 460, 789, 800

だとする（面倒なので四人）最低点は333、最高点は800点になる。最高点から最低点を引くと467点になる。そこで点数/467 - 333　と言う計算をする。

{  \frac{Xi}{max(X) - min(X)} - min(X)  }

数式で書くと上記になる。

0 0.261949 0.976445 1

になる。これをINT8の-127から127の範疇に分布させたい場合、-0.5 してから2と127をかけて四捨五入する。

{  (\frac{Xi}{max(X) - min(X)} - min(X) -0.5) \times 2 \times 127  }

すると

-127 -58 121 127

――になる。ここから元の得点を計算する場合、先程の逆を行う。127で割って2で割って、0.5足して、mix(X)を足し、|max(x) -min(x)|をかけると。

333 469 789 800

――と復元出来る。このケースでは綺麗に復元できたが、実際には誤差が出現する。

ここで最後のx127を省略し-1.0 ～ 1.0のママにするとnfになる。8bitならnf8、4bitならnf4になる。

nf4にしてから逆変換すると誤差が出る

333 445 800 800

――――になり最初と数字が変わる。この誤差を許容出来る用途で使うのがnf4であり、AIの重みづけではこの程度の誤差が許容出来るため使うことが可能になる。

実際にAIの重みづけに使う関数はマイナス値があまり出てこない上に0が多いため（使用する関数による）min(X)を0に置き換えることが出来る。

その場合、演算式は

{  \frac{Xi}{absmax(X)}  }

　になる。ここにINT8の変換処理（127をかけるだけ）を追加する。absmaxは符号を無視した中の最大値をさす。

{  round (\frac{Xi \cdot 127}{absmax(X)} ) }

になる。論文と式が違うのはテンソールに対して実行していないから。ここから8で割り切り捨てる（3bit右シフト）でnf4に変換できる。

この処理は1ブロック64で行うことが多い。ブロックを大きくするとデータの精度が落ちてしまうからである。

例えば

 0 0 10 20

と言うデータの場合、absmaxは20になりnf4に変換すると

 -1.0 -1.0 0 1.0

になり、元に戻すと

  0 0 10 20

になる。この場合誤差は出ない。

 0 0 10 20 100 200

と言うデータがある場合、absmaxは200になる。このようにデータが増えるとabsmaxが大きくなる。この場合、nf4に変換すると

 -1.0 -1.0 -1.0 -0.6961928009986877 0 1.0

になる。元に戻すと

 0 0 9 20 101 200

になり、誤差が出る。

このようにデータが増えるほど誤差が大きくなるため1ブロック64で処理を行うことが多い。

二重量子化

量子化係数（割り算する値）はfp32で保存されるので、1ブロック64の場合、32bit/64 = 0.5bit の追加容量が必要になる。つまりnf4の1パラメータは4.5bitになりINT4やfp4よりサイズが大きくなる。そこからに更にデータ量を減らすために量子化係数を量子化するのが二重量子化になる。例えば量子化係数が800, 532, 311, 222の場合、この係数に対して先程と同じことを行う。

この場合、

{ 量子化係数/ 800}

を行い、各量子化係数の精度をfp8に落とす。これによりデータ量を節約する。量子化係数を256個まとめて量子化することで、1パラメータ当たりの追加データ量を、

{  \frac{8 \times 256 + 16 }{256 \times 64} = 8 \div 64 + 32 \div (256 \cdot 64)  = 0.126953125bit  }

に減らせる。これにより1パラメータ当たり概ね

{ 0.5- 0.127 = 0.373bit }

データが減らせる。全体のデータ量は二重量子化をして居ないnf4に比べて1割弱減る。一方、二重量子化に伴うモデル演算時の計算量の増加は、1パラメータあたり1/64回で済む計算になる。

実際の実装になるflux1-dev-bnb-nf4-v2.safetensorsは、flux1-dev-bnb-nf4.safetensorsにくらべて0.5GBデータが増えているが、それは量子化係数のデータ型をnf4からfp32に変更したからとある。この場合二重量子化をする必要がないしていないと思われる。

このデータをどうやって作成するかと言うとtransformersライブラリにnf4の二重量子化を指定しておくと後は勝手にやってくれる。

# from https://huggingface.co/docs/transformers/ja/main_classes/quantization
from transformers import BitsAndBytesConfig

double_quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

model_double_quant = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=double_quant_config)

どのように実装されているかは、~~transformersのコード~~bitsandbytesライブラリのコードを読む必要がある。