Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

使用 MaxText 在 TPU 虚拟机上运行监督式微调

本教程提供了一个分步指南，介绍如何使用 MaxText（一种基于 JAX 的高性能训练堆栈，适用于大语言模型 (LLM)） Google Cloud 在单个 v6e-8 张量处理单元 (TPU) 虚拟机 (VM) 实例上运行监督式微调 (SFT)。

工作负载摘要

在 Cloud TPU 虚拟机上运行此 SFT 工作负载涉及以下组件：

模型： Gemma 3 4B (gemma3-4b)，一种开放权重的轻量级语言模型。在本教程中，模型权重将从 Hugging Face 格式转换为 MaxText 兼容格式 (mt-format)，并准备好使用 JAX 在单个 Cloud TPU v6e-8 虚拟机上执行。
训练方法：使用 MaxText 进行 SFT。该工作负载使用 JAX 在单个 Cloud TPU v6e-8 实例上执行微调流水线。它使用 MaxText 的内置 SFT 训练器 (maxtext.trainers.post_train.sft.train_sft) 和以对话为中心的指令调优数据集，在 1000 多个步骤中优化模型参数。
数据集和模板： UltraChat 200k (HuggingFaceH4/ultrachat_200k)，一个大规模、高质量的多轮对话数据集（train_sft 拆分，包含 messages 列）。它用于微调模型并提高其对话和指令遵循能力。
基础架构和框架：单个 Cloud TPU v6e-8 虚拟机实例 (ct6e-standard-8t)。该框架是 MaxText (maxtext[tpu-post-train])，由 JAX 提供支持。

目标

设置 Cloud TPU 虚拟机实例。
安装 MaxText 及其依赖项。
将 Hugging Face 模型转换为 MaxText 格式。
在 TPU 上运行 SFT 训练工作负载。
将微调后的模型转换回 Hugging Face 格式以进行部署。

费用

在本文档中，您将使用的以下收费组件： Google Cloud

Cloud TPU

如需根据您的预计使用量来估算费用，请使用价格计算器。

新 Google Cloud 用户可能有资格申请免费试用。

完成本文档中描述的任务后，您可以通过删除所创建的资源来避免继续计费。如需了解详情，请参阅清理。

准备工作

您需要拥有 Hugging Face 访问令牌才能使用本教程。您可以在 Hugging Face 上注册免费账号。拥有账号后，请生成访问令牌：
1. 在欢迎使用 Hugging Face页面上，点击您的账号头像，然后选择访问令牌。
2. 在访问令牌 页面上，点击创建新令牌 。
3. 选择读取令牌类型，然后输入令牌的名称。
4. 系统会显示您的访问令牌。请将令牌保存在安全的位置。

在 Hugging Face 网站上，接受您计划训练的模型的许可协议。本教程使用模型 gemma3-4b。

如需获得完成本教程所需的权限，请让您的管理员为您授予项目的以下 IAM 角色：

TPU 管理员 (roles/tpu.admin)
Service Account User (roles/iam.serviceAccountUser)
Compute Editor (roles/compute.editor)

如需详细了解如何授予角色，请参阅管理对项目、文件夹和组织的访问权限。

您也可以通过自定义角色或其他预定义角色来获取所需的权限。

设置环境

运行以下脚本来设置环境变量：

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export NAME="YOUR_TPU_NAME"
export NETWORK="default"

替换以下内容：

YOUR_PROJECT_ID：您的 Google Cloud 项目 ID
YOUR_ZONE：您要使用的可用区
YOUR_RESERVATION_NAME：您的容量预留
YOUR_TPU_NAME：Cloud TPU 虚拟机实例的名称

运行以下命令，向 Google Cloud 进行身份验证：

gcloud auth login

创建 Cloud TPU 虚拟机

创建一个 Cloud TPU 虚拟机实例，其中包含 8 个 v6e TPU 芯片，并绑定到您的容量预留。

gcloud compute instances create "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --network="${NETWORK}" \
    --tags="${NAME}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

创建虚拟机实例后，使用 SSH 连接到该实例。

gcloud compute ssh "${NAME}" --zone "${ZONE}" --project "${PROJECT}"

在 TPU 虚拟机实例中完成以下步骤。

安装 MaxText

更新 TPU 虚拟机实例中的系统软件包。

sudo apt update && sudo apt upgrade -y --fix-missing

安装 MaxText 所需的 Python 3.12 及其虚拟环境软件包。

sudo apt install -y build-essential cmake ninja-build

使用 uv 加快 Python 软件包的安装速度。

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

创建一个名为 maxtext_venv 的虚拟环境并将其激活。

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

安装 MaxText 及其后训练任务所需的依赖项。

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.2" --resolution=lowest

运行以下命令安装其余所需的依赖项：

install_tpu_post_train_extra_deps

将模型转换为 MaxText 格式

如需以 MaxText 格式训练模型，您必须将其从 Hugging Face 格式转换为 MaxText 格式。

指定您的环境变量，例如您的 Hugging Face 访问令牌、您要使用的模型的名称，以及您要以 MaxText 格式保存模型的目录。

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='gemma3-4b'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

将 YOUR_HF_TOKEN 替换为您之前创建的 Hugging Face 访问令牌。

如需将模型从 Hugging Face 格式转换为 MaxText 格式，请运行以下脚本。此转换大约需要五分钟才能完成。

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

启动训练工作负载

转换过程完成后，您可以启动 SFT 工作负载。

配置 SFT 工作负载训练参数。

# -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export RUN_NAME
export STEPS=1000
export PER_DEVICE_BATCH_SIZE=1

# -- Dataset configuration --
export DATASET_NAME="HuggingFaceH4/ultrachat_200k"
export TRAIN_SPLIT="train_sft"
export TRAIN_DATA_COLUMNS="['messages']"

export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
export TPU_ACCELERATOR_TYPE=v6e-8
export TPU_WORKER_ID=0
TPU_NAME=$(hostname)
export TPU_NAME
export TPU_SKIP_MDS_QUERY=1
export TPU_WORKER_HOSTNAMES=localhost
export TPU_TOPOLOGY=2x4
export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
export TPU_HOST_BOUNDS=1,1,1

启动训练作业。在 v6e-8 虚拟机实例上，这大约需要 10 分钟。

python3 -m maxtext.trainers.post_train.sft.train_sft \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
    steps="${STEPS?}" \
    hf_path="${DATASET_NAME?}" \
    train_split="${TRAIN_SPLIT?}" \
    train_data_columns="${TRAIN_DATA_COLUMNS?}" \
    profiler=xplane

将训练后的模型转换回 Hugging Face 格式

训练工作负载完成后，将模型转换回 Hugging Face 格式。

设置导出路径和训练后的参数。

export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

运行转换以返回到 Hugging Face 格式。

python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name="${MODEL_NAME}" \
    load_parameters_path="${POST_TRAIN_PATH}" \
    base_output_directory="${HF_EXPORT}" \
    scan_layers=True \
    use_multimodal=False \
    weight_dtype=bfloat16

转换完成后，存储在 /dev/shm/gemma3-4b/hf-trained 中的调优后的模型即可使用。由于虚拟机重新启动时您将无法访问 /dev/shm 文件夹的内容，因此您应将调优后的模型移至永久性存储空间或将其上传到 Hugging Face Hub。

清理

为避免产生额外费用，请删除在本教程中创建的资源。

删除 TPU 虚拟机实例

退出 Cloud TPU 虚拟机实例，然后将其删除。

# 1. Delete TPU instance
echo "Deleting TPU instance: ${NAME}..."
gcloud compute instances delete "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --quiet || true

# 2. Delete IAP firewall rule
echo "Deleting firewall rule: ${FIREWALL_RULE_NAME}..."
gcloud compute firewall-rules delete "${FIREWALL_RULE_NAME}" \
    --project="${PROJECT}" \
    --quiet || true

后续步骤

如需详细了解 Cloud TPU，请参阅 Cloud TPU 简介。
如需了解 v6e-8 TPU 的架构和配置详情，请参阅 TPU v6e。