Create an A3 Ultra GKE cluster

This document describes how to deploy an A3 Ultra Google Kubernetes Engine (GKE) cluster. These clusters feature high-performance network topologies for large-scale artificial intelligence (AI) and machine learning (ML) training workloads.

A3 Ultra machine types have NVIDIA H200 SXM GPUs attached and provide the highest network performance in the A3 series. A3 Ultra machine types are ideal for foundation model training and serving.

To learn more about the A3 Ultra accelerator-optimized machine type, see the A3 Ultra machine type section in the Compute Engine documentation.

Before you begin

Before you start, make sure that you have performed the following tasks:

• Enable the Google Kubernetes Engine API.
Enable Google Kubernetes Engine API
• To use the Google Cloud CLI for this task, install and then initialize the gcloud CLI. If you previously installed the gcloud CLI, get the latest version by running the gcloud components update command. Earlier gcloud CLI versions might not support running the commands in this document.
  Note: For existing gcloud CLI installations, make sure to set the compute/region property. If you use primarily zonal clusters, set the compute/zone instead. By setting a default location, you can avoid errors in the gcloud CLI like the following: One of [--zone, --region] must be supplied: Please specify location. You might need to specify the location in certain commands if the location of your cluster differs from the default that you set.

Required roles

To get the permissions that you need to deploy the cluster, ask your administrator to grant you the following IAM roles:

• Kubernetes Engine Admin (roles/container.admin) on the project
• Compute Admin (roles/compute.admin) on the project
• Storage Admin (roles/storage.admin) on the project
• Project IAM Admin (roles/resourcemanager.projectIamAdmin) on the project
• Service Account Admin (roles/iam.serviceAccountAdmin) on the project
• Service Account User (roles/iam.serviceAccountUser) on the project
• Service Usage Consumer (roles/serviceusage.serviceUsageConsumer) on the project
• To consume a single-project reservation: Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1) on the project
• To consume a shared reservation: Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1) on both the owner project and any consumer project where you want to consume the reservation

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.

Choose a consumption option and obtain capacity

To obtain capacity, complete the following steps:

1. Choose a consumption option: Make your choice based on how you want to get and use GPU resources. To learn more, see Choose a consumption option.

   For GKE, consider the following additional information when choosing a consumption option:

   • For more information about flex-start provisioning mode (Preview) and GKE, see About GPU obtainability with flex-start.
   • Flex-start uses best-effort compact placement. To examine your topology, see View the physical topology of nodes in your GKE cluster.
   • You can only get topology information when using Spot VMs if you configure compact placement.
   • For more information about configuring these consumption options in Cluster Toolkit blueprints, see GKE consumption options.

2. Obtain capacity: The process to obtain capacity differs for each consumption option. To learn about the process for your chosen consumption option, see Capacity overview.

Requirements

The following requirements apply to an AI-optimized GKE cluster that uses A3 Ultra instances:

• GPU drivers: The H200 GPUs in A3 Ultra instances require a minimum of the R550 GPU driver version, which is available in GKE 1.31 as the latest driver version. For A3 Ultra, you must set gpu-driver-version=latest with GKE 1.31. For GKE version 1.31.5-gke.1169000 or later, GKE automatically installs the R550 GPU driver versions on A3 Ultra nodes by default.
• Storage: For A3 Ultra node pools, you must set the disk type to hyperdisk-balanced.
• GPUDirect RDMA: To use GPUDirect RDMA with A3 Ultra, use GKE version 1.31.4-gke.1183000 or later.
• Node images: To use GPUDirect RDMA, the GKE nodes must use a Container-Optimized OS node image. Ubuntu and Windows node images aren't supported.

Create an A3 Ultra GKE cluster

To create an A3 Ultra GKE cluster, do the following:

Note: If you create multiple clusters by using the example Cluster Toolkit blueprints, then verify that all VPC network and subnet names are unique per project to prevent errors.

1. Launch Cloud Shell. Although it's possible to use a different environment, we recommend that you use Cloud Shell because the dependencies are already pre-installed for Cluster Toolkit. If you don't want to use Cloud Shell, then follow the instructions to install dependencies to prepare a different environment.
2. Create a Cloud Storage bucket to store the state of the Terraform deployment:

       gcloud storage buckets create gs://BUCKET_NAME \
           --default-storage-class=STANDARD \
           --project=PROJECT_ID \
           --location=COMPUTE_REGION_TERRAFORM_STATE \
           --uniform-bucket-level-access
       gcloud storage buckets update gs://BUCKET_NAME --versioning
       

   Replace the following variables:

   • BUCKET_NAME: the name of the new Cloud Storage bucket.
   • PROJECT_ID: your Google Cloud project ID.
   • COMPUTE_REGION_TERRAFORM_STATE: the compute region where you want to store the state of the Terraform deployment.

