Skip to content

README.md

Latest commit

 

History

History
324 lines (227 loc) · 12.6 KB

README.md

File metadata and controls

324 lines (227 loc) · 12.6 KB
Raw

Pretrain Llama-3.1-70B workloads on A3 Ultra GKE Node pools using MaxText

This recipe outlines the steps for running a Llama-3.1-70B pretraining workload on A3 Ultra GKE Node pools by using the MaxText framework.

Orchestration and deployment tools

For this recipe, the following setup is used:

Test environment

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

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

Training dataset

This recipe uses a mock pre-training dataset provided by the MaxText framework.

Docker container image

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

us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:toolbox-maxtext-2025-01-10-A3U.

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>

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 queue configured for TAS. The default queue created by the cluster toolkit is a3-ultra. The default queue created by the cluster toolkit is a3-ultra. Make sure to verify the name of the local queue in your cluster.

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/a3ultra/llama3-1-70b/maxtext-pretraining-gke
cd $RECIPE_ROOT

Get cluster credentials

From your client, get the credentials for your cluster.

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

Configure and submit a pretraining job

By default, the job executes 50 training steps.

Using 32 nodes (256 GPUs)

To run the recipe with 'bf16' precision, run the following command:

helm  install -f $RECIPE_ROOT/values.yaml \
    --set-file workload_launcher=$REPO_ROOT/src/launchers/maxtext-launcher.sh \
    --set-file workload_config=$REPO_ROOT/src/frameworks/a3ultra/maxtext-configs/llama3-1-70b-256gpus-a3u-bf16.yaml \
    --set queue=${KUEUE_NAME} \
    --set workload.arguments[0]="base_output_directory=gs://${GCS_BUCKET}/maxtext-experiments" \
    $USER-llama-3-1-70b-maxtext \
    $REPO_ROOT/src/helm-charts/a3ultra/jobset

To run the recipe with fp8 precision, run the following command:

helm  install -f $RECIPE_ROOT/values.yaml \
    --set-file workload_launcher=$REPO_ROOT/src/launchers/maxtext-launcher.sh \
    --set-file workload_config=$REPO_ROOT/src/frameworks/a3ultra/maxtext-configs/llama3-1-70b-256gpus-a3u-fp8.yaml \
    --set queue=${KUEUE_NAME} \
    --set workload.arguments[0]="base_output_directory=gs://${GCS_BUCKET}/maxtext-experiments" \
    $USER-llama-3-1-70b-maxtext \
    $REPO_ROOT/src/helm-charts/a3ultra/jobset

Using 64 nodes (512 GPUs)

To run the recipe with 'bf16' precision, run the following command:

helm  install -f $RECIPE_ROOT/values.yaml \
    --set-file workload_launcher=$REPO_ROOT/src/launchers/maxtext-launcher.sh \
    --set-file workload_config=$REPO_ROOT/src/frameworks/a3ultra/maxtext-configs/llama3-1-70b-512gpus-a3u-bf16.yaml \
    --set queue=${KUEUE_NAME} \
    --set workload.arguments[0]="base_output_directory=gs://${GCS_BUCKET}/maxtext-experiments" \
    --set workload.gpus=512 \
    $USER-llama-3-1-70b-maxtext \
    $REPO_ROOT/src/helm-charts/a3ultra/jobset

To run the recipe with fp8 precision, run the following command:

helm  install -f $RECIPE_ROOT/values.yaml \
    --set-file workload_launcher=$REPO_ROOT/src/launchers/maxtext-launcher.sh \
    --set-file workload_config=$REPO_ROOT/src/frameworks/a3ultra/maxtext-configs/llama3-1-70b-512gpus-a3u-fp8.yaml \
    --set queue=${KUEUE_NAME} \
    --set workload.arguments[0]="base_output_directory=gs://${GCS_BUCKET}/maxtext-experiments" \
    --set workload.gpus=512 \
    $USER-llama-3-1-70b-maxtext \
    $REPO_ROOT/src/helm-charts/a3ultra/jobset

Configure job settings

You can overwrite any of the default MaxText configurations. To do this, set the new arguments using --set workload.arguments.

Examples

