recipe

Pretrain GPT3-175B workloads on A3 Mega GKE Node pools

This recipe outlines the steps for running a GPT-3 175B pretraining workload on A3 Mega GKE Node pools by using the NVIDIA NeMo framework.

Orchestration and deployment tools

For this recipe, the following setup is used:

• Orchestration - Google Kubernetes Engine (GKE)
• Job configuration and deployment - A Helm chart is used to configure and deploy the Kubernetes JobSet . This job encapsulates the NVIDIA NeMo Megatron GPT pretraining workload . The chart generates the job's manifest, adhering to best practices for using GPUDirect-TCPXO with Google Kubernetes Engine (GKE), which includes setting optimal values for NVIDIA NCCL and the TCPXO NCCL plugin.

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.1767000 or later.
  • A GPU node pool with 32 a3-megagpu-8g nodes 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 .
  • GPUDirect-TCPXO component versions:
    • NCCL Plugin: v1.0.3
    • RxDM sidecar: v1.0.9
  • NVIDIA NeMo NGC container image : 24.07
• A regional Google Cloud Storage (GCS) bucket to store logs generated by the recipe runs.
• Artifact Registry for storing the Nemo Container image

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

Make sure that you've been granted the following IAM roles:

• Kubernetes Engine Developer (roles/container.developer)
• Cloud Build Editor (roles/cloudbuild.editor)
• Storage Object Admin (roles/storage.objectAdmin)

Training Dataset

Uses a mock pretraining dataset provided by the NeMo framework.

Run the recipe

It's recommended that you use Cloud Shell as your client to complete the steps. Cloud Shell comes pre-installed with the necessary utilities, including kubectl, the Google Cloud SDK, and Helm.

Launch Cloud Shell

In the Google Cloud console, start a Cloud Shell Instance.

Configure environment settings

From your client, complete the following steps:

1. Set the environment variables to match your environment:

export PROJECT_ID=<PROJECT_ID>
export BUILD_REGION=<BUILD_REGION>
export CLUSTER_REGION=<CLUSTER_REGION>
export CLUSTER_NAME=<CLUSTER_NAME>
export GCS_BUCKET=<GCS_BUCKET>
export ARTIFACT_REGISTRY=<ARTIFACT_REGISTRY>
export KUEUE_NAME=<KUEUE_NAME>

Replace the following values:

• <PROJECT_ID>: your Google Cloud project ID
• <BUILD_REGION>: the region where you want to run the Cloud Build
• <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
• <ARTIFACT_REGISTRY>: the full name of your Artifact Registry in the following format: LOCATION-docker.pkg.dev/PROJECT_ID/REPOSITORY
• <KUEUE_NAME>: the name of the Kueue local queue configured for TAS. The default queue created by the cluster toolkit is a3-mega. Make sure to verify the name of the local queue in your cluster.

2. Set the default project:

gcloud config set project $PROJECT_ID

Get the recipe

From your client, 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/a3mega/gpt3_175b/nemo-gke/nemo2507/recipe

Get cluster credentials

From your client, get the credentials for your cluster.

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

Build and push a docker container image to Artifact Registry

To build the container, complete the following steps from your client:

1. Use Cloud Build to build and push the container image.

   cd $REPO_ROOT/src/docker/nemo-24.07
   gcloud builds submit --region=${BUILD_REGION} \
       --config cloudbuild.yml \
       --substitutions _ARTIFACT_REGISTRY=$ARTIFACT_REGISTRY \
       --timeout "2h" \
       --machine-type=e2-highcpu-32 \
       --quiet \
       --async

This command outputs the build ID.

2. You can monitor the build progress by streaming the logs for the build ID. To do this, run the following command. Replace <BUILD_ID> with your build ID.

   BUILD_ID=<BUILD_ID>
   gcloud beta builds log $BUILD_ID --region=$BUILD_REGION

Configure and submit a pretraining job

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:

cd $RECIPE_ROOT
helm install -f values.yaml \
  --set-file workload_launcher=$REPO_ROOT/src/launchers/nemo-10-launcher.sh \
  --set-file workload_config=$REPO_ROOT/src/frameworks/a3mega/nemo-configs/gpt3-175b-256gpus-fp8.yaml \
  --set queue=${KUEUE_NAME} \
  --set "volumes.gcsMounts[0].bucketName=${GCS_BUCKET}" \
  --set workload.image=${ARTIFACT_REGISTRY}/nemo_workload:24.07 \
  $USER-gpt3-175b-nemo \
  $REPO_ROOT/src/helm-charts/a3mega/jobset

Configure job settings