3. Choose a consumption option for your GKE deployment:

   By default, the deployment file is configured to consume specific reservations. To use an alternative consumption model (such as Spot VMs or flex-start provisioning mode), update your deployment YAML file. For more information, see GKE consumption options.

4. In the examples/gke-a3-ultragpu/gke-a3-ultragpu-deployment.yaml deployment file from the GitHub repository, fill in the following settings in the terraform_backend_defaults and vars sections to match the specific values for your deployment:

   • BUCKET: the name of the Cloud Storage bucket that you created in the previous step.
   • PROJECT_ID: your Google Cloud project ID.
   • DEPLOYMENT_NAME: a unique name for the deployment, which must be between 6 and 30 characters in length. If the deployment name isn't unique within a project, then cluster creation fails. The default value is gke-a3-ultragpu.
   • REGION: the compute region for the cluster.
   • ZONE: the compute zone for the node pool of A3 Ultra machines. This zone must match the zone where machines are available in your reservation.
   • AUTHORIZED_CIDR: the IP address range that you want to let connect to the cluster. This CIDR block must include the IP address of the machine that calls Terraform. For more information, see How authorized networks work.
   • In the specific reservation section, configure the following settings:
     • In the reservation_affinity block, in the specific_reservations list, specify the reservation name field:
       • To place the node pool anywhere in the reservation, provide the name of your reservation (RESERVATION_NAME).
       • To target a specific block within your reservation, use the reservation and block names in the following format:

         RESERVATION_NAME/reservationBlocks/BLOCK_NAME

         If you don't know which blocks are available in your reservation, then see View the topology of a reservation.
     • STATIC_NODE_COUNT: the number of A3 Ultra nodes in your cluster.

   To modify advanced settings, edit the examples/gke-a3-ultragpu/gke-a3-ultragpu.yaml file.

5. Generate Application Default Credentials (ADC) to provide access to Terraform. If you're using Cloud Shell, use the gcloud auth application-default login command:

       gcloud auth application-default login
       

6. Deploy the blueprint by using the gcluster deploy command to provision the GKE infrastructure with A3 Ultra machine types:

       cd ~/cluster-toolkit
       ./gcluster deploy -d \
           examples/gke-a3-ultragpu/gke-a3-ultragpu-deployment.yaml \
           examples/gke-a3-ultragpu/gke-a3-ultragpu.yaml
       

   Note: If you encounter an error at the JobSet or Kueue installation stage when you deploy your blueprint, then deploy the blueprint again with the addition of the -w flag (./gcluster deploy -w -d ...) to complete the deployment.
7. When prompted, select (A)pply to deploy the blueprint.
   • The blueprint creates VPC networks, a GPU RDMA VPC network, service accounts, a cluster, and a node pool.
   • To support the fio-bench-job-template Job template in the blueprint, Google Cloud buckets, network storage, and persistent volume resources are created.

Run Dynamic Workload Scheduler sample Jobs

If you provisioned your cluster using Dynamic Workload Scheduler (either flex-start or queued provisioning), then you can validate workload scheduling by submitting sample Kubernetes Jobs.

For more information about Dynamic Workload Scheduler provisioning modes, see Dynamic Workload Scheduler flex-start with queued provisioning.

Submit a sample Job

Select the tab that corresponds to your consumption option:

Flex-start

This task is optional and applies only if you provisioned your cluster with flex-start. Workloads configured for flex-start require the cloud.google.com/gke-flex-start: "true" node selector in their Pod specification.

1. Submit the sample flex-start Job:

   kubectl apply -f examples/dws-sample-workloads/sample-job-flex.yaml
   

2. Monitor the status of your Job and Pods:

   kubectl get jobs
   kubectl describe job JOB_NAME
   kubectl get pods
   kubectl describe pod POD_NAME
   

   Replace the following:

   • JOB_NAME: the name of your Job.
   • POD_NAME: the name of your Pod.

3. When the Job completes, delete the Job resources:

   kubectl delete -f examples/dws-sample-workloads/sample-job-flex.yaml
   

Queued provisioning

This task is optional and applies only if you provisioned your cluster with queued provisioning. Workloads configured for queued provisioning require the following label and annotation in their Pod specification:

• Label: kueue.x-k8s.io/queue-name: dws-local-queue
• Annotation: provreq.kueue.x-k8s.io/maxRunDurationSeconds: "MAX_RUN_DURATION"

