Merge remote-tracking branch 'origin/main' into 24.10-devel-add-ngs-release-testing

Author: DwarKapex

.github/container/Dockerfile.base

@@ -1,8 +1,8 @@
 # syntax=docker/dockerfile:1-labs
-ARG BASE_IMAGE=nvidia/cuda:12.5.0-devel-ubuntu22.04
+ARG BASE_IMAGE=nvidia/cuda:12.6.1-devel-ubuntu22.04
 ARG GIT_USER_NAME="JAX Toolbox"
 ARG [email protected]
-ARG CLANG_VERSION=17
+ARG CLANG_VERSION=18
 
 ###############################################################################
 ## Obtain GCP's NCCL TCPx plugin

.github/container/test-pax.sh

@@ -15,7 +15,8 @@ usage() {
     echo "  -a, --additional-args      Additional fiddle args to pass to paxml/main.py"
     echo "  -b, --batch-per-gpu        Batch size per GPU, defaults to 4."
     echo "  --dtype                    Batch size, defaults to bfloat16."
-    echo "  --enable-te                If set, will run with env var ENABLE_TE=1." 
+    echo "  --enable-te                If set, will run with env var ENABLE_TE=1."
+    echo "  --enable-cudnn-fa          If set, will use cudnn fa."
     echo "  --enable-dropout           If set, will set DROPOUT_PROB to 0.1."
     echo "  --disable-fused-attn       Whether disable TE fused attention."
     echo "  --model-type               One of 126M, 5B, LLaMA70BProxy. Defaults to 126M"
@@ -26,13 +27,13 @@ usage() {
     echo "  --data-parallel            Data parallelism to use. Defaults to 1."
     echo "  --fsdp                     Fully-sharded data parallelism to use. Defaults to 1."
     echo "  --tensor-parallel          Tensor parallelism to use. Defaults to 1."
-    echo "  --pipeline-parallel        Pipeline parallelism to use. Defaults to 1 for no pipelining." 
+    echo "  --pipeline-parallel        Pipeline parallelism to use. Defaults to 1 for no pipelining."
     echo "  -n, --nodes                Number of nodes."
     echo "  -h, --help                 Print usage."
     exit $1
 }
 
-args=$(getopt -o a:b:s:o:n:h --long additional-args:,batch-per-gpu:,dtype:,enable-te,enable-dropout,disable-fused-attn,model-type:,evaluate,steps:,help,multiprocess,output:,data-parallel:,fsdp:,tensor-parallel:,pipeline-parallel:,nodes: -- "$@")
+args=$(getopt -o a:b:s:o:n:h --long additional-args:,batch-per-gpu:,dtype:,enable-te,enable-cudnn-fa,enable-dropout,disable-fused-attn,model-type:,evaluate,steps:,help,multiprocess,output:,data-parallel:,fsdp:,tensor-parallel:,pipeline-parallel:,nodes: -- "$@")
 if [[ $? -ne 0 ]]; then
     exit $1
 fi
@@ -50,6 +51,7 @@ TP=1
 PP=1
 NODES=1
 ENABLE_TE=0
+ENABLE_CUDNN_FA=0
 MODEL_TYPE=126M
 NVTE_FUSED_ATTN=1
 DROPOUT=0
@@ -75,6 +77,10 @@ while [ : ]; do
             ENABLE_TE=1
             shift 1
             ;;
+        --enable-cudnn-fa)
+            ENABLE_CUDNN_FA=1
+            shift 1
+            ;;
         --enable-dropout)
             DROPOUT='0.1'
             shift 1
@@ -128,7 +134,7 @@ while [ : ]; do
             ;;
         --)
             shift;
-            break 
+            break
             ;;
         *)
             echo "UNKNOWN OPTION $1"
@@ -149,6 +155,7 @@ print_var NGPUS
 print_var OUTPUT
 print_var MULTIPROCESS
 print_var ENABLE_TE
+print_var ENABLE_CUDNN_FA
 print_var NVTE_FUSED_ATTN
 print_var EVALUATE
 print_var DROPOUT
@@ -196,10 +203,10 @@ if dcn_factor > 1:
   if dp % dcn_factor == 0:
     dcn_dp = dcn_factor
     dp = int(dp / dcn_factor)
-  elif fsdp % dcn_factor == 0: 
+  elif fsdp % dcn_factor == 0:
     dcn_fsdp = dcn_factor
     fsdp = int(fsdp / dcn_factor)
