MaxText を使用して TPU VM で教師ありファインチューニングを実行する

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

目標

• Cloud TPU VM インスタンスを設定する。
• MaxText とその依存関係をインストールする。
• Hugging Face モデルを MaxText 形式に変換する。
• TPU で SFT トレーニング ワークロードを実行する。
• サービングのために、ファインチューニングしたモデルを Hugging Face 形式に戻す。

費用

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

• Cloud TPU

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

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

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

重要: このチュートリアルを完了するには、v6e-8 TPU にアクセスできる必要があります。この チュートリアルでは、必要なときに容量を使用できるように、プロジェクト Google Cloud でこれらの TPU の予約がすでに行われていることを前提としています。 将来の日時に Cloud TPU 容量を予約する方法については、 をご覧ください。 Cloud TPU の予約について

始める前に

• このチュートリアルを使用するには、Hugging Face アクセス トークンが必要です。Hugging Faceで無料アカウントに登録できます 。アカウントを作成したら、アクセス トークンを生成します。

  1. [Welcome to Hugging Face] ページで、 アカウントのアバターをクリックし、[**Access tokens**] を選択します。
  2. [Access tokens] ページで、[Create new token] をクリックします。
  3. [Read] トークンタイプを選択し、トークンの名前を入力します。
  4. アクセス トークンが表示されます。トークンを安全な場所に保存します。

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

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

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

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

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

環境を設定する

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="ZONE_NAME"
export RESERVATION="RESERVATION_NAME"
export NAME="TPU_MACHINE_NAME"

次のように置き換えます。

• YOUR_PROJECT_ID: 実際の Google Cloud プロジェクト ID
• ZONE_NAME: 使用するゾーン
• RESERVATION_NAME: 容量予約
• TPU_MACHINE_NAME: Cloud TPU VM インスタンスの名前

次のコマンドを実行して、 Google Cloud で認証します。

gcloud auth login

Cloud TPU VM を作成する

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

gcloud alpha compute tpus tpu-vm create $NAME \
    --zone=$ZONE \
    --project=$PROJECT \
    --accelerator-type=v6e-8 \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

VM インスタンスが作成されたら、SSH を使用して接続します。

gcloud compute tpus tpu-vm 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 python3.12 python3.12-venv

uv を使用して、Python パッケージのインストールを高速化します。

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

maxtext_venv という名前の仮想環境を作成して有効にします。

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

MaxText と、トレーニング後のタスクに必要な依存関係をインストールします。

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

次のコマンドを実行して、残りの必要な依存関係をインストールします。

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

トレーニング ワークロードを開始する

変換プロセスが完了したら、SFT ワークロードを開始できます。

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

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   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

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 インスタンスを終了し、削除します。

gcloud alpha compute tpus tpu-vm delete $NAME --zone=$ZONE --project=$PROJECT --quiet

次のステップ

• Cloud TPU の詳細については、 Cloud TPU の概要をご覧ください。
• v6e-8 TPU のアーキテクチャと構成の詳細については、 TPU v6e をご覧ください。