Skip to content

Adding single-host integration test for GRPO and GSPO.#3409

Merged
copybara-service[bot] merged 1 commit intomainfrom
nicogrande/rl-integration-test
Mar 16, 2026
Merged

Adding single-host integration test for GRPO and GSPO.#3409
copybara-service[bot] merged 1 commit intomainfrom
nicogrande/rl-integration-test

Conversation

@NicoGrande
Copy link
Copy Markdown
Collaborator

@NicoGrande NicoGrande commented Mar 13, 2026

Description

Adds an end-to-end integration test for both GRPO and GSPO workflows on qwen3-0.6b (which should work e2e on a single v6e-4 TPU VM).

FIXES: b/490410666, b/490411100, b/490409949, b/490410743

Tests

Tests are skipped for now pending debugging on v6e-4

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov
Copy link
Copy Markdown

codecov Bot commented Mar 13, 2026

Codecov Report

❌ Patch coverage is 12.82051% with 34 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/trainers/post_train/rl/train_rl.py 12.82% 34 Missing ⚠️

📢 Thoughts on this report? Let us know!

@NicoGrande NicoGrande force-pushed the nicogrande/rl-integration-test branch 5 times, most recently from 49adda4 to 7d0bd50 Compare March 13, 2026 18:50
xuefgu
xuefgu reviewed Mar 13, 2026
View reviewed changes
Copy link
Copy Markdown
Collaborator

@xuefgu xuefgu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Many thanks Nico! Just a couple of quick clarifications.

Comment thread src/maxtext/configs/inference/vllm.yml

base_config: "base.yml"
attention: "vllm_rpa"
model_call_mode: "inference"
Copy link
Copy Markdown
Collaborator

@xuefgu xuefgu Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is this new config used below?

Copy link
Copy Markdown
Collaborator Author

@NicoGrande NicoGrande Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for MoE! Its unrelated but I wanted to squeeze it in to this PR since we should always be setting this in the vLLM codepath.

Copy link
Copy Markdown
Collaborator

@dipannita08 dipannita08 Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the entry point for this through train_rl?

Copy link
Copy Markdown
Collaborator Author

@NicoGrande NicoGrande Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, vllm.yml is only used in the MaxText on vLLM adapter class here: https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py#L73

Comment thread tests/integration/single_host_train_rl_test.py
"load_parameters_path=gs://maxtext-model-checkpoints/qwen3-0.6b/2025-10-27/scanned/0/items",
]

def _run_rl_workflow_end_to_end(self, extra_argv):
Copy link
Copy Markdown
Collaborator

@xuefgu xuefgu Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need this vs. directly calling rl_train() for various configs?

Copy link
Copy Markdown
Collaborator Author

@NicoGrande NicoGrande Mar 13, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The biggest benefit is to be able to access some of the internals of rl_train(). For instance, this allows us to assert that the actor model params are being updated.

Copy link
Copy Markdown
Collaborator

@xuefgu xuefgu Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The concern is the gradual/eventual divergence of logic in rl_train() and the tests. Any way we can mitigate that?

Copy link
Copy Markdown
Collaborator Author

@NicoGrande NicoGrande Mar 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Its a good point - I could add a couple of test cases that just calls rl_train() directly to test that API also.

xuefgu
xuefgu approved these changes Mar 13, 2026
View reviewed changes
Copy link
Copy Markdown
Collaborator

@xuefgu xuefgu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deferring to you on how to handle the divergence between tests and application logic. Thanks again for the PR!

@NicoGrande NicoGrande force-pushed the nicogrande/rl-integration-test branch from 7d0bd50 to c6284a2 Compare March 13, 2026 20:02
@NicoGrande NicoGrande force-pushed the nicogrande/rl-integration-test branch from c6284a2 to c45b59c Compare March 13, 2026 20:06
@copybara-service copybara-service Bot merged commit 01fbe6d into main Mar 16, 2026
57 of 62 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@dipannita08 dipannita08 dipannita08 approved these changes

@khatwanimohit khatwanimohit khatwanimohit approved these changes

@xuefgu xuefgu xuefgu approved these changes

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@shralex shralex Awaiting requested review from shralex shralex is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@jesselu-google jesselu-google Awaiting requested review from jesselu-google jesselu-google is a code owner

@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@NicoGrande @dipannita08 @khatwanimohit @xuefgu