使用 MaxText 在 TPU VM 上執行強化學習訓練

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

工作負載摘要

在 Cloud TPU 上執行這項 RL 訓練後工作負載時，需要下列元件：

  • 模型：Llama 3.1 8B Instruct (llama3.1-8b-Instruct)，這款開放權重的 LLM 經過最佳化，適合用於對話和指令遵循任務。在本教學課程中，模型權重會從 Hugging Face 格式轉換為 MaxText 相容格式 (mt-format)，並準備使用 JAX 在 v6e-8 TPU 上執行分散式作業。

  • 增強學習 (RL) 演算法：群組相對策略最佳化 (GRPO)，其中工作負載會使用 MaxText 內建的 RL 訓練器，執行可擴充的訓練後管道。GRPO 會評估同一提示生成的一組輸出內容，並使用相對獎勵有效最佳化 actor 策略，因此不需要個別的值或評論模型。

  • 資料集和範本：Grade School Math 8K (GSM8K) 是高品質資料集，包含各種語言的國小數學文字題。這項工作負載會使用該模型生成展開、計算獎勵，並在強化學習對齊階段提升多步驟推論能力。

  • 基礎架構和架構：單一 Cloud TPU v6e-8 節點 (ct6e-standard-8t)，使用 MaxText (maxtext[tpu-post-train])，這是由 JAX 支援的高效能訓練堆疊。

目標

  • 設定 Cloud TPU VM 執行個體。
  • 安裝 MaxText 及其依附元件。
  • 將 Hugging Face 模型轉換為 MaxText 格式。
  • 在 TPU 上執行 RL 群組相對政策最佳化 (GRPO) 工作負載。
  • 將訓練好的模型轉換回 Hugging Face 格式，以便提供。

費用

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

如要根據預測用量估算費用，請使用 Pricing Calculator。

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

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

事前準備

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

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

  • 在 Hugging Face 網站上，接受要訓練模型的授權協議。本教學課程使用 llama3.1-8b-Instruct 模型。

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

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

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

設定環境

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export TPU_NAME="YOUR_TPU_NAME"

更改下列內容：

  • 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 "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --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 "${TPU_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.3" --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='llama3.1-8b-Instruct'
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?}"

啟動訓練工作負載

轉換程序完成後，即可啟動 RL 工作負載。

  1. 設定強化學習工作負載訓練參數。

    # -- 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 CHIPS_PER_VM=8
export NUM_BATCHES=50
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.rl.train_rl \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    chips_per_vm="${CHIPS_PER_VM?}" \
    num_batches="${NUM_BATCHES?}" \
    num_test_batches=10 \
    rollout_data_parallelism=1 \
    rollout_tensor_parallelism=-1

將訓練好的模型轉換回 Hugging Face 格式

訓練工作負載完成後，將模型轉換回 Hugging Face 格式。

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

    export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
export HF_MODEL_NAME=llama3.1-8b
export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

  2. 執行轉換，改回 Hugging Face 格式。

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

轉換完成後，儲存在 /dev/shm/$MODEL_NAME/hf-trained 的調整後模型即可使用。VM 重新啟動時，您會失去 /dev/shm 資料夾內容的存取權，因此請將微調後的模型移至永久儲存空間，或上傳至 Hugging Face Hub。

清除所用資源

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

刪除 TPU VM 執行個體

刪除 Cloud TPU VM 執行個體。

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

後續步驟

  • 如要進一步瞭解 Cloud TPU，請參閱「Cloud TPU 簡介」。
  • 如要瞭解 v6e-8 TPU 的架構和設定詳細資料，請參閱「TPU v6e」。