在 Google TPU 上微调 Gemma

本教程将介绍如何在 Google Cloud 的 TPU 上微调像 Google Gemma 这样的开源大型语言模型 (LLM)。在我们的示例中，我们将利用 Hugging Face Optimum TPU、🤗 Transformers 和数据集。

Google 的 TPU

Google Cloud TPU 是专门设计的 AI 加速器，经过优化，可用于训练和推理大型 AI 模型。它们非常适用于各种用例，例如聊天机器人、代码生成、媒体内容生成、合成语音、视觉服务、推荐引擎、个性化模型等。

使用 TPU 的优势包括：

旨在以经济高效的方式扩展，适用于各种 AI 工作负载，包括训练、微调和推理。
针对 TensorFlow、PyTorch 和 JAX 进行了优化，并提供各种外形尺寸，包括边缘设备、工作站和基于云的基础设施。
TPU 在 Google Cloud 中可用，并已与 Vertex AI 和 Google Kubernetes Engine (GKE) 集成。

环境设置

对于本示例，单个主机的 v5litepod8 TPU 就足够了。要使用 Pytorch XLA 设置 TPU 环境，此 Google Cloud 指南展示了如何操作。

我们可以使用 ssh 或 gcloud 命令登录远程 TPU。启用端口 8888 的端口转发，例如：

gcloud compute tpus tpu-vm ssh $TPU_NAME \
        --zone=$ZONE \
        -- -L 8888:localhost:8888

一旦我们获得了对 TPU VM 的访问权限，我们就可以克隆包含相关笔记本的 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
# Change directory and launch Jupyter notebook
cd optimum-tpu/examples/language-modeling
jupyter notebook --port 8888

然后我们应该看到熟悉的 Jupyter 输出，显示可以从浏览器访问的地址

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

由于我们将使用受限的 gemma 模型，因此我们需要使用 Hugging Face 令牌登录

!huggingface-cli login --token YOUR_HF_TOKEN

启用 FSDPv2

为了微调 LLM，可能需要将模型分片到 TPU 上，以防止内存问题并提高调优性能。完全分片数据并行 (Fully Sharded Data Parallel) 是一种在 PyTorch 上实现的算法，它允许封装模块以进行分发。在 TPU 上使用 PyTorch/XLA 时，FSDPv2 是一个实用程序，它使用 SPMD (Single Program Multiple Data) 重新表达了著名的 FSDP 算法。在 optimum-tpu 中，可以使用专用辅助函数来使用 FSPDv2。要启用它，可以使用专用函数，该函数应在执行开始时调用。

from optimum.tpu import fsdp_v2


fsdp_v2.use_fsdp_v2()

加载并准备数据集

我们将使用 Dolly，这是一个开源的指令遵循记录数据集，涵盖了 InstructGPT 论文中列出的类别，包括头脑风暴、分类、封闭式问答、生成、信息提取、开放式问答和摘要。

我们将从 Hub 加载数据集

from datasets import load_dataset


dataset = load_dataset("databricks/databricks-dolly-15k", split="train")

我们可以看一个样本

dataset[321]

我们得到一个类似这样的结果

{
    "instruction": "When was the 8088 processor released?",
    "context": "The 8086 (also called iAPX 86) is a 16-bit microprocessor chip designed by Intel between early 1976 and June 8, 1978, when it was released. The Intel 8088, released July 1, 1979, is a slightly modified chip with an external 8-bit data bus (allowing the use of cheaper and fewer supporting ICs),[note 1] and is notable as the processor used in the original IBM PC design.",
    "response": "The Intel 8088 processor was released July 1, 1979.",
    "category": "information_extraction",
}

我们将定义一个格式化函数，它结合了 instruction、context 和 response 字段，并将它们标记化为一个完整的提示。我们将使用与我们打算使用的模型兼容的分词器。

from transformers import AutoTokenizer


model_id = "google/gemma-2b"

tokenizer = AutoTokenizer.from_pretrained(model_id)


def preprocess_function(sample):
    instruction = f"### Instruction\n{sample['instruction']}"
    context = f"### Context\n{sample['context']}" if len(sample["context"]) > 0 else None
    response = f"### Answer\n{sample['response']}"
    # join all the parts together
    prompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])
    prompt += tokenizer.eos_token
    sample["prompt"] = prompt
    return sample

现在可以使用此函数来映射数据集，其中原始列现在可以删除

data = dataset.map(preprocess_function, remove_columns=list(dataset.features))

准备模型进行调优

我们现在可以加载将用于调优的模型。数据集现在已准备好用于微调。

import torch
from transformers import AutoModelForCausalLM


model = AutoModelForCausalLM.from_pretrained(model_id, use_cache=False, torch_dtype=torch.bfloat16)

我们现在将使用参数高效微调 PEFT 和低秩适应 (LoRA) 来高效地在准备好的数据集上微调模型。在 LoraConfig 实例中，我们将定义将进行微调的 nn.Linear 操作。

from peft import LoraConfig


# Set up PEFT LoRA for fine-tuning.
lora_config = LoraConfig(
    r=8,
    target_modules=["k_proj", "v_proj"],
    task_type="CAUSAL_LM",
)

optimum-tpu 的专用函数将帮助我们获取参数，以便我们可以创建训练器实例。

from transformers import TrainingArguments
from trl import SFTTrainer


# Set up the FSDP arguments
fsdp_training_args = fsdp_v2.get_fsdp_training_args(model)

# Set up the trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=data,
    args=TrainingArguments(
        per_device_train_batch_size=64,
        num_train_epochs=32,
        max_steps=-1,
        output_dir="./output",
        optim="adafactor",
        logging_steps=1,
        dataloader_drop_last=True,  # Required for FSDPv2.
        **fsdp_training_args,
    ),
    peft_config=lora_config,
    dataset_text_field="prompt",
    max_seq_length=1024,
    packing=True,
)

一切准备就绪后，调优模型就像调用一个函数一样简单！

trainer.train()

在此之后，我们已成功在 Dolly 数据集上微调了模型。