Skip to content

Commit 6e9c2be

Browse files
hengtaoguohengtaoguo
committed
Enabled layers test against torch for CI only
1 parent 43a0b12 commit 6e9c2be

3 files changed

Lines changed: 42 additions & 10 deletions

File tree

‎.github/workflows/ci_pipeline.yml‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,17 @@ jobs:
199199
is_scheduled_run: ${{ github.event_name == 'schedule' }}
200200
maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
201201

202+
maxtext_cpu_torch_reference_tests:
203+
name: cpu-torch-reference tests
204+
needs: [build_and_upload_maxtext_package]
205+
if: needs.analyze_code_changes.outputs.run_tests == 'true'
206+
uses: ./.github/workflows/run_tests_coordinator.yml
207+
with:
208+
flavor: cpu-torch-reference
209+
base_image: maxtext-unit-test-tpu:py312
210+
is_scheduled_run: ${{ github.event_name == 'schedule' }}
211+
maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
212+
202213
tpu7x-tests:
203214
name: TPU7X tests
204215
needs: [build_and_upload_maxtext_package, gate_test_run]
@@ -284,7 +295,7 @@ jobs:
284295

285296
all_tests_passed:
286297
name: All Required Tests Passed
287-
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]
298+
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]
288299
if: always()
289300
runs-on: ubuntu-latest
290301
steps:
@@ -297,6 +308,7 @@ jobs:
297308
echo "Gate result: ${NEEDS_GATE_TEST_RUN_RESULT}"
298309
echo "TPU Tests (Matrix) result: ${NEEDS_TPU_TESTS_RESULT}"
299310
echo "TPU7X Tests (Matrix) result: ${NEEDS_TPU7X_TESTS_RESULT}"
311+
echo "CPU torch reference tests result: ${NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT}"
300312
echo "GPU Tests (Matrix) result: ${NEEDS_GPU_TESTS_RESULT}"
301313
echo "CPU Tests (Matrix) result: ${NEEDS_CPU_TESTS_RESULT}"
302314
echo "Pathways Unit result: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}"
@@ -317,6 +329,7 @@ jobs:
317329
NEEDS_CPU_TESTS_RESULT: ${{ needs.cpu-tests.result }}
318330
NEEDS_TPU_TESTS_RESULT: ${{ needs.tpu-tests.result }}
319331
NEEDS_TPU7X_TESTS_RESULT: ${{ needs.tpu7x-tests.result }}
332+
NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT: ${{ needs.maxtext_cpu_torch_reference_tests.result }}
320333
NEEDS_GPU_TESTS_RESULT: ${{ needs.gpu-tests.result }}
321334
NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_unit_tests.result }}
322335
NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_integration_tests.result }}
@@ -352,7 +365,7 @@ jobs:
352365

353366
notify_failure:
354367
name: Notify failed build # creates an issue or modifies last open existing issue for failed build
355-
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]
368+
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]
356369
if: ${{ always() }}
357370
runs-on: ubuntu-latest
358371
permissions:
@@ -366,7 +379,7 @@ jobs:
366379

367380
investigate_failure:
368381
name: Investigate failed build # investigates failure of scheduled run and comments on tracking issue
369-
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]
382+
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]
370383
if: ${{ always() && contains(needs.*.result, 'failure') && github.event_name == 'schedule' }}
371384
uses: ./.github/workflows/gemini_investigate.yml
372385
permissions:

‎.github/workflows/run_tests_against_package.yml‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,11 @@ on:
7474
required: false
7575
type: boolean
7676
default: false
77+
install_torch_cpu:
78+
description: 'Install CPU-only torch for layer-wise reference tests'
79+
required: false
80+
type: boolean
81+
default: false
7782
permissions:
7883
contents: read
7984
jobs:
@@ -120,8 +125,14 @@ jobs:
120125
else
121126
install_tpu_pre_train_extra_deps
122127
fi
128+
if [ "${INPUTS_INSTALL_TORCH_CPU}" == "true" ]; then
129+
python3 -m pip install --quiet torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu --no-deps
130+
python3 -c "import torch; print(f'torch {torch.__version__}')"
131+
fi
123132
python3 --version
124133
python3 -m pip freeze
134+
env:
135+
INPUTS_INSTALL_TORCH_CPU: ${{ inputs.install_torch_cpu }}
125136
- name: Copy test assets files
126137
if: ${{ !inputs.maxtext_installed }}
127138
run : gcloud storage cp gs://maxtext-test-assets/* tests/assets