-  elif pp % dcn_factor == 0: 
+  elif pp % dcn_factor == 0:
     dcn_pp = dcn_factor
     pp = int(pp / dcn_factor)
 
@@ -209,12 +216,12 @@ class GPT126MPP(TransformerLmSpmdPipelineAdam):
   USE_REPEATED_LAYER = False
   ICI_MESH_SHAPE = [64,1,1]
   MAX_STEPS = 600000
-  
+
   MAX_SEQ_LEN = 2048
   VOCAB_SIZE = 50304
   PACKED_INPUT = True
   PERCORE_BATCH_SIZE = 4
-  
+
   NUM_LAYERS = 12
   NUM_HEADS = 12
   MODEL_DIMS = 768
@@ -223,14 +230,14 @@ class GPT126MPP(TransformerLmSpmdPipelineAdam):
 
   TRAINABLE_POSITION_EMB = True
   TRAINABLE_PE_MAX_SEQ_LEN = MAX_SEQ_LEN
-  
+
   USE_BIAS = True
   LAYERNORM_EPSILON = 1e-5
   ATTEN_LOGIT_CAP = -1.0
   INIT_STD = 0.023
   SOFTMAX_INIT_STD = 0.023
   ACTIVATION_CLS = layers.GELU
-    
+
   ## optimizer-related
   ADAM_BETA1 = 0.9
   ADAM_BETA2 = 0.95
@@ -255,15 +262,15 @@ class GPT126MPP(TransformerLmSpmdPipelineAdam):
   ## disable eval to avoid including eval
   ## in steps/sec calculation
   EVAL_INTERVAL_STEPS = 100000
-  
+
   def task(self):
     task_p = super().task()
     task_p = configure_gpt3_task(self, task_p)
 
     task_p.train.num_train_steps = self.MAX_STEPS
 
     model_p = task_p.model
-    
+
     ### compute layernorm reductions in fp32. Needed for stable training on GPUs
     stacked_p = model_p.lm_tpl.stacked_transformer_tpl
     if stacked_p.cls == layers.PipelinedTransformer:
@@ -274,13 +281,13 @@ class GPT126MPP(TransformerLmSpmdPipelineAdam):
     transformer_layer_p.ln_tpl.reductions_in_fp32 = True
     transformer_layer_p.tr_fflayer_tpl.ln_tpl.reductions_in_fp32 = True
     task_p.model.lm_tpl.final_ln_tpl.reductions_in_fp32 = True
-    
+
     model_p.params_init = WeightInit.Gaussian(self.INIT_STD)
     softmax_init = WeightInit.Gaussian(self.SOFTMAX_INIT_STD)
     model_p.lm_tpl.softmax_tpl.params_init = softmax_init
-    
+
     model_p.apply_eval_sample_weights = True
-    
+
     ## set input, residual, attention dropout to DROPOUT_PROB, remaining dropout to 0
     stacked_p.dropout_prob = 0.0
     stacked_p.input_dropout_prob = self.DROPOUT_PROB
@@ -316,14 +323,14 @@ class LLaMA70BSyntheticSmall(BaseLLaMA, SyntheticDataset):
 if pp > 1:
   @experiment_registry.register
   class Synthetic126MCI(GPT126MPP, SyntheticDataset):
-    
+
     ICI_MESH_SHAPE = [pp, dp, fsdp, tp]
     DCN_MESH_SHAPE = [dcn_pp, dcn_dp, dcn_fsdp, 1]
     MICROBATCH_SIZE = 2
     NUM_STAGES = pp
     PERCORE_BATCH_SIZE = percore_batch_size
     FRPOP_DTYPE = dtype
-    
+
     def task(self):
       task_p = super().task()
       task_p.train.always_use_train_for_model_init=False
@@ -333,7 +340,7 @@ if pp > 1:
 else:
   @experiment_registry.register
   class Synthetic126MCI(Synthetic126M):
-    
+
     ICI_MESH_SHAPE = [dp, fsdp, tp]
     DCN_MESH_SHAPE = [dcn_dp, dcn_fsdp, 1]
     PERCORE_BATCH_SIZE = percore_batch_size
@@ -343,7 +350,7 @@ else:
 
     ## disable eval
     EVAL_INTERVAL_STEPS = 100000
-    
+
     def task(self):
       task_p = super().task()
 
