README.md

Finetune PaliGemma2 workloads on A4 GKE Node pools with Hugging Face Accelerate

This recipe outlines the steps for running a PaliGemma2 finetune workload on A4 GKE Node pools by using the Hugging Face Accelerate.

Orchestration and deployment tools

For this recipe, the following setup is used:

• Orchestration - Google Kubernetes Engine (GKE)
• Finetuning job configuration and deployment - A Helm chart is used to configure and deploy the Kubernetes Jobset resource.

Test environment

This recipe has been optimized for and tested with the following configuration:

• GKE cluster
  • A regional standard cluster version: 1.32.4-gke.1236000 or later.
  • A GPU node pool with 1, 2 or 4 a4-highgpu-8g provisioned using the DENSE deployment type.
  • Workload Identity Federation for GKE enabled.
  • Cloud Storage FUSE CSI driver for GKE enabled.
  • DCGM metrics enabled.
  • Kueue and JobSet APIs installed.
  • Kueue configured to support Topology Aware Scheduling.
• A regional Google Cloud Storage (GCS) bucket to store logs generated by the recipe runs.

To prepare the required environment, see GKE environment setup guide.

Training dataset

This recipe uses the merve/vqav2-small dataset.

Docker container image

This recipe uses the following Deep Learning Software Layer container image:

nvcr.io/nvidia/pytorch:25.01-py3.

This image is based on NVIDIA NeMo 25.02 and contains the NCCL gIB plugin v1.1.0, bundling all NCCL binaries validated for use with A4 GPUs.

Run the recipe

From your client workstation, complete the following steps:

Configure environment settings

Set the environment variables to match your environment:

export PROJECT_ID=<PROJECT_ID>
export CLUSTER_REGION=<CLUSTER_REGION>
export CLUSTER_NAME=<CLUSTER_NAME>
export GCS_BUCKET=<GCS_BUCKET>
export KUEUE_NAME=<KUEUE_NAME>
export HF_TOKEN=<HF_TOKEN>

Replace the following values:

• <PROJECT_ID>: your Google Cloud project ID.
• <CLUSTER_REGION>: the region where your cluster is located.
• <CLUSTER_NAME>: the name of your GKE cluster.
• <GCS_BUCKET>: the name of your Cloud Storage bucket. Don't include the gs:// prefix.
• <KUEUE_NAME>: the name of the Kueue local queue. The default queue created by the cluster toolkit is a4. Make sure to verify the name of the local queue in your cluster.
• <HF_TOKEN>: your Hugging Face token. You can create one here.

Set the default project:

gcloud config set project $PROJECT_ID

Get the recipe

Clone the gpu-recipes repository and set a reference to the recipe folder.

git clone https://github.com/ai-hypercomputer/gpu-recipes.git
cd gpu-recipes
export REPO_ROOT=`git rev-parse --show-toplevel`
export RECIPE_ROOT=$REPO_ROOT/training/a4/paligemma2
cd $RECIPE_ROOT

Get cluster credentials

gcloud container clusters get-credentials $CLUSTER_NAME --region $CLUSTER_REGION

Configure and submit a pretraining job

Update <HF_TOKEN> (your Hugging Face token) in launcher.sh.

Using 4 nodes (32 GPUs) BF16

The default job setting is 50 training steps and fp8 precision. To execute the job with the default settings, run the following command from your client:

helm install $USER-paligemma2 ${RECIPE_ROOT} -f ${RECIPE_ROOT}/values.yaml \
   --set-file workload_launcher=${RECIPE_ROOT}/launcher.sh \
   --set-file workload_config=${RECIPE_ROOT}/main.py \
   --set workload.image=nvcr.io/nvidia/pytorch:25.01-py3 \
   --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \
   --set volumes.gcsMounts[0].mountPath=/job-logs \
   --set workload.envs[0].value=/job-logs/${user}-paligemma2

Configure job settings

You can overwrite any of the default training configuration envs for this job. To do this, we can set the new env values in launcher.sh

Examples

• To set the number of PER_DEVICE_TRAIN_BATCH_SIZE to 64, update following in launcher.sh.

export PER_DEVICE_TRAIN_BATCH_SIZE=64

Run the previous helm command from client.

Monitor the job

To check the status of pods in your job, run the following command:

kubectl get pods | grep JOB_NAME_PREFIX

Replace the following:

• JOB_NAME_PREFIX - your job name prefix. For example $USER-paligemma2.

To get the logs for one of the pods, run the following command:

kubectl logs POD_NAME

Information about the training job's progress, including crucial details such as loss, step count, and step time, is generated by the rank 0 process. This process runs on the pod whose name begins with JOB_NAME_PREFIX-workload-0-0. For example: user-paligemma2-0-0-s9zrv.

Troubleshooting

This section provides guidance on troubleshooting issues with the training job.

To check the status of the job's pods, use the following command:

kubectl get pods | grep JOB_NAME_PREFIX

Replace JOB_NAME_PREFIX with the prefix of your job name. For example, $USER-paligemma2. This command will list all pods associated with the specified job, along with their current status.

To get the logs from a specific pod, use the following command:

kubectl logs POD_NAME

Replace POD_NAME with the name of the pod you want to inspect.

In this recipe, the training job is orchestrated by the Kubernetes JobSet. If the JobSet encounters a fatal failure, it removes all pods, making it impossible to inspect their logs directly. To analyze logs from a failed job, retrieve them from Cloud Logging using the following filter:

resource.type="k8s_container"
resource.labels.project_id="PROJECT_ID"
resource.labels.location="CLUSTER_REGION"
resource.labels.cluster_name="CLUSTER_NAME"
resource.labels.namespace_name="default"
resource.labels.pod_name=~"^JOB_NAME_PREFIX.*"
severity>=DEFAULT

Replace the following:

• PROJECT_ID: your Google Cloud project ID.
• CLUSTER_REGION: the region where your cluster is located.
• CLUSTER_NAME: the name of your GKE cluster.
• JOB_NAME_PREFIX: the prefix of your job name (e.g., $USER-paligemma2).

This filter will retrieve logs from all containers within pods that match the job with the specified name prefix.

Uninstall the Helm release

You can delete the job and other resources created by the Helm chart. To uninstall Helm, run the following command from your client:

helm uninstall $USER-paligmma2