在 Hugging Face 中微調 Gemma 模型
我們最近宣佈，來自 Google Deepmind 的開放權重語言模型 Gemma，已透過 Hugging Face 向更廣泛的開源社群開放。它提供 20 億和 70 億引數兩種規模，幷包含預訓練和指令微調兩種版本。Gemma 已在 Hugging Face 上線，受 TGI 支援，並可輕鬆在 Vertex Model Garden 和 Google Kubernetes Engine 中進行部署和微調。
Gemma 系列模型也非常適合使用 Colab 提供的免費 GPU 資源進行原型設計和實驗。在這篇文章中，我們將簡要回顧如何使用 Hugging Face Transformers 和 PEFT 庫，在 GPU 和 Cloud TPU 上對 Gemma 模型進行 引數高效微調 (PEFT)，以幫助任何希望在自己資料集上微調 Gemma 模型的使用者。
為什麼選擇 PEFT？
對於語言模型而言，預設的（全權重）訓練方法，即使是中等規模的模型，也往往是記憶體和計算密集型的。一方面，對於依賴 Colab 或 Kaggle 等開放計算平臺進行學習和實驗的使用者來說，這可能 prohibitive (代價高昂)。另一方面，即使對於企業使用者來說，為適應不同領域而調整這些模型的成本也是一個需要最佳化的重要指標。PEFT，即引數高效微調，是一種以低成本實現這一目標的流行技術。
在 GPU 和 TPU 上使用 PyTorch
Hugging Face transformers 中的 Gemma 模型針對 PyTorch 和 PyTorch/XLA 都進行了最佳化。這使得 TPU 和 GPU 使用者都能根據需要訪問和實驗 Gemma 模型。隨著 Gemma 的釋出，我們還改進了 Hugging Face 中 PyTorch/XLA 的 FSDP 體驗。這種透過 SPMD 實現的 FSDP 整合也使得其他 Hugging Face 模型能夠利用 PyTorch/XLA 實現 TPU 加速。在這篇文章中，我們將重點關注 PEFT，特別是針對 Gemma 模型的低秩適應 (LoRA)。想要了解更全面的 LoRA 技術，我們鼓勵讀者參閱 Lialin 等人的《Scaling Down to Scale Up》以及 Belkada 等人這篇優秀的文章。
大型語言模型的低秩適應
低秩適應 (LoRA) 是大型語言模型 (LLM) 的一種引數高效微調技術。它透過凍結原始模型，只訓練分解為低秩矩陣的介面卡層，從而只對模型總引數的一小部分進行微調。 PEFT 庫提供了一個簡單的抽象，允許使用者選擇應該應用介面卡權重的模型層。
from peft import LoraConfig
lora_config = LoraConfig(
r=8,
target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
task_type="CAUSAL_LM",
)
在這個程式碼片段中，我們將所有 `nn.Linear` 層指定為要進行適應的目標層。
在下面的例子中，我們將利用 Dettmers 等人提出的 QLoRA，將基礎模型量化為 4 位精度，以實現更節省記憶體的微調方案。要使用 QLoRA 載入模型，首先需要在環境中安裝 `bitsandbytes` 庫，然後在載入模型時將一個 `BitsAndBytesConfig` 物件傳遞給 `from_pretrained`。
在我們開始之前
為了訪問 Gemma 模型的工件，使用者需要接受同意書。現在讓我們開始實現吧。
學習引用
假設您已提交同意書，您可以從 Hugging Face Hub 訪問模型工件。
我們首先下載模型和分詞器。我們還包含了一個用於純權重化 (weight only) 量化的 `BitsAndBytesConfig`。
import torch
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
model_id = "google/gemma-2b"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['HF_TOKEN'])
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0}, token=os.environ['HF_TOKEN'])
現在，在開始微調之前，我們用一句名言來測試模型。
text = "Quote: Imagination is more"
device = "cuda:0"
inputs = tokenizer(text, return_tensors="pt").to(device)
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
模型給出了一個合理的補全，但帶有一些額外的詞元。
Quote: Imagination is more important than knowledge. Knowledge is limited. Imagination encircles the world.
-Albert Einstein
I
但這並不完全是我們希望的答案格式。讓我們看看能否透過微調來教會模型生成以下格式的答案。
Quote: Imagination is more important than knowledge. Knowledge is limited. Imagination encircles the world.
Author: Albert Einstein
首先，讓我們選擇一個英語名言資料集 Abirate/english_quotes。
from datasets import load_dataset
data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
現在讓我們使用上面提到的 LoRA 配置來微調這個模型。
import transformers
from trl import SFTTrainer
def formatting_func(example):
text = f"Quote: {example['quote'][0]}\nAuthor: {example['author'][0]}<eos>"
return [text]
trainer = SFTTrainer(
model=model,
train_dataset=data["train"],
args=transformers.TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=10,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit"
),
peft_config=lora_config,
formatting_func=formatting_func,
)
trainer.train()
最後，我們準備好用之前使用過的相同提示再次測試模型。
text = "Quote: Imagination is"
device = "cuda:0"
inputs = tokenizer(text, return_tensors="pt").to(device)
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
這次我們得到了我們喜歡的格式的響應。
Quote: Imagination is more important than knowledge. Knowledge is limited. Imagination encircles the world.
Author: Albert Einstein
透過 TPU 上的 SPMD 使用 FSDP 進行加速
如前所述，Hugging Face transformers 現在支援 PyTorch/XLA 最新的 FSDP 實現。這可以極大地加快微調速度。要啟用它，只需向 `transformers.Trainer` 新增一個 FSDP 配置即可。
from transformers import DataCollatorForLanguageModeling, Trainer, TrainingArguments
# Set up the FSDP config. To enable FSDP via SPMD, set xla_fsdp_v2 to True.
fsdp_config = {
"fsdp_transformer_layer_cls_to_wrap": ["GemmaDecoderLayer"],
"xla": True,
"xla_fsdp_v2": True,
"xla_fsdp_grad_ckpt": True
}
# Finally, set up the trainer and train the model.
trainer = Trainer(
model=model,
train_dataset=data,
args=TrainingArguments(
per_device_train_batch_size=64, # This is actually the global batch size for SPMD.
num_train_epochs=100,
max_steps=-1,
output_dir="./output",
optim="adafactor",
logging_steps=1,
dataloader_drop_last = True, # Required for SPMD.
fsdp="full_shard",
fsdp_config=fsdp_config,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
下一步
我們透過這個從源筆記本改編的簡單示例，演示了應用於 Gemma 模型的 LoRA 微調方法。用於 GPU 的完整 colab 可以在這裡找到，而用於 TPU 的完整指令碼可以在這裡找到。我們對這個開源生態系統的新成員為研究和學習帶來的無限可能性感到興奮。我們鼓勵使用者訪問 Gemma 文件，以及我們的釋出部落格，以獲取更多關於訓練、微調和部署 Gemma 模型的示例。