‎.github/workflows/run_tests_coordinator.yml‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ on:
2929
tpu7x-post-training-unit, tpu7x-post-training-integration,
3030
gpu-unit, gpu-integration,
3131
cpu-unit,
32-
cpu-post-training-unit
32+
cpu-post-training-unit,
33+
cpu-torch-reference
3334
)
3435
required: true
3536
type: string
@@ -122,7 +123,8 @@ jobs:
122123
"gpu-unit": "cuda12",
123124
"gpu-integration": "cuda12",
124125
"cpu-unit": "cpu",
125-
"cpu-post-training-unit": "cpu"
126+
"cpu-post-training-unit": "cpu",
127+
"cpu-torch-reference": "cpu"
126128
}')[inputs.flavor] }}
127129
128130
device_name: >-
@@ -138,7 +140,8 @@ jobs:
138140
"gpu-unit": "a100-40gb-4",
139141
"gpu-integration": "a100-40gb-4",
140142
"cpu-unit": "X64",
141-
"cpu-post-training-unit": "X64"
143+
"cpu-post-training-unit": "X64",
144+
"cpu-torch-reference": "X64"
142145
}')[inputs.flavor] }}
143146
144147
cloud_runner: >-
@@ -154,7 +157,8 @@ jobs:
154157
"gpu-unit": "linux-x86-a2-48-a100-4gpu",
155158
"gpu-integration": "linux-x86-a2-48-a100-4gpu",
156159
"cpu-unit": "linux-x86-n2-32",
157-
"cpu-post-training-unit": "linux-x86-n2-32"
160+
"cpu-post-training-unit": "linux-x86-n2-32",
161+
"cpu-torch-reference": "linux-x86-n2-32"
158162
}')[inputs.flavor] }}
159163
# Pytest Marker Mapping
160164
pytest_marker: >-
@@ -170,7 +174,8 @@ jobs:
170174
"gpu-unit": "not cpu_only and not tpu_only and not integration_test and not post_training",
171175
"gpu-integration": "not cpu_only and not tpu_only and integration_test and not post_training",
172176
"cpu-unit": "cpu_only and not post_training",
173-
"cpu-post-training-unit": "cpu_only and post_training"
177+
"cpu-post-training-unit": "cpu_only and post_training",
178+
"cpu-torch-reference": "not post_training"
174179
}')[inputs.flavor] }}
175180
176181
pytest_addopts: >-
@@ -186,7 +191,8 @@ jobs:
186191
"gpu-unit": "",
187192
"gpu-integration": "",
188193
"cpu-unit": "",
189-
"cpu-post-training-unit": "tests/post_training/unit tests/unit"
194+
"cpu-post-training-unit": "tests/post_training/unit tests/unit",
195+
"cpu-torch-reference": ""
190196
}')[inputs.flavor] }}
191197
192198
pytest_extra_args: >-
@@ -202,7 +208,8 @@ jobs:
202208
"gpu-unit": "--ignore=tests/post_training",
203209
"gpu-integration": "--ignore=tests/post_training",
204210
"cpu-unit": "--ignore=tests/post_training",
205-
"cpu-post-training-unit": ""
211+
"cpu-post-training-unit": "",
212+
"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"
206213
}')[inputs.flavor] }}
207214
${{ inputs.additional_pytest_args }}
208215
@@ -223,3 +230,4 @@ jobs:
223230
total_workers: ${{ fromJSON(needs.setup-parameters.outputs.total_workers) }}
224231
maxtext_sha: ${{ inputs.maxtext_sha }}
225232
is_update_hlo: ${{ inputs.is_update_hlo }}
233+
install_torch_cpu: ${{ inputs.flavor == 'cpu-torch-reference' }}

0 commit comments

Comments
 (0)