Train Llama2 with Megatron-LM on A3 Mega virtual machines

Set up your environment

1. Create environment variables for some common parameters

   export CLUSTER_NAME=CLUSTER_NAME
   export REGION=REGION
   export ZONE=ZONE
   export PROJECT_ID=PROJECT_ID
   

   Replace the following:

   • CLUSTER_NAME: the name of your A3 Mega GKE cluster that has GPUDirect-TCPXO and multi-networking enabled.
   • REGION: the region where you created your cluster.
   • ZONE: the zone where you created your cluster.
   • PROJECT_ID: your Google Cloud project ID.

2. Configure the Google Cloud CLI to use your Google Cloud credentials for authentication:

   gcloud auth login
   

   For more information, see Authenticate for using the Google Cloud CLI.

3. Install kubectl and the GKE gcloud CLI plugin:

   sudo apt-get install kubectl
   sudo apt-get install google-cloud-sdk-gke-gcloud-auth-plugin
   

4. Fetch credentials for your GKE cluster:

   gcloud container clusters get-credentials ${CLUSTER_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID}
   

5. If not already installed, install Helm:

   curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3
   chmod 700 get_helm.sh
   ./get_helm.sh && rm get_helm.sh
   sudo chmod +x /usr/local/bin/helm
   

Use topology-aware scheduler to deploy your Pods

You can use the topology-aware scheduler to deploy your GKE Pods to nodes that have a specified GPU topology.

In the following kubectl commands, you will use the files directly from a repository. Alternatively, you can clone the repository locally and the kubectl commands can reference the local files instead.

For more information, see Topology scheduler.

1. Set up the service account:

   kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/service-account.yaml
   

2. Install the topology scheduler scripts in a configmap:

   curl -OL  https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/schedule-daemon.py
   curl -OL  https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/label-nodes-daemon.py
   
   kubectl -n kube-system create configmap topology-scheduler-scripts \
       --from-file=schedule-daemon.py=schedule-daemon.py \
       --from-file=label-nodes-daemon.py=label-nodes-daemon.py
   

3. Install the topology label daemonset and topology scheduler Pod:

   kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/label-nodes-daemon.yaml
   $ kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/topology-scheduler/schedule-daemon.yaml
   

4. Observe the actions of the topology scheduler:

   kubectl -n kube-system logs topology-scheduler-pod
   

Run the workload

Build the Dockerfile and push to the Google Cloud Artifact Registry

1. Create a Cloud Storage bucket and a Docker repository. In the scripts/setup-and-configure-resources.sh script, replace the bucket and repository names with the ones you created, and then run the script:

   bash scripts/setup-and-configure-resources.sh
   

2. Build and push the pytorch-megatron:23.11-py3 image to your repository. Ensure the Docker repository name in the scripts/build-and-push-docker-image.sh file matches the repository name you used in the scripts/setup-and-configure-resources.sh script. You can also edit the Docker image tag name before pushing.

   bash scripts/build-and-push-docker-image.sh
   

Launch Megatron-LM Llama2 benchmark

1. Edit the helm/values.yaml file to specify your Cloud Storage bucket and Docker image created in previous sections. For some example configurations, see sample-configurations.

2. Optional: You can also edit the selected-configuration.sh file to specify any changes you made to the default Helm configuration.

   helm install HELM_EXPERIMENT_NAME helm/ --values helm/values.yaml
   

   Replace HELM_EXPERIMENT_NAME with an arbitrary name for your experiment.

The experiment writes metrics from the Nsight Systems profiling tool to the Cloud Storage bucket specified in the megatron-experiments directory.