Hugging FaceでのGemmaモデルのファインチューニング

Hugging Face は、パラメータ効率的ファインチューニング（PEFT）を用いて Google DeepMind の Gemma モデルをファインチューニングするための詳細なワークフローを発表しました。このアプローチにより、開発者は 20 億パラメータおよび 70 億パラメータの Gemma モデルを、フルウェイトトレーニングに比べてはるかに少ないメモリと計算リソースで特定のデータセットに適応させることができます。

Gemma向け低ランク適応（LoRA）

Low-Rank Adaptation（LoRA）は、Gemma モデル向けに強調されている主要な PEFT 手法です。すべてのモデルパラメータを更新する代わりに、LoRA は元の重みを固定し、低ランク行列で構成された少数のアダプタ層のみを学習します。これにより、ファインチューニングの計算オーバーヘッドが大幅に削減されます。

Gemma に LoRA を実装するには、Hugging Face の PEFT ライブラリを使用して特定の線形層を対象にできます。提供された例では、以下のモジュールが適応対象となります：

  • q_proj
  • o_proj
  • k_proj
  • v_proj
  • gate_proj
  • up_proj
  • down_proj

QLoRAによるメモリ最適化

メモリ使用量をさらに削減するために、Hugging Face は QLoRA の使用を推奨しています。この手法は bitsandbytes ライブラリを用いてベースモデルを 4 ビット精度に量子化します。from_pretrained メソッドに BitsAndBytesConfig を渡すことで、ユーザーは Gemma をメモリ効率の高いフォーマット（例：nf4 量子化タイプと bfloat16 計算 dtype を使用）でロードでき、コンシューマー向け GPU や Google Colab のような無料プラットフォームでもファインチューニングが可能になります。

実装ワークフロー：引用文の生成

trl ライブラリの SFTTrainer を使用して、Hugging Face は Gemma-2b を特定のフォーマット（引用文の後に著者名が続く）で引用文を生成するようにファインチューニングする実用的な例を示しています。

プロセスは以下の通りです:

  1. モデルのロード: BitsAndBytesConfig を使用した 4 ビット量子化でモデルをロードします。
  2. データセットの準備: Abirate/english_quotes データセットを利用し、モデル用にトークナイズします。
  3. 設定: ランク (r) が 8 の LoraConfig を適用します。
  4. トレーニング: 学習率 2e-4 と paged_adamw_8bit オプティマイザでトレーナーを実行します。

ハードウェアアクセラレーション：PyTorch/XLA と TPU 上の FSDP

Hugging Face の transformers ライブラリにある Gemma モデルは、PyTorch と PyTorch/XLA の両方に最適化されており、GPU と Cloud TPU の両方でのデプロイが可能です。

TPU ユーザー向けに、Hugging Face は SPMD（Single Program, Multiple Data）を通じて Fully Sharded Data Parallel（FSDP）の体験を向上させました。transformers.Trainer に fsdp_config を追加することで、GemmaDecoderLayer をラップし、xla_fsdp_v2 を有効にできます。この統合により、TPU ハードウェア上でのファインチューニング速度が大幅に向上します。

TPU 設定サマリー

機能 設定
ラップする層 GemmaDecoderLayer
FSDP モード full_shard
XLA FSDP v2 有効 (True)
XLA FSDP 勾配チェックポイント 有効 (True)

Sources