@@ -374,6 +381,10 @@ export ENABLE_TE=$ENABLE_TE
 export NVTE_FUSED_ATTN=$NVTE_FUSED_ATTN
 export VOCAB_PATH=${VOCAB_PATH:-gs://t5-data/vocabs/cc_all.32000.100extra/sentencepiece.model}
 
+if [[ ${ENABLE_CUDNN_FA} -ne 0 ]]; then
+  ADDITIONAL_ARGS="${ADDITIONAL_ARGS} --fdl.USE_CUDNN_FLASH_ATTENTION=True"
+fi
+
 if [[ ${MODEL_TYPE} == "126M" ]]; then
   CONFIG=ci_configs.Synthetic126MCI
 elif [[ ${MODEL_TYPE} == "5B" ]]; then

.github/workflows/_ci.yaml

@@ -11,6 +11,11 @@ on:
         description: 'Build date in YYYY-MM-DD format'
         required: false
         default: NOT SPECIFIED
+      CUDA_IMAGE:
+        type: string
+        description: CUDA image to use as base, e.g. nvidia/cuda:X.Y.Z-devel-ubuntu22.04
+        default: 'latest'
+        required: false
       MANIFEST_ARTIFACT_NAME:
         type: string
         description: 'Artifact name in current run w/ manifest/patches. Leaving empty uses manifest/patches in current branch'
@@ -37,6 +42,7 @@ jobs:
     uses: ./.github/workflows/_build_base.yaml
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
+      BASE_IMAGE: ${{ inputs.CUDA_IMAGE }}
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       MANIFEST_ARTIFACT_NAME: ${{ inputs.MANIFEST_ARTIFACT_NAME }}
     secrets: inherit

.github/workflows/ci.yaml

@@ -28,6 +28,11 @@ on:
         description: "(used if BUMP_MANIFEST=true) If true: attempt to PR/merge manifest branch"
         default: false
         required: false
+      CUDA_IMAGE:
+        type: string
+        description: CUDA image to use as base, e.g. nvidia/cuda:X.Y.Z-devel-ubuntu22.04
+        default: 'latest'
+        required: false
       SOURCE_OVERRIDES:
         type: string
         description: |
@@ -60,6 +65,7 @@ jobs:
       MANIFEST_ARTIFACT_NAME: ${{ steps.manifest-branch.outputs.MANIFEST_ARTIFACT_NAME }}
       MANIFEST_BRANCH: ${{ steps.manifest-branch.outputs.MANIFEST_BRANCH }}
       MERGE_BUMPED_MANIFEST: ${{ steps.manifest-branch.outputs.MERGE_BUMBED_MANIFEST }}
+      CUDA_IMAGE: ${{ steps.cuda-image.outputs.CUDA_IMAGE }}
     steps:
       - name: Cancel workflow run if the trigger is a draft PR
         id: cancel-if-draft
@@ -114,6 +120,17 @@ jobs:
             exit 1
           fi
       
+      - name: Determine CUDA image to use
+        id: cuda-image
+        shell: bash -x -e {0}
+        run: |
+          if [[ "${{ github.event_name }}" == "workflow_dispatch" ]]; then
+            CUDA_IMAGE="${{ inputs.CUDA_IMAGE }}"
+          else
+            CUDA_IMAGE="latest"
+          fi
+          echo "CUDA_IMAGE=${CUDA_IMAGE}" >> $GITHUB_OUTPUT
+
   bump-manifest:
     needs: metadata
     runs-on: ubuntu-22.04
@@ -177,6 +194,7 @@ jobs:
     with:
       ARCHITECTURE: amd64
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
+      CUDA_IMAGE: ${{ needs.metadata.outputs.CUDA_IMAGE }}
       MANIFEST_ARTIFACT_NAME: ${{ needs.metadata.outputs.MANIFEST_ARTIFACT_NAME }}
       SOURCE_URLREFS: ${{ needs.bump-manifest.outputs.SOURCE_URLREFS }}
     secrets: inherit
@@ -187,6 +205,7 @@ jobs:
     with:
       ARCHITECTURE: arm64
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
+      CUDA_IMAGE: ${{ needs.metadata.outputs.CUDA_IMAGE }}
       MANIFEST_ARTIFACT_NAME: ${{ needs.metadata.outputs.MANIFEST_ARTIFACT_NAME }}
       SOURCE_URLREFS: ${{ needs.bump-manifest.outputs.SOURCE_URLREFS }}
     secrets: inherit