在 Google TPU 上微調 Llama

在 Google Tensor Processing Units (TPU) 上使用單程式多數據 (SPMD) 架構訓練大型語言模型 (LLM) 具有多種優勢。TPU 提供強大的處理能力，不僅能實現良好的訓練時間，還能讓研究人員高效地實驗更大的模型和資料集。SPMD 架構透過將任務分配到多個 TPU 上，最佳化了資源利用率，進而增強了平行處理能力和可擴展性。使用 SPMD 微調模型最簡單的方法是使用完全分片資料平行（Fully Sharded Data Parallel，FSDP）。Pytorch/XLA 最新且高效能的實作是 FSDP v2，它允許對權重、啟用值 (activations) 和輸出進行分片。

此範例展示如何在單主機 TPU 上微調 Meta 的 Llama 模型。關於 TPU 架構的資訊，您可以參考說明文件。

先修條件

我們假設您已經建立了單主機 TPU VM（例如 v5litepod8 設定），並且可以透過 ssh 存取該機器。您需要複製 optimum-tpu 並安裝一些模組。

git clone https://github.com/huggingface/optimum-tpu.git
# Install Optimum TPU
pip install -e . -f https://storage.googleapis.com/libtpu-releases/index.html
# Install TRL and PEFT for training (see later how they are used)
pip install trl peft
# Install Jupyter notebook
pip install -U jupyterlab notebook
# Optionally, install widgets extensions for better rendering
pip install ipywidgets widgetsnbextension
# This will be necessary for the language modeling example
pip install datasets evaluate accelerate
# Change directory and launch Jupyter notebook
cd optimum-tpu/examples/language-modeling
jupyter notebook --port 8888

接著我們應該會看到熟悉的 Jupyter 輸出，顯示可從瀏覽器存取的位址。

http://:8888/tree?token=3ceb24619d0a2f99acf5fba41c51b475b1ddce7cadb2a133

由於我們將使用受限制的 llama 模型，因此需要使用 Hugging Face 權杖 (token) 進行登入。

!huggingface-cli login --token YOUR_HF_TOKEN

啟用 FSDPv2

若要微調大型語言模型，可能需要將模型分片至多個 TPU 上，以避免記憶體問題並提高微調效能。完全分片資料平行是一種在 Pytorch 上實作的演算法，允許對模組進行封裝以進行分佈。在 TPU 上使用 Pytorch/XLA 時，FSDPv2 是一個使用 SPMD (單程式多數據) 重新演繹著名 FSDP 演算法的工具。在 optimum-tpu 中，可以使用專用輔助程式來使用 FSDPv2。若要啟用它，您可以使用專用函式，並應在執行開始時呼叫它。

from optimum.tpu import fsdp_v2


fsdp_v2.use_fsdp_v2()

接著，需要載入 Tokenizer 和模型。在本範例中，我們將選擇 meta-llama/Llama-3.2-1B。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


model_id = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Add custom token for padding Llama
tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token})
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

若要使用 Abirate/english_quotes 資料集微調模型，您可以載入它並取得 quote 欄位。

from datasets import load_dataset


data = load_dataset("Abirate/english_quotes")


def preprocess_function(samples):
    # Add a simple prompt format to the quotes
    prompts = [f"Generate a quote:\n\n{quote}\n" for quote in samples["quote"]]
    # Add EOS token to each prompt
    prompts = [p + tokenizer.eos_token for p in prompts]
    return {"prompt": prompts}


# data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
data = data.map(preprocess_function, batched=True, remove_columns=data["train"].column_names)

然後，您需要指定 FSDP 訓練引數以啟用分片功能，該函式會自動推斷應分片的類別。

fsdp_training_args = fsdp_v2.get_fsdp_training_args(model)

fsdp_training_args 將指定需要分片的 Pytorch 模組。

fsdp_training_args

現在，訓練過程可以像使用標準 Trainer 類別一樣簡單地完成。

from peft import LoraConfig


lora_config = LoraConfig(
    lora_alpha=128,
    lora_dropout=0.05,
    r=256,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
)

from transformers import TrainingArguments
from trl import SFTTrainer


trainer = SFTTrainer(
    model=model,
    train_dataset=data["train"],
    args=TrainingArguments(
        per_device_train_batch_size=32,
        num_train_epochs=10,
        max_steps=-1,
        output_dir="/tmp/output",
        optim="adafactor",
        logging_steps=1,
        dataloader_drop_last=True,  # Required by FSDP v2
        **fsdp_training_args,
    ),
    peft_config=lora_config,
    dataset_text_field="prompt",
    max_seq_length=512,
    packing=True,
)

trainer.train()

Google TPUs

在 Google TPU 上微調 Llama

先修條件

啟用 FSDPv2