feat(test): add modular E2E TPU testing scripts for GPTOSS-20B

[RexBearIU]

Description

This PR implements modular end-to-end (E2E) TPU testing pipeline scripts for the GPTOSS-20B model, matching the exact modular pattern used by Gemma3-4B in tests/end_to_end/tpu/gemma3/4b.

This change is made to establish clean, modular end-to-end regression validation for GPTOSS-20B training, post-training, reinforcement learning, and conversion pipelines.

The PR introduces the following modular testing scripts:

• test_gpt_oss_to_mt.sh: Converts Hugging Face weights to scanned and unscanned Orbax checkpoints, manages dependency setup (consolidated to run once), and executes forward-pass logit verification with robust KL divergence tolerance (--max_kl_div=0.01), matching Gemma3 test patterns.
• test_gpt_oss.sh: Refactored baseline pre-training script on synthetic datasets with dynamic device and hardware configurations.
• test_gpt_oss_sft.sh: Validates the Supervised Fine-Tuning (SFT) pipeline on conversational datasets, with verification decoding.
• test_gpt_oss_rl.sh: Validates Reinforcement Learning (RL) using GRPO, matching MoE vLLM parameters.
• test_gpt_oss_to_hf.sh: Converts the MaxText-compatible checkpoint back to Hugging Face format (running as the final pipeline step).

All scripts support seamless running in TPU environments to facilitate automated verification and testing.

Important

Dependency / Blocking: This PR is blocked by and depends on #4560. The checkpoint conversion and forward pass logit checkers in this PR depend on the --pure-nnx formatting introduced in #4560 to load and verify pure-NNX checkpoints correctly.

Tests

Every script has been meticulously checked and verified on actual TPU VM architectures:

• Logit Checker Correctness: Successfully converted pure-NNX checkpoint and validated forward-pass logit equivalence:
  • Max Absolute Difference: 0.36 (well within standard limits)
  • Max KL Divergence: 0.0018 (fully passes the robust --max_kl_div=0.01 threshold)
  • Rank agreement percentage: 100.0%
  • Overlap count: 10/10
• Decoding Quality: Completed successfully, loading the checkpoint and running inference to generate coherent text (Index 0, 1, 2 all pass).
• Pre-training Validation: Successfully compiled and completed synthetic pre-training.

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]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!