在 Hugging Face 上微調 Gemma 模型

Hugging Face 已推出一套使用參數效率微調 (PEFT) 來微調 Google DeepMind 的 Gemma 模型的詳細工作流程。此方法讓開發者能夠將 20 億與 70 億參數的 Gemma 模型調整至特定資料集，且所需的記憶體與計算資源遠低於完整權重訓練。

針對 Gemma 的低秩適應 (LoRA)

低秩適應 (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 位元精度。透過將 BitsAndBytesConfig 傳入 from_pretrained 方法，使用者可以以記憶體友善的格式載入 Gemma（例如使用 nf4 量化類型與 bfloat16 計算資料型別），使微調能在消費級 GPU 或如 Google Colab 等免費平台上執行。

實作工作流程：學習產生引言

透過 trl 套件中的 SFTTrainer，Hugging Face 示範了一個實務應用：將 Gemma-2b 微調以產生特定格式的引言（引言後接作者）。

此流程包括：

  1. 模型載入： 透過 BitsAndBytesConfig 以 4 位元量化載入模型。
  2. 資料集準備： 使用 Abirate/english_quotes 資料集並為模型進行斷詞（tokenize）。
  3. 設定： 套用 LoraConfig，其秩 (r) 為 8。
  4. 訓練： 使用學習率 2e-4 與 paged_adamw_8bit 優化器執行訓練器。

硬體加速：PyTorch/XLA 與 TPU 上的 FSDP

Hugging Face transformers 套件中的 Gemma 模型已針對 PyTorch 與 PyTorch/XLA 進行最佳化，讓其可部署於 GPU 與雲端 TPU 上。

對於 TPU 使用者，Hugging Face 透過 SPMD（單程式多資料）提升了 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