For more information, see Workload manifest requirements for queued provisioning.

1. Submit the sample queued provisioning Job:

   kubectl apply -f examples/dws-sample-workloads/sample-job-flex-queue.yaml
   

2. Monitor the status of your Job and Pods:

   kubectl get jobs
   kubectl describe job JOB_NAME
   kubectl get pods
   kubectl describe pod POD_NAME
   

   Replace the following:

   • JOB_NAME: the name of your Job.
   • POD_NAME: the name of your Pod.

3. When the Job completes, delete the Job resources:

   kubectl delete -f examples/dws-sample-workloads/sample-job-flex-queue.yaml
   

Test network performance

Validate the networking performance and GPU communication of your provisioned clusters by running NCCL/gIB tests, which are NVIDIA Collective Communications Library (NCCL) tests that are optimized for the Google Cloud environment.

Select the tab for your cluster provisioning model:

Standard

For clusters provisioned with standard or future reservations or Spot VMs, deploy the standard NCCL test JobSet:

1. Deploy the NCCL test JobSet:

   kubectl create -f examples/gke-a3-ultragpu/nccl-jobset-example.yaml
   

2. Confirm that all nccl-test Pods have reached the Completed state:

   kubectl get pods
   

3. Inspect the benchmark results in the primary JobSet Pod logs:

   kubectl logs JOBSET_POD_NAME
   

   Replace JOBSET_POD_NAME with the name of the primary JobSet Pod.

4. When the test completes and Jobs show a Completed status, clean up the test resources:

   kubectl delete -f examples/gke-a3-ultragpu/nccl-jobset-example.yaml
   

Flex-start

For clusters configured with Dynamic Workload Scheduler flex-start, deploy the NCCL test by using the JobSet manifest:

1. Deploy the NCCL test JobSet:

   kubectl create -f examples/gke-a3-ultragpu/nccl-jobset-flex.yaml
   

2. Monitor the status of the JobSet Pods:

   kubectl get pods
   

3. Inspect the benchmark results in the primary JobSet Pod logs:

   kubectl logs JOBSET_POD_NAME
   

   Replace JOBSET_POD_NAME with the name of the primary JobSet Pod.

4. When the test completes and Jobs show a Completed status, clean up the test resources:

   kubectl delete -f examples/gke-a3-ultragpu/nccl-jobset-flex.yaml
   

Queued provisioning

For clusters configured with flex-start with queued provisioning, deploy the NCCL test by using the queued provisioning JobSet manifest:

1. Deploy the NCCL test JobSet:

   kubectl create -f examples/gke-a3-ultragpu/nccl-jobset-flex-queue.yaml
   

2. Monitor the status of the JobSet Pods:

   kubectl get pods
   

3. Inspect the benchmark results in the primary JobSet Pod logs:

   kubectl logs JOBSET_POD_NAME
   

   Replace JOBSET_POD_NAME with the name of the primary JobSet Pod.

4. When the test completes and Jobs show a Completed status, clean up the test resources:

   kubectl delete -f examples/gke-a3-ultragpu/nccl-jobset-flex-queue.yaml
   

Run reproducible benchmarks

You can reproduce pre-training benchmarks for large machine learning open models on A3 Ultra instances on GKE.

Each recipe provides you with the instructions to complete the following tasks:

• Prepare your environment.
• Run the benchmark.
• Analyze the benchmarks results. This review includes the benchmark results and detailed logs for further analysis.

To view all the available recipes, see the GPU recipes GitHub repository.

Note: The following table lists pre-training benchmark recipes that are tested on A3 Ultra instances on GKE clusters that you create by using Cluster Toolkit.

Models	Framework	Recipe
Llama-3.1-70B	MaxText	32 node workload
Llama-3.1-70B	NeMo	32 node workload
Mixtral-8-7B	NeMo	32 node workload

Clean up resources

To avoid recurring charges for the resources that you used on this page, use the gcluster destroy command to clean up the resources that are provisioned by Cluster Toolkit, including the VPC networks and the GKE cluster:

cd ~/cluster-toolkit
./gcluster destroy CLUSTER_NAME

Replace CLUSTER_NAME with the name of your cluster.

What's next

• To learn about scheduling workloads on your GKE clusters by using Topology-Aware Scheduling (TAS) and Kueue, see Schedule GKE workloads with Topology Aware Scheduling.
• To learn about managing common events that are relevant to GKE clusters and AI workloads, see Manage AI-optimized GKE clusters.
• For information about testing your environment for proper setup and optimization, see Cluster networking optimization overview.
• To explore all supported consumption models for GKE clusters, see GKE consumption options.