モデルの量子化

bitsandbytes の統合

Accelerate は bitsandbytes 量子化をモデルに導入します。これにより、数行のコードで任意の PyTorch モデルを 8 ビットまたは 4 ビットで読み込むことができます。

bitsandbytes を使用して Transformers のモデルを利用する場合は、このドキュメントに従ってください。

bitsandbytesの量子化がどのように機能するかを詳しく知るには、8ビット量子化と4ビット量子化に関するブログ記事をご覧ください。

前提条件

以下の要件をインストールする必要があります。

• bitsandbytesライブラリをインストールします

pip install bitsandbytes

非CUDAデバイスについては、こちらのbitsandbytesインストールガイドを参照してください。

• ソースから最新のaccelerateをインストールする

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

• minGPTとhuggingface_hubをインストールして、サンプルを実行します

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
pip install huggingface_hub

仕組み

まず、モデルを初期化する必要があります。メモリを節約するために、コンテキストマネージャーinit_empty_weights()を使用して空のモデルを初期化できます。

minGPTライブラリからGPT2モデルを取り上げてみましょう。

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    empty_model = GPT(model_config)

次に、モデルの重みのパスを取得する必要があります。パスは、state_dictファイル（例: “pytorch_model.bin”）または、シャード化されたチェックポイントを含むフォルダーです。

from huggingface_hub import snapshot_download
weights_location = snapshot_download(repo_id="marcsun13/gpt2-xl-linear-sharded")

最後に、BnbQuantizationConfigで量子化設定を設定する必要があります。

8ビット量子化の例を以下に示します。

from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)

4ビット量子化の例を以下に示します。

from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")

選択した設定で空のモデルを量子化するには、load_and_quantize_model()を使用する必要があります。

from accelerate.utils import load_and_quantize_model
quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config)

8ビットモデルの保存と読み込み

accelerate を使用して、save_model() で 8 ビットモデルを保存できます。

from accelerate import Accelerator
accelerate = Accelerator()
new_weights_location = "path/to/save_directory"
accelerate.save_model(quantized_model, new_weights_location)

quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")

4ビットモデルのシリアライズは現在サポートされていません。

モジュールをCPUとディスクにオフロード

GPUにモデル全体を格納する十分なスペースがない場合、一部のモジュールをCPU/ディスクにオフロードできます。 これは内部でbig model inferenceを使用しています。詳細については、このドキュメントを参照してください。

8ビット量子化の場合、選択されたモジュールは8ビット精度に変換されます。

4ビット量子化の場合、選択されたモジュールは、ユーザーがBnbQuantizationConfigに渡したtorch_dtypeで保持されます。4ビットのシリアライズが可能になった時点で、これらのオフロードされたモジュールを4ビットに変換するサポートを追加する予定です。

モジュールをCPU/ディスクにオフロードするには、カスタムのdevice_mapを渡すだけです。オフロードされたモジュールは、必要に応じてGPUにディスパッチされます。以下に例を示します。

device_map = {
    "transformer.wte": 0,
    "transformer.wpe": 0,
    "transformer.drop": 0,
    "transformer.h": "cpu",
    "transformer.ln_f": "disk",
    "lm_head": "disk",
}

量子化モデルのファインチューニング

これらのモデルで純粋な8bitまたは4bitのトレーニングを行うことはできません。ただし、パラメータ効率の良いファインチューニング手法（PEFT）を活用して、これらのモデルをトレーニングし、例えばアダプターを上に構築してトレーニングすることは可能です。詳細については、peftライブラリを参照してください。

現在、量子化モデルの上にアダプターを追加することはできません。ただし、Transformersモデルでのアダプターが公式にサポートされているため、量子化モデルのファインチューニングを行うことができます。Transformersモデルのファインチューニングを行いたい場合は、代わりにこのドキュメントに従ってください。4ビットのTransformersモデルをファインチューニングする方法については、このデモをご覧ください。

モデルを学習用に読み込む際にdevice_mapを渡す必要はありません。モデルは自動的にGPUに読み込まれます。device_map=autoは推論でのみ使用してください。

デモ例 - Google Colab上でGPT2 1.5bを実行

GPT2モデルで量子化モデルを実行するためのGoogle Colabのデモをご覧ください。GPT2-1.5BモデルのチェックポイントはFP32形式で、6GBのメモリを使用します。量子化後、8ビットモジュールでは1.6GB、4ビットモジュールでは1.2GBのメモリを使用します。