From cca23e897e79a128c52331b936ab977769c0ce50 Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Tue, 30 Sep 2025 20:26:55 +0400 Subject: [PATCH 01/13] GRPO integration test Signed-off-by: Vladimir Suvorov --- .../grpo_llama3_demo_test.py | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 tests/integration_tests/grpo_llama3_demo_test.py diff --git a/tests/integration_tests/grpo_llama3_demo_test.py b/tests/integration_tests/grpo_llama3_demo_test.py new file mode 100644 index 0000000000..7aaa532baa --- /dev/null +++ b/tests/integration_tests/grpo_llama3_demo_test.py @@ -0,0 +1,103 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Integration test for GRPO Llama3 demo. + +This test runs the grpo_llama3_demo.py script with minimal configuration +to verify the GRPO training pipeline works correctly on TPU. + +ATTENTION: This test should only be run on TPU (v4-8, v5p-8, or v6e-8). + +Usage: + pytest tests/integration_tests/grpo_llama3_demo_test.py +""" + +import os +import subprocess +import sys +import tempfile +import pytest + +from MaxText.globals import MAXTEXT_REPO_ROOT + + +@pytest.mark.integration_test +@pytest.mark.tpu_only +@pytest.mark.scheduled_only +def test_grpo_llama3_demo(): + """ + Test the GRPO Llama3 demo script with minimal configuration. + + This test verifies that the demo can: + 1. Load the dataset + 2. Initialize the model + 3. Run a few training steps + 4. Complete evaluation + """ + + # Create a temporary directory for outputs + with tempfile.TemporaryDirectory() as temp_dir: + # Set environment variables for the test + env = os.environ.copy() + env["HOME"] = temp_dir + env["SKIP_JAX_PRECOMPILE"] = "1" + + # Build the command to run the demo + demo_script = os.path.join( + MAXTEXT_REPO_ROOT, + "src/MaxText/examples/grpo_llama3_demo.py" + ) + + # We'll modify the script to run with minimal steps for testing + # by setting environment variables that the script can read + env["GRPO_TEST_MODE"] = "1" + env["GRPO_NUM_BATCHES"] = "2" # Minimal batches for testing + env["GRPO_NUM_TEST_BATCHES"] = "2" # Minimal test batches + env["GRPO_MAX_STEPS"] = "2" # Minimal training steps + + # Run the demo script + result = subprocess.run( + [sys.executable, demo_script], + env=env, + capture_output=True, + text=True, + timeout=600, # 10 minute timeout + ) + + # Check if the script executed successfully + if result.returncode != 0: + print("STDOUT:", result.stdout) + print("STDERR:", result.stderr) + raise AssertionError( + f"GRPO Llama3 demo failed with return code {result.returncode}" + ) + + # Verify expected outputs in the logs + assert "HBM usage before loading model:" in result.stdout, \ + "Model initialization not found in output" + assert "HBM usage after loading ref model:" in result.stdout, \ + "Reference model loading not found in output" + assert "HBM usage after loading policy model:" in result.stdout, \ + "Policy model loading not found in output" + assert "Pre GRPO Training:" in result.stdout, \ + "Pre-training evaluation not found in output" + assert "Post GRPO Training:" in result.stdout, \ + "Post-training evaluation not found in output" + + print("GRPO Llama3 demo test passed successfully!") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From f40b3dfd01ab54281591125d5088b757b6ee70f7 Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 02:15:20 +0400 Subject: [PATCH 02/13] fix vllm installation Signed-off-by: Vladimir Suvorov --- .../examples/install_tunix_vllm_requirement.sh | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/src/MaxText/examples/install_tunix_vllm_requirement.sh b/src/MaxText/examples/install_tunix_vllm_requirement.sh index 0a89f96675..19a6205455 100644 --- a/src/MaxText/examples/install_tunix_vllm_requirement.sh +++ b/src/MaxText/examples/install_tunix_vllm_requirement.sh @@ -29,17 +29,9 @@ pip install aiohttp==3.12.15 pip install keyring keyrings.google-artifactregistry-auth # Install vLLM for Jax and TPUs from the artifact registry -VLLM_TARGET_DEVICE="tpu" pip install --no-cache-dir --pre \ - --index-url https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/ \ - --extra-index-url https://pypi.org/simple/ \ - --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \ - --extra-index-url https://download.pytorch.org/whl/nightly/cpu \ - --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \ - --find-links https://storage.googleapis.com/libtpu-wheels/index.html \ - --find-links https://storage.googleapis.com/libtpu-releases/index.html \ - --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \ - --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html \ - vllm==0.10.2rc2.dev59+gdcb28a332.tpu +git clone https://github.com/vllm-project/vllm.git /tmp/vllm +cd /tmp/vllm +VLLM_TARGET_DEVICE="tpu" pip install -e . # Install tpu-commons from the artifact registry pip install --no-cache-dir --pre \ From 4be48e2f1b9430c8cc2777c38282968f892d3c4c Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 02:17:47 +0400 Subject: [PATCH 03/13] Fix Signed-off-by: Vladimir Suvorov --- tests/integration_tests/grpo_llama3_demo_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration_tests/grpo_llama3_demo_test.py b/tests/integration_tests/grpo_llama3_demo_test.py index 7aaa532baa..afa068eef0 100644 --- a/tests/integration_tests/grpo_llama3_demo_test.py +++ b/tests/integration_tests/grpo_llama3_demo_test.py @@ -35,7 +35,6 @@ @pytest.mark.integration_test @pytest.mark.tpu_only -@pytest.mark.scheduled_only def test_grpo_llama3_demo(): """ Test the GRPO Llama3 demo script with minimal configuration. From 5094b2b49f7f843f846b206cc0d79c8fca4de70c Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 05:02:08 +0400 Subject: [PATCH 04/13] Fix Signed-off-by: Vladimir Suvorov --- .github/workflows/RunTests.yml | 20 +++- .github/workflows/grpo_integration_test.yml | 116 ++++++++++++++++++++ 2 files changed, 133 insertions(+), 3 deletions(-) create mode 100644 .github/workflows/grpo_integration_test.yml diff --git a/.github/workflows/RunTests.yml b/.github/workflows/RunTests.yml index a2851b13eb..2a5802ebbe 100644 --- a/.github/workflows/RunTests.yml +++ b/.github/workflows/RunTests.yml @@ -103,6 +103,20 @@ jobs: container_resource_option: "--privileged" is_scheduled_run: ${{ github.event_name == 'schedule' }} + tpu_grpo_integration_test: + needs: tpu_image + uses: ./.github/workflows/run_tests_internal.yml + with: + device_type: tpu + device_name: v4-8 + cloud_runner: linux-x86-ct4p-240-4tpu + pytest_marker: 'integration_test and tpu_only' + pytest_addopts: 'tests/integration_tests/grpo_llama3_demo_test.py' + xla_python_client_mem_fraction: 0.75 + tf_force_gpu_allow_growth: false + container_resource_option: "--privileged" + is_scheduled_run: ${{ github.event_name == 'schedule' }} + gpu_unit_tests: needs: gpu_image uses: ./.github/workflows/run_tests_internal.yml @@ -133,7 +147,7 @@ jobs: clean_up: if: ${{ always() }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] + needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_grpo_integration_test] name: "Clean up" runs-on: ["self-hosted"] permissions: @@ -152,7 +166,7 @@ jobs: notify_failure: name: Notify failed build # creates an issue or modifies last open existing issue for failed build - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] + needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_grpo_integration_test] if: ${{ always() }} runs-on: ubuntu-latest permissions: @@ -184,7 +198,7 @@ jobs: name: Close issue after 3 successful builds # This job runs only if all the preceding test jobs succeeded if: ${{ success() && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] + needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_grpo_integration_test] runs-on: ubuntu-latest permissions: issues: write diff --git a/.github/workflows/grpo_integration_test.yml b/.github/workflows/grpo_integration_test.yml new file mode 100644 index 0000000000..56c640a7b1 --- /dev/null +++ b/.github/workflows/grpo_integration_test.yml @@ -0,0 +1,116 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# GRPO Integration Test Workflow +# This workflow runs the GRPO Llama3 demo integration test on TPU + +name: GRPO Integration Test + +on: + pull_request: + paths: + - 'src/MaxText/examples/grpo_llama3_demo.py' + - 'tests/integration_tests/grpo_llama3_demo_test.py' + - 'src/MaxText/integration/tunix/**' + - 'src/MaxText/experimental/rl/**' + - '.github/workflows/grpo_integration_test.yml' + workflow_dispatch: + schedule: + # Run every 6 hours to catch regressions + - cron: '0 */6 * * *' + +concurrency: + group: > + ${{ + github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || + github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || + github.run_id + }} + cancel-in-progress: true + +jobs: + prelim: + runs-on: ["self-hosted"] + steps: + - name: Test gsutil installation + run: which gsutil >/dev/null 2>&1 || { echo >&2 "gsutil is required but not installed. Aborting"; exit 24;} + - name: Cleanup old docker images + run: docker system prune --all --force + + tpu_image: + needs: prelim + uses: ./.github/workflows/build_upload_internal.yml + with: + device_type: tpu + device_name: v4-8 + cloud_runner: linux-x86-n2-16-buildkit + build_mode: jax_ai_image + base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest + + grpo_integration_test: + needs: tpu_image + runs-on: ["self-hosted", "tpu", "v4-8"] + container: + image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu + env: + XLA_PYTHON_CLIENT_MEM_FRACTION: 0.75 + TF_FORCE_GPU_ALLOW_GROWTH: false + TPU_SKIP_MDS_QUERY: '' + options: "--privileged" + steps: + - uses: actions/checkout@v4 + - name: Run GRPO Integration Test + run: | + python3 -m pip install -e . --no-dependencies && + LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' python3 -m pytest -v -m "grpo_test" tests/integration_tests/grpo_llama3_demo_test.py --durations=0 + + clean_up: + if: ${{ always() }} + needs: [grpo_integration_test] + name: "Clean up" + runs-on: ["self-hosted"] + permissions: + contents: read + issues: write + steps: + - name: Authenticate gcloud + run: | + # configure registries as root and as runner + gcloud auth configure-docker --quiet + gcloud auth configure-docker us-docker.pkg.dev --quiet + - name: Delete the tpu image + run: gcloud container images delete "gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu" --force-delete-tags --quiet + + notify_failure: + name: Notify failed GRPO test + needs: [grpo_integration_test] + if: ${{ failure() }} + runs-on: ubuntu-latest + permissions: + issues: write + steps: + - name: Create failure issue + uses: actions/github-script@v7 + with: + script: | + const issue = await github.rest.issues.create({ + owner: context.repo.owner, + repo: context.repo.repo, + title: `GRPO Integration Test Failed - ${new Date().toISOString().split('T')[0]}`, + body: `The GRPO integration test failed in workflow run [${context.runId}](${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}). + + Please check the logs and fix any issues with the GRPO Llama3 demo.`, + labels: ['failed-build', 'grpo-test'] + }); + console.log(`Created issue #${issue.data.number}`); From 6969a05fef29f2ce920a10ffa32852d82458d89c Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 22:29:25 +0400 Subject: [PATCH 05/13] minus grpo Signed-off-by: Vladimir Suvorov --- .github/workflows/grpo_integration_test.yml | 116 -------------------- 1 file changed, 116 deletions(-) delete mode 100644 .github/workflows/grpo_integration_test.yml diff --git a/.github/workflows/grpo_integration_test.yml b/.github/workflows/grpo_integration_test.yml deleted file mode 100644 index 56c640a7b1..0000000000 --- a/.github/workflows/grpo_integration_test.yml +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright 2023–2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# GRPO Integration Test Workflow -# This workflow runs the GRPO Llama3 demo integration test on TPU - -name: GRPO Integration Test - -on: - pull_request: - paths: - - 'src/MaxText/examples/grpo_llama3_demo.py' - - 'tests/integration_tests/grpo_llama3_demo_test.py' - - 'src/MaxText/integration/tunix/**' - - 'src/MaxText/experimental/rl/**' - - '.github/workflows/grpo_integration_test.yml' - workflow_dispatch: - schedule: - # Run every 6 hours to catch regressions - - cron: '0 */6 * * *' - -concurrency: - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true - -jobs: - prelim: - runs-on: ["self-hosted"] - steps: - - name: Test gsutil installation - run: which gsutil >/dev/null 2>&1 || { echo >&2 "gsutil is required but not installed. Aborting"; exit 24;} - - name: Cleanup old docker images - run: docker system prune --all --force - - tpu_image: - needs: prelim - uses: ./.github/workflows/build_upload_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-n2-16-buildkit - build_mode: jax_ai_image - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest - - grpo_integration_test: - needs: tpu_image - runs-on: ["self-hosted", "tpu", "v4-8"] - container: - image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu - env: - XLA_PYTHON_CLIENT_MEM_FRACTION: 0.75 - TF_FORCE_GPU_ALLOW_GROWTH: false - TPU_SKIP_MDS_QUERY: '' - options: "--privileged" - steps: - - uses: actions/checkout@v4 - - name: Run GRPO Integration Test - run: | - python3 -m pip install -e . --no-dependencies && - LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' python3 -m pytest -v -m "grpo_test" tests/integration_tests/grpo_llama3_demo_test.py --durations=0 - - clean_up: - if: ${{ always() }} - needs: [grpo_integration_test] - name: "Clean up" - runs-on: ["self-hosted"] - permissions: - contents: read - issues: write - steps: - - name: Authenticate gcloud - run: | - # configure registries as root and as runner - gcloud auth configure-docker --quiet - gcloud auth configure-docker us-docker.pkg.dev --quiet - - name: Delete the tpu image - run: gcloud container images delete "gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu" --force-delete-tags --quiet - - notify_failure: - name: Notify failed GRPO test - needs: [grpo_integration_test] - if: ${{ failure() }} - runs-on: ubuntu-latest - permissions: - issues: write - steps: - - name: Create failure issue - uses: actions/github-script@v7 - with: - script: | - const issue = await github.rest.issues.create({ - owner: context.repo.owner, - repo: context.repo.repo, - title: `GRPO Integration Test Failed - ${new Date().toISOString().split('T')[0]}`, - body: `The GRPO integration test failed in workflow run [${context.runId}](${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}). - - Please check the logs and fix any issues with the GRPO Llama3 demo.`, - labels: ['failed-build', 'grpo-test'] - }); - console.log(`Created issue #${issue.data.number}`); From da99ec4f1be8b3ab3e584f251943301624e5405f Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 22:34:03 +0400 Subject: [PATCH 06/13] minus grpo Signed-off-by: Vladimir Suvorov --- .github/workflows/RunTests.yml | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/.github/workflows/RunTests.yml b/.github/workflows/RunTests.yml index 2a5802ebbe..2fe22f28c1 100644 --- a/.github/workflows/RunTests.yml +++ b/.github/workflows/RunTests.yml @@ -103,20 +103,6 @@ jobs: container_resource_option: "--privileged" is_scheduled_run: ${{ github.event_name == 'schedule' }} - tpu_grpo_integration_test: - needs: tpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'integration_test and tpu_only' - pytest_addopts: 'tests/integration_tests/grpo_llama3_demo_test.py' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - gpu_unit_tests: needs: gpu_image uses: ./.github/workflows/run_tests_internal.yml From 33653a3c1f75fb2d905c6ad1d1e9ec004e4f07e1 Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 22:35:06 +0400 Subject: [PATCH 07/13] minus grpo Signed-off-by: Vladimir Suvorov --- .github/workflows/RunTests.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/RunTests.yml b/.github/workflows/RunTests.yml index 2fe22f28c1..a2851b13eb 100644 --- a/.github/workflows/RunTests.yml +++ b/.github/workflows/RunTests.yml @@ -133,7 +133,7 @@ jobs: clean_up: if: ${{ always() }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_grpo_integration_test] + needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] name: "Clean up" runs-on: ["self-hosted"] permissions: @@ -152,7 +152,7 @@ jobs: notify_failure: name: Notify failed build # creates an issue or modifies last open existing issue for failed build - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_grpo_integration_test] + needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] if: ${{ always() }} runs-on: ubuntu-latest permissions: @@ -184,7 +184,7 @@ jobs: name: Close issue after 3 successful builds # This job runs only if all the preceding test jobs succeeded if: ${{ success() && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_grpo_integration_test] + needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] runs-on: ubuntu-latest permissions: issues: write From aa9dca37c3fb2ff55eecce373ad2261f07ed14fd Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 22:43:47 +0400 Subject: [PATCH 08/13] fix linter Signed-off-by: Vladimir Suvorov --- .../grpo_llama3_demo_test.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/tests/integration_tests/grpo_llama3_demo_test.py b/tests/integration_tests/grpo_llama3_demo_test.py index afa068eef0..a846f76b79 100644 --- a/tests/integration_tests/grpo_llama3_demo_test.py +++ b/tests/integration_tests/grpo_llama3_demo_test.py @@ -38,34 +38,34 @@ def test_grpo_llama3_demo(): """ Test the GRPO Llama3 demo script with minimal configuration. - + This test verifies that the demo can: 1. Load the dataset 2. Initialize the model 3. Run a few training steps 4. Complete evaluation """ - + # Create a temporary directory for outputs with tempfile.TemporaryDirectory() as temp_dir: # Set environment variables for the test env = os.environ.copy() env["HOME"] = temp_dir env["SKIP_JAX_PRECOMPILE"] = "1" - + # Build the command to run the demo demo_script = os.path.join( - MAXTEXT_REPO_ROOT, + MAXTEXT_REPO_ROOT, "src/MaxText/examples/grpo_llama3_demo.py" ) - + # We'll modify the script to run with minimal steps for testing # by setting environment variables that the script can read env["GRPO_TEST_MODE"] = "1" env["GRPO_NUM_BATCHES"] = "2" # Minimal batches for testing env["GRPO_NUM_TEST_BATCHES"] = "2" # Minimal test batches env["GRPO_MAX_STEPS"] = "2" # Minimal training steps - + # Run the demo script result = subprocess.run( [sys.executable, demo_script], @@ -73,8 +73,9 @@ def test_grpo_llama3_demo(): capture_output=True, text=True, timeout=600, # 10 minute timeout + check=False, ) - + # Check if the script executed successfully if result.returncode != 0: print("STDOUT:", result.stdout) @@ -82,7 +83,7 @@ def test_grpo_llama3_demo(): raise AssertionError( f"GRPO Llama3 demo failed with return code {result.returncode}" ) - + # Verify expected outputs in the logs assert "HBM usage before loading model:" in result.stdout, \ "Model initialization not found in output" @@ -94,7 +95,7 @@ def test_grpo_llama3_demo(): "Pre-training evaluation not found in output" assert "Post GRPO Training:" in result.stdout, \ "Post-training evaluation not found in output" - + print("GRPO Llama3 demo test passed successfully!") From 2d0a7313dc26cfc3eded0544a0e6ec5242adee78 Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Thu, 2 Oct 2025 22:48:17 +0400 Subject: [PATCH 09/13] fix linter Signed-off-by: Vladimir Suvorov --- .../grpo_llama3_demo_test.py | 24 ++++++------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/tests/integration_tests/grpo_llama3_demo_test.py b/tests/integration_tests/grpo_llama3_demo_test.py index a846f76b79..ec0546dd78 100644 --- a/tests/integration_tests/grpo_llama3_demo_test.py +++ b/tests/integration_tests/grpo_llama3_demo_test.py @@ -54,10 +54,7 @@ def test_grpo_llama3_demo(): env["SKIP_JAX_PRECOMPILE"] = "1" # Build the command to run the demo - demo_script = os.path.join( - MAXTEXT_REPO_ROOT, - "src/MaxText/examples/grpo_llama3_demo.py" - ) + demo_script = os.path.join(MAXTEXT_REPO_ROOT, "src/MaxText/examples/grpo_llama3_demo.py") # We'll modify the script to run with minimal steps for testing # by setting environment variables that the script can read @@ -80,21 +77,14 @@ def test_grpo_llama3_demo(): if result.returncode != 0: print("STDOUT:", result.stdout) print("STDERR:", result.stderr) - raise AssertionError( - f"GRPO Llama3 demo failed with return code {result.returncode}" - ) + raise AssertionError(f"GRPO Llama3 demo failed with return code {result.returncode}") # Verify expected outputs in the logs - assert "HBM usage before loading model:" in result.stdout, \ - "Model initialization not found in output" - assert "HBM usage after loading ref model:" in result.stdout, \ - "Reference model loading not found in output" - assert "HBM usage after loading policy model:" in result.stdout, \ - "Policy model loading not found in output" - assert "Pre GRPO Training:" in result.stdout, \ - "Pre-training evaluation not found in output" - assert "Post GRPO Training:" in result.stdout, \ - "Post-training evaluation not found in output" + assert "HBM usage before loading model:" in result.stdout, "Model initialization not found in output" + assert "HBM usage after loading ref model:" in result.stdout, "Reference model loading not found in output" + assert "HBM usage after loading policy model:" in result.stdout, "Policy model loading not found in output" + assert "Pre GRPO Training:" in result.stdout, "Pre-training evaluation not found in output" + assert "Post GRPO Training:" in result.stdout, "Post-training evaluation not found in output" print("GRPO Llama3 demo test passed successfully!") From 7a63428d76dfbe045ca78023343f1f77983140ee Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Fri, 3 Oct 2025 00:37:08 +0400 Subject: [PATCH 10/13] - grpo in integration tests Signed-off-by: Vladimir Suvorov --- .../grpo_llama3_demo_test.py | 93 ------------------- 1 file changed, 93 deletions(-) delete mode 100644 tests/integration_tests/grpo_llama3_demo_test.py diff --git a/tests/integration_tests/grpo_llama3_demo_test.py b/tests/integration_tests/grpo_llama3_demo_test.py deleted file mode 100644 index ec0546dd78..0000000000 --- a/tests/integration_tests/grpo_llama3_demo_test.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright 2023–2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Integration test for GRPO Llama3 demo. - -This test runs the grpo_llama3_demo.py script with minimal configuration -to verify the GRPO training pipeline works correctly on TPU. - -ATTENTION: This test should only be run on TPU (v4-8, v5p-8, or v6e-8). - -Usage: - pytest tests/integration_tests/grpo_llama3_demo_test.py -""" - -import os -import subprocess -import sys -import tempfile -import pytest - -from MaxText.globals import MAXTEXT_REPO_ROOT - - -@pytest.mark.integration_test -@pytest.mark.tpu_only -def test_grpo_llama3_demo(): - """ - Test the GRPO Llama3 demo script with minimal configuration. - - This test verifies that the demo can: - 1. Load the dataset - 2. Initialize the model - 3. Run a few training steps - 4. Complete evaluation - """ - - # Create a temporary directory for outputs - with tempfile.TemporaryDirectory() as temp_dir: - # Set environment variables for the test - env = os.environ.copy() - env["HOME"] = temp_dir - env["SKIP_JAX_PRECOMPILE"] = "1" - - # Build the command to run the demo - demo_script = os.path.join(MAXTEXT_REPO_ROOT, "src/MaxText/examples/grpo_llama3_demo.py") - - # We'll modify the script to run with minimal steps for testing - # by setting environment variables that the script can read - env["GRPO_TEST_MODE"] = "1" - env["GRPO_NUM_BATCHES"] = "2" # Minimal batches for testing - env["GRPO_NUM_TEST_BATCHES"] = "2" # Minimal test batches - env["GRPO_MAX_STEPS"] = "2" # Minimal training steps - - # Run the demo script - result = subprocess.run( - [sys.executable, demo_script], - env=env, - capture_output=True, - text=True, - timeout=600, # 10 minute timeout - check=False, - ) - - # Check if the script executed successfully - if result.returncode != 0: - print("STDOUT:", result.stdout) - print("STDERR:", result.stderr) - raise AssertionError(f"GRPO Llama3 demo failed with return code {result.returncode}") - - # Verify expected outputs in the logs - assert "HBM usage before loading model:" in result.stdout, "Model initialization not found in output" - assert "HBM usage after loading ref model:" in result.stdout, "Reference model loading not found in output" - assert "HBM usage after loading policy model:" in result.stdout, "Policy model loading not found in output" - assert "Pre GRPO Training:" in result.stdout, "Pre-training evaluation not found in output" - assert "Post GRPO Training:" in result.stdout, "Post-training evaluation not found in output" - - print("GRPO Llama3 demo test passed successfully!") - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) From faf2bef93dde64ad7e9e1521a0e2d7495d79e270 Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Sat, 4 Oct 2025 02:42:51 +0400 Subject: [PATCH 11/13] fix path Signed-off-by: Vladimir Suvorov --- src/MaxText/examples/grpo_llama3_demo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/examples/grpo_llama3_demo.py b/src/MaxText/examples/grpo_llama3_demo.py index ee5997b902..a1e070ef76 100644 --- a/src/MaxText/examples/grpo_llama3_demo.py +++ b/src/MaxText/examples/grpo_llama3_demo.py @@ -398,7 +398,7 @@ def get_ref_maxtext_model(config): config_ref = pyconfig.initialize( [ "", - f"{HOME}/maxtext/src/MaxText/configs/base.yml", + os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"), ], base_output_directory="dummy", # This is not used in Tunix. run_name="test-tunix-maxtext-llama3.1-8b", From c2049cc4728198c8751c03bc42e434bee6cda913 Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Sat, 4 Oct 2025 03:02:53 +0400 Subject: [PATCH 12/13] fix path Signed-off-by: Vladimir Suvorov --- src/MaxText/examples/grpo_llama3_demo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/examples/grpo_llama3_demo.py b/src/MaxText/examples/grpo_llama3_demo.py index a1e070ef76..758364f4e7 100644 --- a/src/MaxText/examples/grpo_llama3_demo.py +++ b/src/MaxText/examples/grpo_llama3_demo.py @@ -398,7 +398,7 @@ def get_ref_maxtext_model(config): config_ref = pyconfig.initialize( [ "", - os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"), + os.path.join(MAXTEXT_ASSETS_ROOT, "configs", "base.yml"), ], base_output_directory="dummy", # This is not used in Tunix. run_name="test-tunix-maxtext-llama3.1-8b", From 2347d0a8f22797a03ea0e004754eccd771849e78 Mon Sep 17 00:00:00 2001 From: Vladimir Suvorov Date: Sat, 4 Oct 2025 03:05:02 +0400 Subject: [PATCH 13/13] fix path Signed-off-by: Vladimir Suvorov --- src/MaxText/examples/grpo_llama3_demo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/examples/grpo_llama3_demo.py b/src/MaxText/examples/grpo_llama3_demo.py index 758364f4e7..a1e070ef76 100644 --- a/src/MaxText/examples/grpo_llama3_demo.py +++ b/src/MaxText/examples/grpo_llama3_demo.py @@ -398,7 +398,7 @@ def get_ref_maxtext_model(config): config_ref = pyconfig.initialize( [ "", - os.path.join(MAXTEXT_ASSETS_ROOT, "configs", "base.yml"), + os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"), ], base_output_directory="dummy", # This is not used in Tunix. run_name="test-tunix-maxtext-llama3.1-8b",