24GB コンシューマ GPU で RLHF による 20B LLM のファインチューニング

TL;DR

Hugging Face は trl ライブラリが peft と 8‑bit 量子化に対応したことを発表しました。これにより、20 B パラメータの LLM を単一の 24 GB コンシューマ GPU で RLHF ファインチューニングできるようになりました。マルチ GPU のモデルパラレル構成が不要になり、大規模 RL ファインチューニングが手頃でアクセスしやすくなります。

なぜ RLHF には効率的なファインチューニングが必要か

Reinforcement Learning with Human Feedback (RLHF) は通常、次の 3 つのステップで構成されます。(1) 指示に対する教師ありファインチューニング、(2) 人間のアノテーションから報酬モデルを学習、(3) 報酬モデルを用いた PPO ベースの RL ファインチューニング。RL のステップでは、各 GPU 上にモデルの 2 つのコピー（アクティブとリファレンス）が必要になるため、10 B パラメータを超えるモデルでは単一デバイスのメモリをすぐに超えてしまいます。

TRL: PPO ベースの RL 用ライブラリ

trl は言語モデルの PPO 訓練用のハイレベル API を提供します。🤗 Accelerate を活用して単一デバイスまたは分散環境で実行できます。PPO ループでは、KL 正則化報酬を計算するためにアクティブモデル（更新中）とリファレンスモデル（凍結）の両方が必要で、実質的にメモリ使用量が 2 倍になります。

PEFT と 8 ビット量子化でメモリフットプリントを削減

8‑Bit 行列乗算

  • 8‑bit 量子化 (LLM.int8()) はパラメータあたり 1 バイトで重みを保存し、float32 に比べてモデルサイズを 4 分の 1 に削減します。
  • この手法は各線形層を外れ値処理用の float16 部分と大量部分の int8 に分割し、精度を保ちつつ高速化します。

PEFT による低ランク適応 (LoRA)

  • LoRA は事前学習済み重みを凍結し、注意ブロックの query と value 投影に低ランク行列 (A と B) を注入します。
  • 訓練可能なのはアダプタパラメータだけなので、オプティマイザのメモリ使用量が劇的に減少します。
  • 前方・後方パスは余分な行列乗算のためおおよそ 2 倍遅くなりますが、メモリ節約によりコンシューマハードウェアで 20 B モデルの訓練が可能になります。

24GB GPU で 20B モデルのエンドツーエンドパイプライン

Step 1 – 8‑bit 精度でモデルをロード

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neox-20b",
    load_in_8bit=True,
    device_map="auto",
)

8‑bit でロードするとメモリは約 80 GB (float32) から約 20 GB に削減され、24 GB カードに余裕で収まります。

Step 2 – PEFT で学習可能な LoRA アダプタを付与

from peft import get_peft_model, LoraConfig
config = LoraConfig(r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], bias="none")
model = get_peft_model(model, config)

低ランク行列だけがオプティマイザの状態に保存されるため、オプティマイザのメモリは数ギガバイトから数百メガバイトに削減されます。

Step 3 – リファレンスとアクティブのロジットに同一モデルを使用

PEFT の disable_adapters コンテキストマネージャで LoRA 層を一時的に無効化し、同じ基盤モデルからリファレンスロジットを取得します:

with model.disable_adapter():
    ref_logits = model(input_ids)
# アダプタ有効時にアクティブロジットを計算
active_logits = model(input_ids)

2 つ目のフルモデルコピーは不要となり、さらにメモリ使用量が削減されます。

トレーニングスクリプトの概要

ブログ記事では、20 B GPT‑NeoX モデルでフルワークフローを示す 3 つのスクリプトへのリンクが提供されています:

  1. clm_finetune_peft_imdb.py – IMDB 感情データセットで LoRA アダプタを因果言語モデルとしてファインチューニング（1 エポック）。
  2. merge_peft_adapter.py – LoRA 重みをベースモデルにマージし、推論またはさらなる訓練に使用。
  3. gpt-neo-20b_sentiment_peft.py – IMDB 感情分類器を報酬モデルとして使用し、ポジティブな映画レビューを生成する PPO ファインチューニング。

すべてのスクリプトは NVIDIA RTX 4090（24 GB）上で実行されました。完全な訓練は 🤗 research クラスタの単一 A100 でもテストされています。

結果

  • 損失曲線は IMDB での教師あり LoRA ファインチューニング 1 エポック後に安定した収束を示します。
  • PPO 中は平均報酬が着実に上昇し、モデルがよりポジティブなレビューを生成できるようになることが分かります。
  • パイプライン全体が単一の 24 GB GPU で動作し、RLHF がマルチ GPU クラスタに限定されなくなったことを実証しました。

コミュニティへの影響

  • 参入障壁の低減 – 研究者や開発者はコンシューマグレードのハードウェアで RLHF を試すことができます。
  • オープンソースの再現性 – すべてのコードとアダプタは Hugging Face Hub にホストされており、ファインチューニング成果物の共有が容易です。
  • スケーラブルな基盤 – 同様の手法はデータ並列を加えることで、より大規模なモデルにも拡張可能です。

未解決の質問と今後の課題

  • マルチ GPU スケーリング – 複数 GPU にまたがるデータ並列での統合はどの程度うまく機能するか？
  • 訓練速度 – LoRA のオーバーヘッドを軽減するため、より高速なカーネルや混合精度戦略の検討が必要です。
  • より広範な RL アルゴリズム – PPO がデフォルトですが、他の RL 手法（例: DPO）を組み込むことで適用範囲が広がります。

参考文献

Sources