Update UnitTests.yml

Author: parambole

.github/workflows/UnitTests.yml

@@ -19,6 +19,9 @@
 # This workflow will run a small MaxText training workload on a GPU runner
 # using a custom Docker image with all dependencies pre-installed.
 
+# This workflow will run a small MaxText training workload on a GPU runner
+# using a custom Docker image with all code and dependencies pre-installed.
+
 name: MaxText Custom Image Workload
 
 on:
@@ -32,21 +35,23 @@ jobs:
     name: "Run MaxText Training Workload"
     runs-on: ["linux-x86-a3-megagpu-h100-8gpu"]
     container:
-      # Use your newly built custom image
+      # Use your custom image which contains the source code and dependencies.
       image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/maxtext-gpu-custom:latest
 
     steps:
       - name: Run MaxText Training
+        # The Docker image's working directory is /deps, but the code is in /deps/src.
+        working-directory: /deps/src
         env:
           NVTE_FRAMEWORK: jax
           TF_FORCE_GPU_ALLOW_GROWTH: "true"
         run: |
-          # The working directory is /deps, so this path is correct.
+          # Run the main training script from the /deps/src directory.
           python MaxText/train.py MaxText/configs/base.yml \
             run_name="maxtext-ci-test-${{ github.run_id }}" \
             steps=5 \
             enable_checkpointing=false \
-            attention='cudnet_flash_te' \
+            attention='cudnn_flash_te' \
             dataset_type='synthetic'
 
 # name: SDXL Workload Training on GPU