使用 MaxText 在 TPU VM 上執行監督式微調

本教學課程提供逐步指南，說明如何使用 MaxText，在 Google Cloud 的單一 v6e-8 Tensor Processing Unit (TPU) 虛擬機器 (VM) 執行個體上執行監督式微調 (SFT)。MaxText 是以 JAX 為基礎的高效能大型語言模型 (LLM) 訓練堆疊。

工作負載摘要

在 Cloud TPU VM 上執行這項 SFT 工作負載時，會涉及下列元件：

  • 模型：Gemma 3 4B (gemma3-4b)，開放權重的輕量語言模型。在本教學課程中，模型權重會從 Hugging Face 格式轉換為與 MaxText 相容的格式 (mt-format)，並準備使用 JAX 在單一 Cloud TPU v6e-8 VM 上執行。
  • 訓練方法：使用 MaxText 進行 SFT。工作負載會使用 JAX，在單一 Cloud TPU v6e-8 執行個體上執行微調 pipeline。這項工作會使用 MaxText 內建的 SFT 訓練器 (maxtext.trainers.post_train.sft.train_sft) 和對話式指令微調資料集，在 1,000 個步驟中最佳化模型參數。
  • 資料集和範本：UltraChat 200k (HuggingFaceH4/ultrachat_200k)，這是大規模的高品質多輪對話資料集 (以 messages 資料欄分割 train_sft)。用於微調模型，提升對話和指令遵循能力。
  • 基礎架構和架構：單一 Cloud TPU v6e-8 VM 執行個體 (ct6e-standard-8t)。架構為 MaxText (maxtext[tpu-post-train])，由 JAX 提供支援。

目標

  • 設定 Cloud TPU VM 執行個體。
  • 安裝 MaxText 及其依附元件。
  • 將 Hugging Face 模型轉換為 MaxText 格式。
  • 在 TPU 上執行 SFT 訓練工作負載。
  • 將微調後的模型轉換回 Hugging Face 格式，以便提供。

費用

在本文件中，您將使用下列 Google Cloud計費元件：

您可以使用 Pricing Calculator，根據預測用量估算費用。

初次使用 Google Cloud 的使用者可能符合免費試用期資格。

完成本文所述工作後，您可以刪除建立的資源，避免繼續計費，詳情請參閱「清除所用資源」一節。

事前準備

  • 您需要 Hugging Face 存取權杖，才能使用本教學課程。你可以在 Hugging Face 註冊免費帳戶。擁有帳戶後，請產生存取權杖：

    1. 在「Welcome to Hugging Face」(歡迎使用 Hugging Face) 頁面中，按一下帳戶顯示圖片，然後選取「Access tokens」(存取權杖)。
    2. 在「Access tokens」頁面，按一下「Create new token」。
    3. 選取「Read」權杖類型，並輸入權杖名稱。
    4. 系統會顯示存取權杖。將權杖儲存在安全的地方。

  • 在 Hugging Face 網站上，接受您要訓練的模型授權協議。本教學課程使用 gemma3-4b 模型。

如要取得完成本教學課程所需的權限，請要求管理員在專案中授予您下列 IAM 角色：

如要進一步瞭解如何授予角色，請參閱「管理專案、資料夾和組織的存取權」。

您或許也能透過自訂角色或其他預先定義的角色，取得必要權限。

設定環境

執行下列指令碼，設定環境變數：

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 VM 執行個體的名稱

執行下列指令，向 Google Cloud 進行驗證：

gcloud auth login

建立 Cloud TPU VM

建立 Cloud TPU VM 執行個體，並搭載 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}"

建立 VM 執行個體後，請使用 SSH 連線至該執行個體。

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

在 TPU VM 執行個體中完成下列步驟。

安裝 MaxText

更新 TPU VM 執行個體中的系統套件。

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 工作負載。

  1. 設定 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

  2. 啟動訓練工作。在 v6e-8 VM 執行個體上，這項作業約需 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 格式。

  1. 設定匯出路徑和訓練參數。

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

  2. 執行轉換，改回 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 的調整後模型即可使用。VM 重新啟動時，您會失去 /dev/shm 資料夾內容的存取權，因此請將微調後的模型移至永久儲存空間，或上傳至 Hugging Face Hub。

清除所用資源

為避免產生額外費用，請刪除在本教學課程中建立的資源。

刪除 TPU VM 執行個體

退出 Cloud TPU VM 執行個體，然後刪除該執行個體。

# 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」。