Multislice and elastic training on TPUs using Ray Train on GKE

This tutorial shows you how to train large language models (LLMs) like Llama 3 70B on Google Kubernetes Engine (GKE) using MaxText, Ray Train, and Multislice Trillium TPUs. This tutorial provides a complete, end-to-end walkthrough, from configuring the necessary secondary data center networking to submitting and successfully running a distributed training workload across 32 physical TPU chips.

This tutorial is for Platform admins, operators, and AI specialists who want to learn how to overcome the memory and networking challenges of training 70-billion parameter models on distributed, multi-host TPU slices.

Background

The combination of GKE, KubeRay, MaxText, and TPUs provides a powerful and scalable platform for large-scale model training. This section describes the key technologies used in this guide:

JAX

JAX is a Python library for accelerator-oriented array computation and program transformation, utilizing the XLA compiler to create highly optimized code that scales efficiently on accelerators.

MaxText

MaxText is a high-performance, open-source LLM framework designed for scalability and customizability. MaxText is built on top of JAX and is optimized to run efficiently on Cloud TPUs.

TPUs

Tensor Processing Units (TPUs) are custom-designed accelerators created by Google to optimize machine learning workloads. Unlike general-purpose CPUs or parallel-processing GPUs, TPUs are highly specialized for the massive matrix and tensor computations at the foundation of deep learning, making them efficient at this specific task. The primary advantage of TPUs is performance at scale.

This tutorial uses TPU Trillium, the sixth generation of TPUs, in a Multislice deployment pattern. Cloud TPU Multislice is where two or more Cloud TPU slices communicate over the data center network (DCN). Multislice enables full-stack, cost-effective, large scale training with near-linear scaling up to tens of thousands of TPU chips. For more information about Multislice, see Cloud TPU Multislice Overview.

KubeRay

KubeRay is a Kubernetes operator that provides a unified way to deploy, manage, and monitor Ray applications on Kubernetes. The KubeRay operator is installed and managed through the Ray on GKE add-on, which is the recommended way to deploy and manage Ray clusters on GKE.

GKE Dynamic Resource Allocation Network (DRANET)

GKE DRANET (Dynamic Resource Allocation Network) is a feature that dynamically attaches high-performance network devices to Pods, bypassing standard Kubernetes networking and enabling high performance over the DCN.

Objectives

This tutorial shows you how to do the following:

1. Set up a GKE cluster with two multi-host TPU node pools.
2. Configure a secondary DCN for cross-slice TPU communication.
3. Configure KubeRay to manage the distributed training environment.
4. Deploy a RayCluster custom resource by using Dynamic Resource Allocation (DRA) for network attachments.
5. Create a Python training script by utilizing Ray Train's JaxTrainer to orchestrate the MaxText training loop across the TPU slices.
6. Run a baseline Llama 3 8B training job.
7. Scale up to Llama 3 70B utilizing 2D sharding (Tensor Parallelism and FSDP) over the DCN.

Before you begin

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

Install the Google Cloud CLI.

If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

To initialize the gcloud CLI, run the following command:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the required APIs:

Roles required to enable APIs

To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

gcloud services enable container.googleapis.com cloudbuild.googleapis.com

Install the Google Cloud CLI.

If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

To initialize the gcloud CLI, run the following command:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the required APIs:

Roles required to enable APIs

To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

gcloud services enable container.googleapis.com cloudbuild.googleapis.com

• Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/container.admin, roles/iam.serviceAccountAdmin, roles/cloudbuild.builds.editor

  gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

  Replace the following:

  • PROJECT_ID: Your project ID.
  • USER_IDENTIFIER: The identifier for your user account. For example, myemail@example.com.
  • ROLE: The IAM role that you grant to your user account.

• Because this tutorial utilizes TPU Trillium (v6e), select a region or zone with availability. For more information, see Cloud TPU quotas.

Prepare your environment

In this tutorial, you use Cloud Shell. Cloud Shell comes preinstalled with the gcloud, helm, and kubectl command-line tools that are used in this tutorial.

1. Go to the Google Cloud console.

2. At the top of the Google Cloud console window, click the Activate Cloud Shell button.

   A Cloud Shell session opens inside a new frame in the Google Cloud console and displays a command-line prompt.

3. In your terminal, clone the kubernetes-engine-samples repository:

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
   

4. Change to the directory containing the sample files:

   cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/maxtext
   

5. Create and activate a Python virtual environment:

   python3 -m venv ray-env
   source ray-env/bin/activate
   

6. Install the Ray CLI:

   pip install "ray[default]==2.55.0"
   

7. Set the following environment variables:

   export PROJECT_ID=$(gcloud config get project)
   export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
   export GS_BUCKET=GS_BUCKET
   export KSA_NAME=KSA_NAME
   export NAMESPACE=default
   export CLUSTER_NAME=CLUSTER_NAME
   export REGION=REGION
   export ZONE=ZONE
   export CLUSTER_VERSION=1.35.2-gke.1842000
   

   Replace the following:

   • GS_BUCKET: the name of the Cloud Storage bucket.
   • KSA_NAME: the name of the Kubernetes Service Account.
   • CLUSTER_NAME: the name of the new cluster.
   • REGION: the region where your TPU Trillium capacity is available.
   • ZONE: the zone where your TPU Trillium capacity is available. For more information, see TPU availability in GKE.

Configure cluster networking for Cloud TPU Multislice

Within a multi-host TPU slice, TPU devices communicate over the high-speed inter-chip interconnects. However, when running Multislice jobs, the TPU slices must communicate with each other over the DCN. Standard Kubernetes Pod networks can bottleneck this traffic. The ct6e-standard-4t machine type is backed by multiple physical network interface cards (NICs). To achieve the best performance, you create two additional VPC networks and use GKE DRANET to connect them directly to the Ray Pods.

1. Create the two additional VPC networks with a large maximum training unit (MTU):

   gcloud compute networks create ${CLUSTER_NAME}-net-1 \
     --subnet-mode=custom \
     --mtu=8896
   
   gcloud compute networks create ${CLUSTER_NAME}-net-2 \
     --subnet-mode=custom \
     --mtu=8896
   

2. Create the dedicated subnets:

   gcloud compute networks subnets create tpu-subnet-1 \
     --network=${CLUSTER_NAME}-net-1 \
     --region=${REGION} \
     --range=10.50.0.0/16
   
   gcloud compute networks subnets create tpu-subnet-2 \
     --network=${CLUSTER_NAME}-net-2 \
     --region=${REGION} \
     --range=10.60.0.0/16
   

Create a GKE cluster

You can configure KubeRay on TPUs in a GKE Autopilot or Standard cluster. We recommend that you use a Autopilot cluster for a fully managed Kubernetes experience. To choose the GKE mode of operation that's the best fit for your workloads, see About GKE modes of operation.

To use GKE managed DRANET, your cluster must use version 1.35.2-gke.1842000 or later for Autopilot mode, or 1.34.1-gke.1829001 or later for Standard mode. This tutorial uses version 1.35.2-gke.1842000.