Run SFT for Gemma 4 31B on multi-host TPU v6e

This tutorial shows you how to run supervised fine-tuning (SFT) on a multi-host Tensor Processing Unit (TPU) v6e cluster by using MaxText and Cluster Toolkit. You use Cluster Toolkit to execute a multi-host training workload and export the results back to Hugging Face format for serving.

Objectives

  • Install Cluster Toolkit and its dependencies.
  • Deploy a Cluster Toolkit cluster.
  • Convert a Hugging Face model to MaxText format.
  • Run an SFT training workload on the TPU.
  • Convert the fine-tuned model back to Hugging Face format for serving.

Costs

In this document, you use the following billable components of Google Cloud:

To generate a cost estimate based on your projected usage, use the pricing calculator.

New Google Cloud users might be eligible for a free trial.

When you finish the tasks that are described in this document, you can avoid continued billing by deleting the resources that you created. For more information, see Clean up.

Before you begin

Authenticate your user account:

gcloud auth login

Generate Application Default Credentials (ADC):

gcloud auth application-default login

If you're running in a headless or remote SSH session without a browser, append --no-launch-browser to both commands.

You need a Hugging Face access token to use this tutorial. You can sign up for a free account at Hugging Face. After you have an account, generate an access token:

  1. On the Welcome to Hugging Face page, click your account avatar and select Access tokens.
  2. On the Access tokens page, click Create new token.
  3. Select the Read token type and enter a name for your token.
  4. Your access token is displayed. Save the token in a safe place.
  • On the Hugging Face website, accept the license agreement for the model that you plan to train. This tutorial uses the model gemma4-31b.

To get the permissions that you need to complete this tutorial, ask your administrator to grant you the following IAM roles on your project:

For more information about granting roles, see Manage access to projects, folders, and organizations.

You might also be able to get the required permissions through custom roles or other predefined roles.

Set up your environment variables

Set up your environment variables by running the following script:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="gke-tpu-v6e"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

gcloud config set project "${PROJECT}"
gcloud auth application-default set-quota-project "${PROJECT}"

Replace the following:

  • YOUR_PROJECT_ID: the ID of your Google Cloud project.
  • YOUR_REGION: the region where you want to deploy your cluster.
  • YOUR_ZONE: the zone where you want to deploy your cluster.
  • YOUR_RESERVATION_NAME: the name of your reservation.
  • YOUR_HF_TOKEN: your Hugging Face access token.
  • YOUR_BUCKET_NAME: a globally unique name for a Cloud Storage bucket.

Install Cluster Toolkit dependencies

To complete this tutorial from a Linux or macOS client or workstation, follow the relevant steps in Install dependencies in the Cluster Toolkit documentation.

If you're using Cloud Shell, then you can skip this section.

Install Cluster Toolkit

Install the prebuilt bundle for Cluster Toolkit by following the instructions at Install Cluster Toolkit.

Create your Cluster Toolkit cluster

To create and deploy a Cluster Toolkit cluster with 32 v6e TPU chips, complete the following steps:

  1. Create a Cloud Storage bucket:

    gcloud storage buckets create "gs://${GCS_BUCKET}" --project="${PROJECT}" --location="${REGION}" || true

  2. Stage the blueprint configuration and copy the example deployment blueprint to an isolated working directory:

    mkdir -p tmp
cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

  3. By default, your cluster node pool service account doesn't have the required permissions to write to your Cloud Storage bucket. To allow the node pool service account to write to your Cloud Storage bucket, you must grant it the Storage Admin role. To grant this role, edit the file tmp/gke-tpu-v6e-advanced.yaml by updating the node_pool_service_account module:

    - id: node_pool_service_account
  source: modules/project/service-account
  settings:
    name: gke-np-sa
    project_roles:
    - logging.logWriter
    - monitoring.metricWriter
    - monitoring.viewer
    - stackdriver.resourceMetadata.writer
    - storage.admin            # Change from storage.objectViewer
    - artifactregistry.reader

  4. Use the gcluster deploy command to deploy your Cluster Toolkit cluster by using the blueprint tmp/gke-tpu-v6e-advanced.yaml and passing the required variables by using the --vars flag:

    ./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
    --vars "project_id=${PROJECT},deployment_name=${CLUSTER_NAME},region=${REGION},zone=${ZONE},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -l IGNORE \
    -w

  5. Fetch Google Kubernetes Engine cluster credentials for kubectl:

    gcloud container clusters get-credentials "${CLUSTER_NAME}" --location="${REGION}" --project="${PROJECT}"

  6. Configure container registry authentication:

    gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION}-docker.pkg.dev" --quiet

  7. Grant the Storage Admin role (roles/storage.admin) to your Google Kubernetes Engine service accounts:

    gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet || true
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet || true

Convert the model to MaxText format

To train the model in MaxText format, you must convert it from Hugging Face format to MaxText format.

  1. To simplify subsequent commands, use the gcluster job config command to configure your default project, cluster, and location:

    # Configure gcluster Defaults
./gcluster job config set project "${PROJECT}"
./gcluster job config set cluster "${CLUSTER_NAME}"
./gcluster job config set location "${REGION}"

  2. To convert the model from Hugging Face format to MaxText format and store it in your Cloud Storage bucket, run the following script, which uses the gcluster job submit command:

    ./gcluster job submit --name hf-to-mt \
    --compute-type "${TPU_TYPE}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_maxtext \
        model_name=${MODEL_NAME} \
        hf_access_token=${HF_TOKEN} \
        base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/ \
        scan_layers=True \
        use_multimodal=False \
        skip_jax_distributed_system=true \
        checkpoint_storage_use_zarr3=0 \
        checkpoint_storage_use_ocdbt=0 \
        hardware=cpu \
        --lazy_load_tensors=True"

To check the status of the conversion job, run the gcluster job logs command:

# Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs hf-to-mt --main-only -f

Start the training workload

After the conversion process has completed, you can start the SFT workload by running the following command:

./gcluster job submit --name sft \
    --compute-type "${TPU_TYPE}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME}" \
    --await-job-completion \
    --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
      python3 -m maxtext.trainers.post_train.sft.train_sft \
      run_name=sft \
      base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/trained/ \
      model_name=${MODEL_NAME} \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
      hf_access_token=${HF_TOKEN} \
      dataset_type=hf \
      hf_path=HuggingFaceH4/ultrachat_200k \
      per_device_batch_size=1 steps=1000 \
      profiler=xplane \
      checkpoint_storage_use_zarr3=0 \
      checkpoint_storage_use_ocdbt=0 \
      skip_jax_distributed_system=False"

To check the status of the training job, run the following command:

# Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs sft --main-only -f

Convert the trained model back into Hugging Face format

After the training workload has completed, convert the model back to Hugging Face format:

./gcluster job submit --name mt-to-hf \
    --compute-type "${TPU_TYPE}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_huggingface \
        model_name=${MODEL_NAME?} \
        hf_access_token=${HF_TOKEN?} \
        load_parameters_path=gs://${GCS_BUCKET?}/${MODEL_NAME}/trained/sft/checkpoints/1000/model_params/ \
        base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
        skip_jax_distributed_system=true \
        hardware=cpu \
        scan_layers=True \
        use_multimodal=False \
        weight_dtype=bfloat16 \
        --override_model_architecture"

To check the status of the conversion job, run the following command:

# Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs mt-to-hf --main-only -f

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

Clean up

To avoid incurring additional charges, use the gcluster destroy command to delete the resources created during this tutorial.

./gcluster destroy "${CLUSTER_NAME}" --robust
gcloud storage rm -r "gs://${GCS_BUCKET}"

# To delete the local deployment folder
rm -rf .ghpc "${CLUSTER_NAME}"

What's next