Merge pull request #4371 from AI-Hypercomputer:hengtaoguo-test

PiperOrigin-RevId: 944691011

Author: Google-ML-Automation

.github/workflows/ci_pipeline.yml

@@ -199,6 +199,17 @@ jobs:
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
+  maxtext_cpu_torch_reference_tests:
+    name: cpu-torch-reference tests
+    needs: [build_and_upload_maxtext_package]
+    if: needs.analyze_code_changes.outputs.run_tests == 'true'
+    uses: ./.github/workflows/run_tests_coordinator.yml
+    with:
+      flavor: cpu-torch-reference
+      base_image: maxtext-unit-test-tpu:py312
+      is_scheduled_run: ${{ github.event_name == 'schedule' }}
+      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+
   tpu7x-tests:
     name: TPU7X tests
     needs: [build_and_upload_maxtext_package, gate_test_run]
@@ -284,7 +295,7 @@ jobs:
 
   all_tests_passed:
     name: All Required Tests Passed
-    needs: [build_and_upload_maxtext_package, gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
+    needs: [build_and_upload_maxtext_package, gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
     if: always()
     runs-on: ubuntu-latest
     steps:
@@ -297,6 +308,7 @@ jobs:
           echo "Gate result: ${NEEDS_GATE_TEST_RUN_RESULT}"
           echo "TPU Tests (Matrix) result: ${NEEDS_TPU_TESTS_RESULT}"
           echo "TPU7X Tests (Matrix) result: ${NEEDS_TPU7X_TESTS_RESULT}"
+          echo "CPU torch reference tests result: ${NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT}"
           echo "GPU Tests (Matrix) result: ${NEEDS_GPU_TESTS_RESULT}"
           echo "CPU Tests (Matrix) result: ${NEEDS_CPU_TESTS_RESULT}"
           echo "Pathways Unit result: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}"
@@ -317,6 +329,7 @@ jobs:
           NEEDS_CPU_TESTS_RESULT: ${{ needs.cpu-tests.result }}
           NEEDS_TPU_TESTS_RESULT: ${{ needs.tpu-tests.result }}
           NEEDS_TPU7X_TESTS_RESULT: ${{ needs.tpu7x-tests.result }}
+          NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT: ${{ needs.maxtext_cpu_torch_reference_tests.result }}
           NEEDS_GPU_TESTS_RESULT: ${{ needs.gpu-tests.result }}
           NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_unit_tests.result }}
           NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_integration_tests.result }}
@@ -352,7 +365,7 @@ jobs:
 
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
+    needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:
@@ -366,7 +379,7 @@ jobs:
 
   investigate_failure:
     name: Investigate failed build # investigates failure of scheduled run and comments on tracking issue
-    needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check, notify_failure]
+    needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check, notify_failure]
     if: ${{ always() && contains(needs.*.result, 'failure') && github.event_name == 'schedule' }}
     uses: ./.github/workflows/gemini_investigate.yml
     permissions:

.github/workflows/run_tests_against_package.yml

@@ -74,6 +74,11 @@ on:
         required: false
         type: boolean
         default: false
+      install_torch_cpu:
+        description: 'Install CPU-only torch for layer-wise reference tests'
+        required: false
+        type: boolean
+        default: false
 permissions:
   contents: read
 jobs:
@@ -120,8 +125,14 @@ jobs:
           else
             install_tpu_pre_train_extra_deps
           fi
+          if [ "${INPUTS_INSTALL_TORCH_CPU}" == "true" ]; then
+            python3 -m pip install --quiet torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu --no-deps
+            python3 -c "import torch; print(f'torch {torch.__version__}')"
+          fi
           python3 --version
           python3 -m pip freeze
+        env:
+          INPUTS_INSTALL_TORCH_CPU: ${{ inputs.install_torch_cpu }}
       - name: Copy test assets files
         if: ${{ !inputs.maxtext_installed }}
         run : gcloud storage cp gs://maxtext-test-assets/* tests/assets

.github/workflows/run_tests_coordinator.yml

@@ -29,7 +29,8 @@ on:
             tpu7x-post-training-unit, tpu7x-post-training-integration,
             gpu-unit, gpu-integration,
             cpu-unit,
-            cpu-post-training-unit
+            cpu-post-training-unit,
+            cpu-torch-reference
           )
         required: true
         type: string
@@ -122,7 +123,8 @@ jobs:
           "gpu-unit": "cuda12",
           "gpu-integration": "cuda12",
           "cpu-unit": "cpu",
-          "cpu-post-training-unit": "cpu"
+          "cpu-post-training-unit": "cpu",
+          "cpu-torch-reference": "cpu"
         }')[inputs.flavor] }}
 
       device_name: >-
@@ -138,7 +140,8 @@ jobs:
           "gpu-unit": "a100-40gb-4",
           "gpu-integration": "a100-40gb-4",
           "cpu-unit": "X64",
-          "cpu-post-training-unit": "X64"
+          "cpu-post-training-unit": "X64",
+          "cpu-torch-reference": "X64"
         }')[inputs.flavor] }}
 
       cloud_runner: >-
@@ -154,7 +157,8 @@ jobs:
           "gpu-unit": "linux-x86-a2-48-a100-4gpu",
           "gpu-integration": "linux-x86-a2-48-a100-4gpu",
           "cpu-unit": "linux-x86-n2-32",
-          "cpu-post-training-unit": "linux-x86-n2-32"
+          "cpu-post-training-unit": "linux-x86-n2-32",
+          "cpu-torch-reference": "linux-x86-n2-32"
         }')[inputs.flavor] }}
       # Pytest Marker Mapping
       pytest_marker: >-
@@ -170,7 +174,8 @@ jobs:
             "gpu-unit": "not cpu_only and not tpu_only and not integration_test and not post_training",
             "gpu-integration": "not cpu_only and not tpu_only and integration_test and not post_training",
             "cpu-unit": "cpu_only and not post_training",
-            "cpu-post-training-unit": "cpu_only and post_training"
+            "cpu-post-training-unit": "cpu_only and post_training",
+            "cpu-torch-reference": "not post_training"
           }')[inputs.flavor] }}
 
       pytest_addopts: >-
@@ -186,7 +191,8 @@ jobs:
             "gpu-unit": "",
             "gpu-integration": "",
             "cpu-unit": "",
-            "cpu-post-training-unit": "tests/post_training/unit tests/unit"
+            "cpu-post-training-unit": "tests/post_training/unit tests/unit",
+            "cpu-torch-reference": ""
           }')[inputs.flavor] }}
 
       pytest_extra_args: >-
@@ -202,7 +208,8 @@ jobs:
             "gpu-unit": "--ignore=tests/post_training",
             "gpu-integration": "--ignore=tests/post_training",
             "cpu-unit": "--ignore=tests/post_training",
-            "cpu-post-training-unit": ""
+            "cpu-post-training-unit": "",
+            "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py"
           }')[inputs.flavor] }}
         ${{ inputs.additional_pytest_args }}
 
@@ -223,3 +230,4 @@ jobs:
       total_workers: ${{ fromJSON(needs.setup-parameters.outputs.total_workers) }}
       maxtext_sha: ${{ inputs.maxtext_sha }}
       is_update_hlo: ${{ inputs.is_update_hlo }}
+      install_torch_cpu: ${{ inputs.flavor == 'cpu-torch-reference' }}