Fix MaxText metrics/baselines

Author: olupton

.github/workflows/_test_maxtext.yaml

@@ -373,24 +373,6 @@ jobs:
       ARTIFACT_NAME: ${{ inputs.ARTIFACT_NAME }}
       FW_NAME: ${{ inputs.FW_NAME }}
 
-  summary:
-    name: test-maxtext-summary
-    runs-on: ubuntu-22.04
-    needs: [single-process-multi-device, maxtext-multinode]
-    if: "!cancelled()"
-    steps:
-      - name: Generate TensorBoard query URL
-        run: |
-          (
-          cat << EOF
-
-          ## MaxText training
-
-          [view metrics](https://${{ vars.HOSTNAME_TENSORBOARD }}/#scalars&regexInput=${{ inputs.FW_NAME }}-${GITHUB_RUN_ID}&_smoothingWeight=0&tagFilter=seqs_per)
-
-          EOF
-          ) | tee $GITHUB_STEP_SUMMARY
-
   outcome:
     name: test-maxtext-outcome
     needs: sitrep

.github/workflows/baselines/MAXTEXT/upstream/1DP1FSDP1TP1PP.json

@@ -1 +1,33 @@
-{"start_step":1,"end_step":9,"step_interval":1,"loss_values":[0.19917848706245422,0,0,0,0,0,0,0,0],"step_times":[0.27129199107487995,0.17545133332411447,0.1774536669254303,0.18130967020988464,0.17997999986012778,0.17623033126195273,0.17701533436775208,0.17688766618569693,0.1763359953959783],"step_time_avg":0.1879951098450908,"e2e_time_seconds":89.195,"run_urls":["https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/7729826963/artifacts","https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/7729993999/artifacts","https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/7731914601/artifacts"],"date":"2024-01-31"}
+{
+  "start_step": 1,
+  "end_step": 9,
+  "step_interval": 1,
+  "loss_values": [
+    24.950599670410156,
+    23.6308536529541,
+    22.22606086730957,
+    20.85379981994629,
+    19.61219596862793,
+    18.5745849609375,
+    17.77853012084961,
+    17.22124481201172,
+    16.864944458007812
+  ],
+  "step_times": [
+    0.34216299653053284,
+    0.20161199569702148,
+    0.20115399360656738,
+    0.19551700353622437,
+    0.19348999857902527,
+    0.19705399870872498,
+    0.19354699552059174,
+    0.1978529989719391,
+    0.19385899603366852
+  ],
+  "step_time_avg": 0.2129165530204773,
+  "e2e_time_seconds": 34.555,
+  "run_urls": [
+    "https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/14054620516/artifacts"
+  ],
+  "date": "2025-03-25"
+}

.github/workflows/baselines/MAXTEXT/upstream/1DP1FSDP8TP1PP.json

@@ -1 +1,33 @@
-{"start_step":1,"end_step":9,"step_interval":1,"loss_values":[0.20010541379451752,3.576278402306343e-07,0,0,0,0,0,0,0],"step_times":[0.2453316698471705,0.15532933175563812,0.15473033487796783,0.15553300082683563,0.15387233098347983,0.1558946669101715,0.15286600093046823,0.1544460008541743,0.15370899935563406],"step_time_avg":0.1646347040379489,"e2e_time_seconds":184.883,"run_urls":["https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/7729826963/artifacts","https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/7729993999/artifacts","https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/7731914601/artifacts"],"date":"2024-01-31"}
+{
+  "start_step": 1,
+  "end_step": 9,
+  "step_interval": 1,
+  "loss_values": [
+    24.9504337310791,
+    23.63066291809082,
+    22.226167678833008,
+    20.853334426879883,
+    19.611804962158203,
+    18.574153900146484,
+    17.778240203857422,
+    17.22144317626953,
+    16.864635467529297
+  ],
+  "step_times": [
+    0.31094300746917725,
+    0.21168699860572815,
+    0.16145099699497223,
+    0.1525229960680008,
+    0.153779998421669,
+    0.15510499477386475,
+    0.1530109941959381,
+    0.1535159945487976,
+    0.15335600078105927
+  ],
+  "step_time_avg": 0.17837466465102303,
+  "e2e_time_seconds": 32.688,
+  "run_urls": [
+    "https://api.github.com/repos/NVIDIA/JAX-Toolbox/actions/runs/14054620516/artifacts"
+  ],
+  "date": "2025-03-25"
+}

.github/workflows/baselines/MAXTEXT/upstream/1DP2FSDP4TP1PP_single_process.json

@@ -46,7 +46,10 @@ elif [[ "$TYPE" == "rosetta-t5x" ]]; then
     )
     OUTPUT_DIR=T5X_MGMN/rosetta
 elif [[ "$TYPE" == "upstream-maxtext" ]]; then
-    CONFIGS=("1DP1FSDP1TP1PP" "1DP1FSDP8TP1PP" "1DP2FSDP4TP1PP_single_process" "1DP4FSDP2TP1PP" "1DP8FSDP1TP1PP" "2DP2FSDP2TP1PP" "4DP2FSDP2TP1PP")
+    CONFIGS=(
+	"1DP2FSDP4TP1PP_single_process"
+	"2DP2FSDP2TP1PP"
+    )
     OUTPUT_DIR=MAXTEXT/upstream
 else
     usage
@@ -60,8 +63,10 @@ bash ${UTIL_DIR}/download_artifacts.sh ${ALL_WF_RUNS[@]}
 URLS=()
 for WORKFLOW_RUN in ${ALL_WF_RUNS[@]}; do
   for CFG in ${CONFIGS[@]}; do
-    if [[ $(find . -mindepth 1 -maxdepth 2 -type d -name $CFG | wc -l) -ne 1 ]]; then
-      echo "Expected one artifact to have a '$CFG' dir under '$PWD', but found $(find . -mindepth 1 -maxdepth 2 -type d -name $CFG)"
+    CFG=$TYPE-$WORKFLOW_RUN-$CFG
+    ARTS=$(find . -mindepth 1 -maxdepth 2 -type d -name $CFG)
+    if (( $(echo ${ARTS} | wc -l) != 1 )); then
+      echo "Expected one artifact to have a '$CFG' dir under '$PWD', but found ${ARTS}"
       exit 1
     fi
   done

.github/workflows/baselines/MAXTEXT/upstream/1DP4FSDP2TP1PP.json

@@ -40,7 +40,7 @@ def main():
         if not os.path.exists(searchpath):
             searchpath = os.path.join(args.test_config, "summaries/train")
         if not os.path.exists(searchpath):
-            searchpath = os.path.join(args.test_config, "logdir/tensorboard")
+            searchpath = os.path.join(args.test_config, "logdir/tensorboard/logdir")
         assert os.path.exists(searchpath), f"Neither {args.test_config}/train nor {args.test_config}/summaries/train nor {args.test_config}/logdir/tensorboard dirs exist"
         event_files = glob.glob(os.path.join(searchpath, "events*"))
         assert len(event_files) > 0, f"{searchpath} did not contain a tensorboard events file"