Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

MaxText を使用して TPU VM で強化学習トレーニングを実行する

このチュートリアルでは、大規模言語モデル（LLM）用の高性能 JAX ベースのトレーニングスタックである MaxText を使用して、 Google Cloud 上の単一の v6e-8 Tensor Processing Unit（TPU）仮想マシン（VM）インスタンスで強化学習（RL）トレーニングを実行する手順について説明します。

ワークロードの概要

Cloud TPU でこの RL 事後トレーニングワークロードを実行するには、次のコンポーネントが必要です。

モデル: Llama 3.1 8B Instruct（llama3.1-8b-Instruct）。会話と指示に従うタスクに最適化されたオープンウェイト LLM です。このチュートリアルでは、モデルの重みが Hugging Face 形式から MaxText 互換形式（mt-format）に変換され、JAX を使用して v6e-8 TPU で分散実行できるように準備されます。
強化学習（RL）アルゴリズム: Group Relative Policy Optimization（GRPO）。ワークロードは、MaxText の組み込み RL トレーナーを使用して、スケーラブルな事後トレーニングパイプラインを実行します。GRPO では、同じプロンプトに対して生成された出力のグループを評価し、相対報酬を使用してアクターポリシーを効率的に最適化することで、個別の値モデルやクリティックモデルが不要になります。
データセットとテンプレート: Grade School Math 8K（GSM8K）。言語的に多様な小学生の数学の文章問題で構成される高品質のデータセットです。このワークロードでは、ロールアウトの生成、報酬の計算、RL アライメントフェーズでのマルチステップの推論機能の改善に使用されます。
インフラストラクチャとフレームワーク: MaxText（maxtext[tpu-post-train]）を使用する単一の Cloud TPU v6e-8 ノード（ct6e-standard-8t）。JAX を搭載した高性能トレーニングスタックです。

目標

Cloud TPU VM インスタンスを設定する。
MaxText とその依存関係をインストールする。
Hugging Face モデルを MaxText 形式に変換する。
TPU で RL Group Relative Policy Optimization（GRPO）ワークロードを実行する。
トレーニング済みモデルをサービング用に Hugging Face 形式に変換する。

費用

このドキュメントでは、課金対象である次のコンポーネントを使用します。 Google Cloud

Cloud TPU

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。

新規の Google Cloud ユーザーの方は、無料トライアルをご利用いただける場合があります。

このドキュメントに記載されているタスクの完了後、作成したリソースを削除すると、それ以上の請求は発生しません。詳細については、クリーンアップをご覧ください。

始める前に

このチュートリアルを使用するには、Hugging Face アクセストークンが必要です。Hugging Faceで無料アカウントに登録できます。アカウントを取得したら、アクセストークンを生成します。
1. [Welcome to Hugging Face] ページで、アカウントのアバターをクリックし、[アクセストークン] を選択します。
2. [アクセストークン] ページで、[Create new token] をクリックします。
3. [Read] トークンタイプを選択し、トークンの名前を入力します。
4. アクセストークンが表示されます。トークンを安全な場所に保存します。

Hugging Face ウェブサイトで、トレーニングするモデルのライセンス契約に同意します。このチュートリアルでは、モデル llama3.1-8b-Instruct を使用します。

このチュートリアルを完了するために必要な権限を取得するには、プロジェクトに対する次の IAM ロールを付与するよう管理者に依頼してください。

TPU 管理者 (roles/tpu.admin)
サービスアカウントユーザー（roles/iam.serviceAccountUser）
Compute 編集者（roles/compute.editor）

ロールの付与については、プロジェクト、フォルダ、組織へのアクセス権の管理をご覧ください。

必要な権限は、カスタムロールや他の事前定義ロールから取得することもできます。

環境を設定する

次のスクリプトを実行して、環境変数を設定します。

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 を作成する

容量予約にバインドされた 8 個の v6e TPU チップを使用して、Cloud TPU VM インスタンスを作成します。

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 形式に変換するには、次のスクリプトを実行します。この変換には約 5 分かかります。

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 ワークロードを開始できます。

RL ワークロードのトレーニングパラメータを構成します。

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

トレーニングジョブを開始します。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 形式に変換します。

エクスポートとトレーニング済みパラメータのパスを設定します。

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

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 をご覧ください。