Add v6e-8 GSM8K RL convergence benchmark config

[AntonyMei]

Add a Qwen3-4B non-agentic GRPO recipe for GSM8K on a single v6e-8 host. The config pins the v6e-8 trainer/rollout allocation, uses a one-epoch benchmark budget, and keeps the explicit step count and LR schedule horizon aligned with the computed optimizer updates.

Tests

Tested on a v6e-8 TPU node. The model quality (mean rewards) improves significantly (from ~0.75 to ~1.05) after ~200 training steps.

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.

[AntonyMei]

Command used for training:

export BASE_OUTPUT_DIRECTORY="gs://yixuanm-maxtext-logs/ConvergenceTesting"
export RUN_NAME="non-agentic-0706-test-sync"
export MAXTEXT_CKPT_PATH="gs://yixuanm-maxtext-logs/AgenticGRPOTesting/qwen3-4b-ckpt/0/items/"
export NEW_MODEL_DESIGN=1
export SKIP_JAX_PRECOMPILE=1
export TPU_LOG_DIR=~/tpu_logs

python3 -m maxtext.trainers.post_train.rl.train_rl \
  src/maxtext/configs/post_train/rl_gsm8k_qwen3_4b_v6e8.yml \
  load_parameters_path=${MAXTEXT_CKPT_PATH?} \
  run_name=${RUN_NAME?} \
  base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
  enable_tunix_perf_metrics=True

Results can be viewed with

tensorboard --logdir gs://yixuanm-maxtext-logs/ConvergenceTesting/non-agentic-0706-test-sync/tensorboard --load_fast=false

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!