在 Google TPU 上微調 Llama

在 Google 張量處理單元 (TPU) 上使用單程式多資料 (SPMD) 訓練大型語言模型 (LLM) 具有諸多優勢。TPU 提供具有競爭力的處理能力，可實現良好的訓練時間，並允許研究人員高效地試驗更大的模型和資料集。SPMD 架構透過在多個 TPU 上分配任務來最佳化資源利用，從而增強並行性和可伸縮性。使用 SPMD 調整模型最簡單的方法是使用完全分片資料並行 (FSDP)。Pytorch/XLA 最新且效能最佳的實現是 FSDP v2，它允許對權重、啟用和輸出進行分片。

本示例展示瞭如何在單主機 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

為了微調 LLM，可能需要將模型分片到 TPU 上，以防止記憶體問題並提高微調效能。完全分片資料並行是一種已在 Pytorch 上實現的演算法，它允許包裝模組以分發它們。在 TPU 上使用 Pytorch/XLA 時，FSDPv2 是一種利用 SPMD（單程式多資料）重新表達著名的 FSDP 演算法的實用程式。在 `optimum-tpu` 中，可以使用專用助手來使用 FSPDv2。要啟用它，可以使用專用函式，該函式應在執行開始時呼叫

from optimum.tpu import fsdp_v2


fsdp_v2.use_fsdp_v2()

然後，需要載入分詞器和模型。在本例中，我們將選擇 `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()