使用 Hugging Face 微調 Gemma 模型

發表於2024-03-01
Hugging FaceGemma模型

我們最近宣佈了，來自 Google Deepmind 開放權重的語言模型 Gemma現已透過 Hugging Face 面向更廣泛的開源社群開放。該模型提供了兩個規模的版本：20 億和 70 億引數，包括預訓練版本和經過指令調優的版本。它在 Hugging Face 平臺上提供支援，可在 Vertex Model Garden 和 Google Kubernetes Engine 中輕鬆部署和微調。

Gemma 模型系列同樣非常適合利用 Colab 提供的免費 GPU 資源進行原型設計和實驗。在這篇文章中，我們將簡要介紹如何在 GPU 和 Cloud TPU 上，使用 Hugging Face Transformers 和 PEFT 庫對 Gemma 模型進行引數高效微調（PEFT），這對想要在自己的資料集上微調 Gemma 模型的使用者尤其有用。

為什麼選擇 PEFT?

即使對於中等大小的語言模型，常規的全引數訓練也會非常佔用記憶體和計算資源。對於依賴公共計算平臺進行學習和實驗的使用者來說，如 Colab 或 Kaggle，成本可能過高。另一方面，對於企業使用者來說，調整這些模型以適應不同領域的成本也是一個需要最佳化的重要指標。引數高效微調（PEFT）是一種以低成本實現這一目標的流行方法。

瞭解更多 PEFT 請參考文章：🤗 PEFT: 在低資源硬體上對十億規模模型進行引數高效微調

在 GPU 和 TPU 上使用 PyTorch 進行 Gemma 模型的高效微調

在 Hugging Face 的 transformers 中，Gemma 模型已針對 PyTorch 和 PyTorch/XLA 進行了最佳化，使得無論是 TPU 還是 GPU 使用者都可以根據需要輕鬆地訪問和試驗 Gemma 模型。隨著 Gemma 的釋出，我們還改善了 PyTorch/XLA 在 Hugging Face 上的 FSDP 使用體驗。這種 FSDP 透過 SPMD 的整合還讓其他 Hugging Face 模型能夠透過 PyTorch/XLA 利用 TPU 加速。本文將重點介紹 Gemma 模型的 PEFT 微調，特別是低秩適應（LoRA）。

想要深入瞭解 LoRA 技術，我們推薦閱讀 Lialin 等人的 "Scaling Down to Scale Up" 以及 Belkada 等人的 精彩文章。

使用低秩適應技術 (LoRA) 對大語言模型進行微調

低秩適應（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 層視為要適應的目標層。

在以下示例中，我們將利用 QLoRA，出自 Dettmers 等人，透過 4 位精度量化基礎模型，以實現更高的記憶體效率微調協議。透過首先在您的環境中安裝 bitsandbytes 庫，然後在載入模型時傳遞 BitsAndBytesConfig 物件，即可載入具有 QLoRA 的模型。

開始之前

要訪問 Gemma 模型檔案，使用者需先填寫 同意表格。

我們繼續。

微調 Gemma，讓它學會並生成一些“名言金句”

假設您已提交同意表格，您可以從 Hugging Face Hub 獲取模型檔案。

我們首先下載模型和分詞器 (tokenizer)，其中包含了一個 BitsAndBytesConfig 用於僅限權重的量化。

import torch
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'])

在開始微調前，我們先使用一個相當熟知的名言來測試一下 Gemma 模型：

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))

模型完成了一個合理的補全，儘管有一些額外的 token:

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

首先，我們選擇一個英文“名人名言”資料集：

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]}"
    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 模型的示例。

相關文章