メモリを節約！4bit量子化で大規模言語モデルを読み込む方法

はじめに

ファインチューニングなどで学習させた大規模言語モデル(LLM)を、実際に使うためには、まず保存されたモデルデータをプログラムに読み込む必要があります。しかし、高性能なモデルほどサイズが大きく、そのまま読み込もうとすると大量のメモリ(特にGPUメモリ)を消費してしまい、実行環境によってはメモリ不足エラーが発生することもあります。

この記事では、そのような問題を回避し、保存済みのモデルと、それに対応する「トークナイザー」を効率的に読み込むためのPythonコードを解説します。特に、量子化という技術を使ってモデルのメモリ使用量を削減する方法に焦点を当てます。Google Colaboratory (Colab) でGoogle Drive上のファイルを扱うことを想定したコードです。

コードの各部分が何を行っているのか、なぜそのような処理が必要なのかを理解することで、大規模モデルをよりスムーズに扱えるようになるでしょう。

なお、本記事は下記の記事の「3.モデルとトークナイザーのロード（量子化）」の補足記事でもあります。

Google ColabでLLMファインチューニング【実践】:WordPressデータからブログAIアシスタント作成

今回のコード

# --- 設定 ---
# 保存したモデルとトークナイザーのパス
model_load_path = '/content/drive/MyDrive/llm_model'
# ----------------

# 量子化設定 (メモリ削減のため4bit量子化を使用)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16 # T4 GPUではbf16は限定的サポートだが、利用を試みる
)

# モデルのロード
try:
    model = AutoModelForCausalLM.from_pretrained(
        model_load_path,
        quantization_config=bnb_config,
        device_map="auto", # 自動的にGPUに割り当て
        trust_remote_code=True # モデルによっては必要
    )
    print("モデルのロード完了。")
except Exception as e:
    print(f"モデルのロードに失敗しました: {e}")
    print("メモリ不足の可能性があります。ランタイムのタイプを確認するか、バッチサイズ等を調整してください。")
    raise

# トークナイザーのロード
try:
    tokenizer = AutoTokenizer.from_pretrained(model_load_path)
    # パディングトークンの設定 (Llama系ではEOSトークンを使うのが一般的)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id
    # Llama-3系は右パディングが推奨されることが多い
    tokenizer.padding_side = 'right'
    print("トークナイザーのロード完了。")
except Exception as e:
    print(f"トークナイザーのロードに失敗しました: {e}")
    raise

print("Google Driveから量子化モデルとトークナイザーをロードしました。")

コード解説

それでは、コードをステップごとに見ていきましょう。

1. パス設定

# Python
# --- 設定 ---
# 保存したモデルとトークナイザーのパス
model_load_path = '/content/drive/MyDrive/llm_model'
# ----------------

何をしているか まず、読み込みたいモデルとトークナイザーが保存されている場所（ディレクトリパス）を変数 model_load_path に設定しています。ここでは、Google Colab環境でマウントされたGoogle Drive内の llm_model というフォルダを指定しています。
ポイント Google Colabでは、学習させたモデルなどの大きなファイルをGoogle Driveに保存しておき、必要な時に読み込むという使い方をよく行います。/content/drive/MyDrive/ は、ColabでGoogle Driveをマウントした際の標準的なパスです。実際に使用する際は、ご自身のファイル保存場所に合わせてこのパスを変更してください。

2. 量子化設定 (BitsAndBytesConfig)

# Python
# 量子化設定 (メモリ削減のため4bit量子化を使用)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16 # T4 GPUではbf16は限定的サポートだが、利用を試みる
)

何をしているか ここでは、モデルを読み込む際の「量子化」に関する設定を行っています。量子化とは、モデルのパラメータ（数値データ）をより少ないビット数で表現することで、モデル全体のサイズを圧縮し、メモリ使用量を削減する技術です。BitsAndBytesConfig は、この量子化設定を行うためのクラス（Hugging Faceの transformers ライブラリと連携する bitsandbytes ライブラリの一部）です。
ポイント解説
- load_in_4bit=True モデルを4ビット精度で読み込むことを指定します。これにより、メモリ使用量が大幅に削減されます（通常、標準的な32ビット浮動小数点数に比べて約1/8になります）。
- bnb_4bit_use_double_quant=True 「二重量子化」という技術を有効にします。これは、量子化の過程で発生するわずかな情報量をさらに圧縮し、メモリ効率をもう少しだけ高めるためのオプションです。
- bnb_4bit_quant_type="nf4" 4ビット量子化の種類として “nf4” (NormalFloat 4) を指定しています。これは、情報の損失を抑えつつ効率的に量子化できる、比較的新しく標準的な形式の一つです。
- bnb_4bit_compute_dtype=torch.bfloat16 量子化されたモデルが内部で計算を行う際のデータ型を指定します。torch.bfloat16 は、 float32 (32ビット)よりもメモリ効率が良く、 float16 (16ビット)よりも数値の安定性（扱える数値の範囲）が高いという特徴があります。ただし、GPUによってはサポートが限定的な場合があるため、環境に合わせて調整が必要なこともあります。

3. モデルのロードとエラーハンドリング

# Python
# モデルのロード
try:
    model = AutoModelForCausalLM.from_pretrained(
        model_load_path,
        quantization_config=bnb_config,
        device_map="auto", # 自動的にGPUに割り当て
        trust_remote_code=True # モデルによっては必要
    )
    print("モデルのロード完了。")
