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 CONTROL_PLANE_LOCATION= CONTROL_PLANE_LOCATION 
   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.
   • CONTROL_PLANE_LOCATION : the Compute Engine location of the control plane of your cluster. Provide a region for regional clusters, or a zone for zonal clusters.
   • 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} \
     --location=${CONTROL_PLANE_LOCATION} \
     --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.