You can overwrite any of the default NeMo configurations for this job. To do this, we can set the new arguments using --set workload.arguments.

Examples

• To set the number of training steps to 100, run the following command from your client:

  cd $RECIPE_ROOT
  helm install -f values.yaml \
    --set-file workload_launcher=$REPO_ROOT/src/launchers/nemo-10-launcher.sh \
    --set-file workload_config=$REPO_ROOT/src/frameworks/a3mega/nemo-configs/gpt3-175b-256gpus-fp8.yaml \
    --set queue=${KUEUE_NAME} \
    --set "volumes.gcsMounts[0].bucketName=${GCS_BUCKET}" \
    --set workload.image=${ARTIFACT_REGISTRY}/nemo_workload:24.07 \
    --set "workload.arguments[0]=trainer.max_steps=100" \
    $USER-gpt3-175b-nemo \
    $REPO_ROOT/src/helm-charts/a3mega/jobset

Monitor the JobSet

To check the status of pods in the JobSet, run the following command from your client:

kubectl get pods -l app.kubernetes.io/instance=$USER-gpt3-175b-nemo

To get the logs of pods in the JobSet, run the following command from your client:

kubectl logs -l app.kubernetes.io/instance=$USER-gpt3-175b-nemo

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

kubectl logs POD_NAME

Analyze results

When completed, the job creates several artifacts, including logs and traces, and places them in the configured Google Cloud Storage bucket as follows:

gs://${GCS_BUCKET}/nemo-experiments/<JOB_ID>
├── hparams.yaml
├── lightning_logs.txt
├── nemo_error_logs.txt
├── nemo_log_globalrank-[RANK]_localrank-[LOCAL].txt
├── dllogger
│   ├── rank-0
│   │   ├── dllogger.json
...

• hparams.yaml: the NeMo configuration used by the pretraining script. This includes the combined configuration file and the command line overrides.
• lightning_logs.txt: the log files generated by PyTorch Lightning, which is used by NeMo
• nemo_error_logs.txt: the warning and error logs generated by NeMo
• nemo_log_globalrank-[RANK]_localrank-[LOCAL].txt: the NeMo logs for each rank
• dllogger/: The log captured by [NVIDIA DLLogger](https://github.com/NVIDIA/dllogger) : DLLogger is configured to store logs on the rank 0 node. The log is in JSON format and includes loss, step_time, and other key metrics for each training step.

Here is an example of an entry in the DLLogger log:

DLLL {
  "timestamp": "1727463426.984522",
  "datetime": "2024-09-27 18:57:06.984522",
  "elapsedtime": "1645.119157",
  "type": "LOG",
  "step": 47,
  "data": {
    "reduced_train_loss": 7.890198707580566,
    "lr": 0.000036782606912311167,
    "global_step": 47,
    "consumed_samples": 98304,
    "train_backward_timing in s": 0.00004844665454584174,
    "grad_norm": 8.800999641418457,
    "train_step_timing in s": 25.115550994873047,
    "epoch": 0
  }
}

The DLLogger log can be used to calculate the Model FLOPS Utilization (MFU) metric, as described in the next section.

Calculate training performance metrics (MFU, TFLOPS, Average Step Time)

This section explains how to calculate key training performance metrics, such as Model FLOPS Utilization (MFU), using the dllogger.json file generated during training.

We provide a tool called training_metrics to help you compute these metrics. This tool can calculate the following metrics:

• MFU: Model FLOPS Utilization
• Average training step time: the average time taken for each training step
• TFLOPS per GPU: the number of Tera Floating Point Operations per second achieved by each GPU

To calculate training performance metrics using the training_metrics tool, complete the following steps command from your client:

1. Download the dllogger.json file. The dllogger.json file is generated during the training session.

   To download the file, run the following command. Replace <JOB_ID> with the ID of your training session.

   gcloud storage cp gs://${GCS_BUCKET}/nemo-experiments/<JOB_ID>/dllogger/rank-0/dllogger.json \
       /path/to/your/local/dllogger.json

2. Run the process_training_results.py script

   cd $REPO_ROOT/src/utils/training_metrics
   python3 process_training_results.py --file /path/to/your/local/dllogger.json \
   --batch_size 2048 \
   --num_accelerators 256 \
   --accelerator_type h100 \
   --precision fp8 \
   --model_type gpt3-175b \
   --start_step=10 \
   --end_step=30

Note: The batch_size, num_accelerators, precision, model_type and accelerator_type are the specific values for this recipe running the default configuration. Average step time is computed by default using the steps 10 to 30.

For more detailed information and advanced usage instructions of this tool, see the full documentation

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-gpt3-175b-nemo