Add upstream pax 2-node Grok test

.github/container/test-pax.sh

@@ -382,6 +382,10 @@ elif [[ ${MODEL_TYPE} == "5B" ]]; then
 elif [[ ${MODEL_TYPE} == "LLaMA70BProxy" ]]; then
   CONFIG=ci_configs.LLaMA70BSyntheticSmall
   ADDITIONAL_ARGS="--fdl.DCN_MESH_SHAPE=[1,${NODES},1] --fdl.ICI_MESH_SHAPE=[${DP},${FSDP},${TP}] ${ADDITIONAL_ARGS} --fdl.PERCORE_BATCH_SIZE=${BATCH_PER_GPU}"
+## hard-code ICI mesh shape for Grok
+elif [[ ${MODEL_TYPE} == "GrokProxy" ]]; then
+  CONFIG=paxml.tasks.lm.params.nvidia.Grok_Proxy
+  ADDITIONAL_ARGS="--fdl.DCN_MESH_SHAPE=[1,${NODES},1,1] --fdl.ICI_MESH_SHAPE=[1,1,8,1] ${ADDITIONAL_ARGS} --fdl.PERCORE_BATCH_SIZE=${BATCH_PER_GPU} --fdl.NUM_LAYERS=2"
 else
   echo "Unsupported model ${MODEL_TYPE}"
   exit 1

.github/workflows/_test_upstream_pax.yaml

@@ -210,6 +210,11 @@ jobs:
             BATCH_SIZE: 4
             EVALUATE: true
             ADDITIONAL_ARGS: "--model-type LLaMA70BProxy --evaluate"
+          - TEST_NAME: Grok
+            PARALLEL_CONFIG: [1, 2, 8, 1] ## note: only used to compute num nodes
+            BATCH_SIZE: 4
+            EVALUATE: true
+            ADDITIONAL_ARGS: "--model-type GrokProxy"
       fail-fast: false
 
     runs-on: ubuntu-22.04