except Exception as e:
    print(f"モデルのロードに失敗しました: {e}")
    print("メモリ不足の可能性があります。ランタイムのタイプを確認するか、バッチサイズ等を調整してください。")
    raise

何をしているか AutoModelForCausalLM.from_pretrained 関数を使って、指定したパス (model_load_path) から実際にモデルを読み込んでいます。この際、先ほど定義した量子化設定 (bnb_config) を適用しています。また、try-except ブロックを使って、ロード中にエラーが発生した場合の処理を定義しています。
ポイント解説
- AutoModelForCausalLM Hugging Face Transformersライブラリが提供するクラスで、様々な種類のテキスト生成モデル（Causal LM）を統一的に扱うことができます。from_pretrained メソッドは、保存済みのモデルファイルやHugging Face Hub上のモデル名を指定するだけで、簡単にモデルをロードできる便利な機能です。
- quantization_config=bnb_config ここで先ほどの量子化設定を適用し、モデルが4ビットでロードされるように指示しています。
- device_map="auto" モデルの各部分をどのデバイス(CPUやGPU)に配置するかを自動で決定するオプションです。通常、利用可能なGPUがあれば、そのGPU上にモデルがロードされます。これにより、GPUを使った高速な計算が可能になります。
- trust_remote_code=True 一部のモデルは、Hugging Face Hub上で公開されているカスタムコードを必要とします。このオプションをTrueにすることで、そのコードの実行を許可します。セキュリティに関するオプションなので、信頼できる提供元からのモデルを使用する際に指定します。
- try-except モデルのロードは、特にメモリ不足(OOM: Out of Memory)エラーが発生しやすい処理です。try ブロック内の処理でエラーが発生した場合、except ブロックが実行され、エラーメッセージが表示されます。これにより、何が原因でロードに失敗したのかを知る手がかりが得られます。raise は、エラー情報を表示した上でプログラムを意図的に停止させる命令です。

4. トークナイザーのロード

# Python
# トークナイザーのロード
try:
    tokenizer = AutoTokenizer.from_pretrained(model_load_path)
    # パディングトークンの設定 (Llama系ではEOSトークンを使うのが一般的)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id
    # Llama-3系は右パディングが推奨されることが多い
    tokenizer.padding_side = 'right'
    print("トークナイザーのロード完了。")
except Exception as e:
    print(f"トークナイザーのロードに失敗しました: {e}")
    raise

何をしているか モデルと同様に、対応する「トークナイザー」を指定したパスからロードしています。トークナイザーは、人間が使う自然言語のテキストを、モデルが理解できる数値の列（トークンID列）に変換したり、その逆の変換を行ったりする重要なコンポーネントです。ロード後に、パディングに関する設定を行っています。ここでも try-except によるエラーハンドリングを行っています。
ポイント解説
- AutoTokenizer AutoModelForCausalLM と同様に、様々なモデルに対応するトークナイザーを統一的に扱えるクラスです。from_pretrained メソッドで、モデルと同じパスを指定すれば、通常は対応するトークナイザーがロードされます。
- パディングトークンの設定
  - 複数のテキストデータをまとめてモデルに入力する場合、各テキストの長さを揃える必要があります。その際に、短いテキストの後ろ（または前）に追加される特殊なトークンが「パディングトークン (pad_token)」です。
  - if tokenizer.pad_token is None: トークナイザーによっては、デフォルトでパディングトークンが設定されていない場合があります。この if 文でそれをチェックしています。
  - tokenizer.pad_token = tokenizer.eos_token パディングトークンが未設定の場合、文の終わりを示す「EOSトークン (eos_token)」をパディングトークンとして代用しています。これはLlama系のモデルでよく行われる設定です。
  - model.config.pad_token_id = model.config.eos_token_id モデル側の設定にも、パディングトークンIDとしてEOSトークンIDを使うことを反映させています。
  - tokenizer.padding_side = 'right' パディングをトークン列の右側（つまり末尾）に追加するように設定しています。モデルによっては左側パディングが推奨される場合もありますが、Llama-3系などでは右側パディングが推奨されることが多いです。
- try-except トークナイザーのファイルが見つからないなどの理由でロードに失敗した場合に、エラーメッセージを表示します。

5. 完了メッセージ

# Python
print("Google Driveから量子化モデルとトークナイザーをロードしました。")

何をしているか 全ての処理がエラーなく完了した場合に、モデルとトークナイザーが正常にロードされたことを示すメッセージを表示します。

まとめ

今回は、Google Driveに保存された大規模言語モデルとトークナイザーを、4ビット量子化を適用して効率的にロードするPythonコードを解説しました。

BitsAndBytesConfig を使って量子化設定を定義し、メモリ使用量を削減する。
AutoModelForCausalLM と AutoTokenizer の from_pretrained メソッドで、モデルとトークナイザーを簡単にロードする。
device_map="auto" でGPUへの割り当てを自動化する。
トークナイザーのパディング設定を行い、モデルへの入力準備を整える。
try-except でエラーハンドリングを行い、特にメモリ不足エラーなどに対応する。

このコードによって、計算リソースが限られた環境でも大規模モデルを扱いやすくなります。ロードされたモデルとトークナイザーは、この後、テキスト生成などの推論タスクに使用することができます。

※この記事の作成には、一部生成AI（Google Geminiなど）を利用しています。