Run multi-host supervised fine-tuning on the Qwen3-14b model using MaxText

This tutorial provides a step-by-step guide for running supervised fine-tuning (SFT) on the Qwen3-14b model using MaxText on Cloud TPU. You'll learn how to build a specialized container image, provision a Google Kubernetes Engine (GKE) cluster with Pathways by using Accelerated Processing Kit (XPK), and execute a multi-host training workload.

Objectives

• Learn how to build a custom MaxText container image optimized for post-training.
• Provision a GKE cluster using XPK with Pathways enabled.
• Convert the Qwen3 14b model from Hugging Face format to MaxText format.
• Run a multi-host SFT training workload on Cloud 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:

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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

• Verify that your user account or service account has the following roles:
  • roles/compute.admin, to create the build VM
  • roles/artifactregistry.admin, to manage the Docker repository
  • roles/storage.admin, to manage the data bucket
  • roles/container.admin, to create and manage the Google Kubernetes Engine cluster
  • roles/iam.serviceAccountAdmin, to create the workload service account
  • roles/resourcemanager.projectIamAdmin, to set Identity and Access Management (IAM) policies
  • roles/iam.serviceAccountUser, to act as the service account
• Install and initialize the Google Cloud CLI.
• Verify that you've installed Python 3.12 or later on your workstation.

• You need a Hugging Face access token to use this tutorial. You can sign up for a free account at Hugging Face. Once 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.

Set up the environment

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="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

Replace the following:

• YOUR_PROJECT_ID: your Google Cloud project ID
• YOUR_REGION: the region that you want to use
• YOUR_ZONE: the zone that you want to use
• YOUR_CLUSTER_NAME: a name for your Google Kubernetes Engine cluster
• YOUR_GCS_BUCKET: a unique name for your Cloud Storage bucket
• YOUR_RESERVATION_NAME: your capacity reservation
• YOUR_HF_TOKEN: your Hugging Face access token

Prepare your MaxText container image

To prepare your MaxText container image, including installing required dependencies, complete the following steps:

1. Create a Cloud Storage bucket:

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

2. Create an Artifact Registry repository:

   gcloud artifacts repositories create maxtext-images \
       --repository-format=docker \
       --location=$REGION \
       --project=$PROJECT \
       --description="Docker repository for MaxText images in $REGION" || true

3. Create a file in the root directory of your repository with the filename cloudbuild.yaml and the following contents:

   steps:
     - name: 'gcr.io/cloud-builders/docker'
       entrypoint: 'bash'
       args:
         - '-c'
         - |
           set -euo pipefail
   
           # 0. Install prerequisites (if needed)
           apt-get update && apt-get install -y curl || apk add curl || true
   
           # 1. Install uv
           curl -LsSf https://astral.sh/uv/install.sh | sh
           source $$HOME/.local/bin/env
   
           # 2. Setup Python environment and install MaxText runner
           uv venv --python 3.12 --seed maxtext_venv
           source maxtext_venv/bin/activate
           uv pip install maxtext[runner]==0.2.1 --resolution=lowest
   
           # 3. Build the Docker image (Cloud Build has Docker pre-configured)
           build_maxtext_docker_image WORKFLOW=post-training
   
           # 4. Tag the image properly
           docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}
   
   # Cloud Build automatically pushes images listed here
   images:
     - '${_CLOUD_IMAGE_NAME}'
   
   options:
     # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
     machineType: 'E2_HIGHCPU_32'

4. Use Cloud Build to build your MaxText Docker image:

   gcloud builds submit . \
       --project=$PROJECT \
       --region=$REGION \
       --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

Create your Google Kubernetes Engine cluster

To run SFT training on the Qwen3 14b model, you need a Google Kubernetes Engine cluster equipped with TPU chips. Install Accelerated Processing Kit (XPK) and create a GKE cluster with Pathways support.

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

Prepare the model for training

Convert the base model to MaxText format by using a CPU-based workload. Don't run this task on multiple machines in parallel. The following command includes a check to ensure that the conversion runs on only one TPU node.

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

Track the progress of the model conversion

To track the progress of the conversion, do the following:

1. To list the pods scheduled on your GKE cluster, run the command kubectl get pod.
2. Find the pod named qwen-hf-to-mt-slice-job-0-0-HASH.
3. To inspect the pod's output in real time, run the command kubectl logs -f POD_NAME.

Start the training workload

After the conversion process has completed, you can start the SFT fine-tuning workload by using XPK.

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

Monitor the training workload

Monitor the status of your workload by using the XPK command-line interface (CLI).

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

To view logs and TPU utilization, use the Google Cloud console. You can also view logs by running the following command:

kubectl logs -f qwen-training-pathways-head-0-0-HASH

Replace HASH with the numeric hash in your pod's name. To verify the value of this hash, run the command kubectl get pod and examine the returned list of pods.

Convert the trained model back into Hugging Face format

After your training workload has completed, convert the checkpoints back to Hugging Face format.

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

To track the progress of the conversion, run the command kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH, replacing HASH with the numeric hash in your pod's name.

After the conversion has completed, your tuned model stored in gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ is ready to be used.

Clean up

To avoid incurring additional charges, delete the resources created during this tutorial, including your Google Kubernetes Engine cluster, your Cloud Storage bucket, and your Artifact Registry repository.

Caution: If you plan to continue to use the model that you trained in this tutorial, save a version that you can access before deleting the Cloud Storage bucket containing it.

To delete the resources that you created for this tutorial, run the following command:

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

What's next

• For more information about Cloud TPU, see Introduction to Cloud TPU.
• For architecture and configuration details for the v6e-32 TPU, see TPU v6e.