To set the number of training steps to 100, run the following command:

helm  install -f $RECIPE_ROOT/values.yaml \
    --set-file workload_launcher=$REPO_ROOT/src/launchers/maxtext-launcher.sh \
    --set-file workload_config=$REPO_ROOT/src/frameworks/a3ultra/maxtext-configs/llama3-1-70b-256gpus-a3u-bf16.yaml \
    --set queue=${KUEUE_NAME} \
    --set workload.arguments[0]="base_output_directory=gs://${GCS_BUCKET}/maxtext-experiments steps=100" \
    $USER-llama-3-1-70b-maxtext \
    $REPO_ROOT/src/helm-charts/a3ultra/jobset

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-llama-3-1-70b-maxtext.

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-llama-3-1-70b-maxtext-workload-0-0-s9zrv.

Analyze results

When completed, the job creates tensorboard logs in the following location:

gs://${GCS_BUCKET}/maxtext-experiments/$JOB_ID/tensorboard/$JOB_ID/
├── events.out.tfevents....
...

To inspect the text logs generated by MaxText, retrieve them from any Pod in the job using the following command: kubectl logs "<pod_name>"

Here is an example of an entry in :

completed step: 12, seconds: 15.516, TFLOP/s/device: 508.371, Tokens/s/device: 1055.949, total_weights: 4194304, loss: 0.000

The logs will show you the step time in seconds and the TFLOP/s/device.

Calculate training performance metrics (MFU)

This section explains how to calculate the Model FLOPS Utilization (MFU), using the logs from the pods. Using the example logs from the previous step, and considering the number of TFLOP/s/device of 508.371, you can compute the MFU using the following formula:

           TFLOP/s/device       508.371
MFU =   ------------------- =  --------- = 0.514 = 51.4%
             MAX TFLOP H200       989

MAX TFLOP H200:

  • BF16: 989
  • FP8: 1979

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-llama-3-1-70b-maxtext. 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-llama-3-1-70b-maxtext).

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 RELEASE_NAME

Replace RELEASE_NAME with the name of your release. E.g.

  • $USER-llama-3-1-70b-maxtext
  • $USER-llama-3-1-70b-maxtext-64nodes
  • $USER-llama-3-1-70b-maxtext-fp8
  • $USER-llama-3-1-70b-maxtext-fp8-64nodes

Running the recipe on a cluster that doesn't use the default configuration.

If you created your cluster using the GKE environment setup guide, it's configured with default settings that include the names for networks and subnetworks used for communication between:

  • The host to external services.
  • GPU-to GPU communication.

For clusters with this default configuration, the Helm chart can automatically generate the required networking annotations in a Pod's metadata. Therefore, you can use the streamlined command to install the chart, as described in the the Configure and submit a pretraining job section.

To configure the correct networking annotations for a cluster that uses non-default names for GKE Network resources, provide the names of the GKE Network resources in you cluster when installing the chart. Use the following example command. Be sure to replace the example values with the actual names of your cluster's GKE Network resources:

helm install -f $RECIPE_ROOT/values.yaml \
    --set-file workload_launcher=$REPO_ROOT/src/launchers/maxtext-launcher.sh \
    --set-file workload_config=$REPO_ROOT/src/frameworks/a3ultra/maxtext-configs/llama3-1-70b-256gpus-a3u-bf16.yaml \
    --set workload.arguments[0]="base_output_directory=gs://${GCS_BUCKET}/maxtext-experiments" \
    --set queue=${KUEUE_NAME} \
    --set network.subnetworks[0]=default \
    --set network.subnetworks[1]=gvnic-1 \
    --set network.subnetworks[2]=rdma-0 \
    --set network.subnetworks[3]=rdma-1 \
    --set network.subnetworks[4]=rdma-2 \
    --set network.subnetworks[5]=rdma-3 \
    --set network.subnetworks[6]=rdma-4 \
    --set network.subnetworks[7]=rdma-5 \
    --set network.subnetworks[8]=rdma-6 \
    --set network.subnetworks[9]=rdma-7 \
    $USER-llama-3-1-70b-maxtext \
    $REPO_ROOT/src/helm-charts/a3ultra/jobset