Copybara import of gpu-recipes:

  - 2963f1f33f1822d384dc3a7167f990421bcf60ab Add aotc option.

GitOrigin-RevId: 2963f1f33f1822d384dc3a7167f990421bcf60ab

Author: Copybara

src/helm-charts/a3ultra/nemo-training/templates/nemo-launcher-job.yaml

@@ -18,6 +18,7 @@
 
 {{ $nodes := div .Values.workload.gpus 8 | max 1 }}
 {{ $gpusPerNode := min .Values.workload.gpus 8 }}
+{{ $aotc := default 0 .Values.workload.aotc }}
 
 {{- $root := . -}}
 
@@ -152,6 +153,11 @@ spec:
         - name: SSD_MOUNT_PATH
           value: "{{ $root.Values.volumes.ssdMountPath }}"
 
+        {{- if $aotc }}
+        - name: AOTC
+          value: "{{ $aotc }}"
+        {{- end }}
+
         # The following settings are specific to the Torch distributed launcher:
         {{- range $gcs := $root.Values.volumes.gcsMounts }}
         - name: GCS_FUSE_BUCKET
@@ -253,9 +259,10 @@ spec:
           cat /etc/workload-configuration/nemo-configuration.yaml | sed 's/^/| /'
           echo ""
 
+          touch /workspace/workload_arguments.txt
           echo "Detected the following additional workload arguments:"
           {{- range $root.Values.workload.arguments }}
-          echo "{{ . }}"
+          echo "{{ . }}" | tee -a /workspace/workload_arguments.txt
           {{- end }}
 
 
@@ -275,22 +282,31 @@ spec:
             nvidia-smi dmon -d 20 -s pum &
           fi
 
-          OMP_NUM_THREADS=12 torchrun \
-          --nproc-per-node="$GPUS_PER_NODE" \
-          --nnodes="$NNODES" \
-          --node_rank="$NODE_RANK" \
-          --rdzv_id="$JOB_IDENTIFIER" \
-          --master_addr="$MASTER_ADDR" \
-          --master_port="$MASTER_PORT" \
-          ${TORCH_DISTRIBUTED_TARGET} \
-          --config-path="/etc/workload-configuration" \
-          --config-name="nemo-configuration.yaml" \
-          +trainer.num_nodes="$NNODES" \
-          +exp_manager.version="$JOB_IDENTIFIER" \
-          +exp_manager.dllogger_logger_kwargs.json_file="/gcs/nemo-experiments/$JOB_IDENTIFIER/dllogger/rank-$NODE_RANK/dllogger.json" \
-          {{- range $root.Values.workload.arguments }}
-          {{ . }} \
-          {{- end }}
+          if [[ "{{ $aotc }}" == "0" ]]; then
+            OMP_NUM_THREADS=12 torchrun \
+            --nproc-per-node="$GPUS_PER_NODE" \
+            --nnodes="$NNODES" \
+            --node_rank="$NODE_RANK" \
+            --rdzv_id="$JOB_IDENTIFIER" \
+            --master_addr="$MASTER_ADDR" \
+            --master_port="$MASTER_PORT" \
+            ${TORCH_DISTRIBUTED_TARGET} \
+            --config-path="/etc/workload-configuration" \
+            --config-name="nemo-configuration.yaml" \
+            +trainer.num_nodes="$NNODES" \
+            +exp_manager.version="$JOB_IDENTIFIER" \
+            +exp_manager.dllogger_logger_kwargs.json_file="/gcs/nemo-experiments/$JOB_IDENTIFIER/dllogger/rank-$NODE_RANK/dllogger.json" \
+            {{- range $root.Values.workload.arguments }}
+            {{ . }} \
+            {{- end }}
+
+          else
+            echo "Using AOTC"
+            export ENV_FILE=/workspace/env_vars.txt
+            env > /workspace/env_vars.txt
+            python -m aotc.executor
+
+          fi
 
           echo "Copying log files"
           cp -r /workspace/nemo_experiments/megatron_gpt/$JOB_IDENTIFIER/* /gcs/nemo-experiments/$JOB_IDENTIFIER/

training/a3ultra/llama-3.1-70b/nemo-pretraining-gke/README.md

@@ -108,6 +108,8 @@ gcloud container clusters get-credentials $CLUSTER_NAME --region $CLUSTER_REGION
 
 ### Build and push a docker container image to Artifact Registry
 
+Note:If you'd like to use AotC-based [library](https://github.com/AI-Hypercomputer/aotc) image, you may skip this step.
+
 To build the container, complete the following steps from your client:
 
 1. Use Cloud Build to build and push the container image.
@@ -169,7 +171,21 @@ for this job. To do this, we can set the new arguments using `--set workload.arg
       --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \
       --set workload.arguments="{trainer.max_steps=100}" \
       $USER-llama-3-1-70b-nemo-fp8 \
-      $REPO_ROOT/src/helm-charts/a3mega/nemo-training
+      $REPO_ROOT/src/helm-charts/a3ultra/nemo-training
+  ```
+
+- To use the AotC-based image, run the following command from your client:
+
+  ```bash
+  cd $RECIPE_ROOT
+  export IMAGE=us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo@sha256:7a84264e71f82f225be639dd20fcf9104c80936c0f4f38f94b88dfb60303c70e
+  helm install -f values.yaml \
+      --set-file nemo_config=$REPO_ROOT/src/frameworks/a3ultra/nemo-configs/llama-3.1-70b-256gpus-a3ultra-fp8.yaml \
+      --set workload.image=${IMAGE} \
+      --set volumes.gcsMounts[0].bucketName=${GCS_BUCKET} \
+      --set-string workload.aotc=true \
+      $USER-llama-3-1-70b-nemo-fp8 \
+      $REPO_ROOT/src/helm-charts/a3ultra/nemo-training
   ```
 
 ### Monitor the job

training/a3ultra/llama-3.1-70b/nemo-pretraining-gke/values.yaml

@@ -25,6 +25,7 @@ volumes:
 workload:
   torchDistributedTarget: "/opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py"
   gpus: 256 # This should be one of: {<= 8,  multiple of 8}
+  aotc: false
 
 network:
   ncclSettings: