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