Speed up cpu-unit CI

Author: gagika

.github/workflows/run_pathways_tests.yml

@@ -102,7 +102,7 @@ jobs:
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
           export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
           # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
-          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0
+          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not AotJaxprIdenticalTest and not CompileThenLoad" --durations=0
         env:
           PYTHONPATH: "${{ github.workspace }}/src"
     services:

.github/workflows/run_tests_against_package.yml

@@ -19,6 +19,11 @@ name: Run Tests Against MaxText Package
 on:
   workflow_call:
     inputs:
+      flavor:
+        description: 'Test flavor name (e.g. cpu-unit, tpu-unit) - used for artifact naming'
+        required: false
+        type: string
+        default: ''
       device_type:
         required: true
         type: string
@@ -164,6 +169,11 @@ jobs:
           if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then
             $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist
             SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP} -n auto"
+            # On scheduled runs, record per-shard durations so future splits balance by time (LPT).
+            # Merge artifacts offline and commit as .test_durations at repo root.
+            if [ "${INPUTS_IS_SCHEDULED_RUN}" == "true" ]; then
+              SPLIT_ARGS="${SPLIT_ARGS} --store-durations --durations-path=.test_durations.${INPUTS_WORKER_GROUP}"
+            fi
           else
             SPLIT_ARGS=""
           fi
@@ -195,3 +205,11 @@ jobs:
           # If scheduled, upload to scheduled flag only. If PR, upload to regular flag only.
           flags: ${{ inputs.is_scheduled_run == 'true' && 'scheduled' || 'regular' }}
           verbose: true
+      - name: Upload test durations artifact
+        if: ${{ inputs.is_scheduled_run == 'true' && inputs.total_workers > 1 }}
+        uses: actions/upload-artifact@v4
+        continue-on-error: true
+        with:
+          name: test-durations-${{ inputs.flavor }}-${{ inputs.worker_group }}
+          path: .test_durations.*
+          if-no-files-found: ignore

.github/workflows/run_tests_coordinator.yml

@@ -57,10 +57,11 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        worker_group: ${{ fromJSON(contains(inputs.flavor, 'cpu-unit') && '[1, 2]' || '[1]') }}
+        worker_group: ${{ fromJSON(contains(inputs.flavor, 'cpu-unit') && '[1, 2, 3, 4]' || '[1]') }}
 
     uses: ./.github/workflows/run_tests_against_package.yml
     with:
+      flavor: ${{ inputs.flavor }}
       # Infrastructure Mapping
       device_type: >-
         ${{ fromJSON('{
@@ -148,5 +149,5 @@ jobs:
       is_scheduled_run: ${{ inputs.is_scheduled_run }}
       maxtext_installed: ${{ inputs.maxtext_installed }}
       worker_group: ${{ matrix.worker_group }}
-      total_workers: ${{ contains(inputs.flavor, 'cpu-unit') && 2 || 1 }}
+      total_workers: ${{ contains(inputs.flavor, 'cpu-unit') && 4 || 1 }}
       maxtext_sha: ${{ inputs.maxtext_sha }}

tests/unit/train_compile_test.py

@@ -682,9 +682,8 @@ def test_moe_llama4_17b_16e(self):
         )
     )
 
-  @pytest.mark.cpu_only
-  def test_moe_gpt_oss_20b_sparse_matmul(self):
-    compiled_trainstep_file = "/tmp/test_moe_gpt_oss_20b_sparse_matmul.pickle"
+  def _run_moe_gpt_oss_20b(self, suffix, matmul_args):
+    compiled_trainstep_file = f"/tmp/test_moe_gpt_oss_20b_{suffix}.pickle"
     train_compile_main(
         (
             "",
@@ -698,33 +697,18 @@ def test_moe_gpt_oss_20b_sparse_matmul(self):
             "dtype=bfloat16",
             "weight_dtype=bfloat16",
             "scan_layers=True",
-            "sparse_matmul=True",
-            "megablox=True",
             "attention=flash",
+            *matmul_args,
         )
     )
 
+  @pytest.mark.cpu_only
+  def test_moe_gpt_oss_20b_sparse_matmul(self):
+    self._run_moe_gpt_oss_20b("sparse_matmul", ["sparse_matmul=True", "megablox=True"])
+
   @pytest.mark.cpu_only
   def test_moe_gpt_oss_20b_dense_matmul(self):
-    compiled_trainstep_file = "/tmp/test_moe_gpt_oss_20b_dense_matmul.pickle"
-    train_compile_main(
-        (
-            "",
-            get_test_config_path(),
-            f"compiled_trainstep_file={compiled_trainstep_file}",
-            "compile_topology=v5p-16",
-            "compile_topology_num_slices=1",
-            "model_name=gpt-oss-20b",
-            "per_device_batch_size=1",
-            "max_target_length=1024",
-            "dtype=bfloat16",
-            "weight_dtype=bfloat16",
-            "scan_layers=True",
-            "sparse_matmul=False",
-            "capacity_factor=-1",
-            "attention=flash",
-        )
-    )
+    self._run_moe_gpt_oss_20b("dense_matmul", ["sparse_matmul=False", "capacity_factor=-1"])
 
   @pytest.mark.cpu_only
   def test_gpt3_6b(self):