在 Google TPU 上微调 Llama 模型

使用单程序多数据 (SPMD) 在 Google 张量处理单元 (TPU) 上训练大型语言模型 (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://localhost:8888/tree?token=3ceb24619d0a2f99acf5fba41c51b475b1ddce7cadb2a133

由于我们将使用门控 llama 模型，因此我们需要使用 Hugging Face 令牌登录

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