GPT-OSS: Add user guide and tests

[shuningjin]

Description

Main change: for gpt-oss

• add user guide:
  • end_to_end/tpu/gpt_oss/run_gpt_oss.md
  • use gpt-oss-20b on v5p-8 (sparse_matmul + megablox) as example
• add test script
  • end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh
  • end_to_end/tpu/gpt_oss/120b/test_gpt_oss.sh
• pyconfig validation

Other change: format output from forward_logit_checker, correct typos in other files

Tests

single script for tpu: checkpoint conversion, logit check, train, sft, decode (instead of splitting conversion to another)

gpt-oss-20b, end-to-end test on v5p-8 cluster, takes around 1h

• cmd: https://paste.googleplex.com/5491652676419584
• log: https://cloudlogging.app.goo.gl/zbAppeG8m3DjGhim8
• to faciliate the test, I have uploaded the dequantized checkpoint and golden logits to GCS
• note: gpt-oss with flash attention requires jax>=0.7.2

The test will also be onboarded to XLML: GoogleCloudPlatform/ml-auto-solutions#913

Additional Notes

HF tokenizer:

• Here we directly use the HF tokenizer, instead of storing it as assets in the repo. This practice seems fine, as long as the HF access token is handled correctly.

HF access token:

• For MaxText native training, the access token is not needed because the HF model/tokenizer (gpt-oss) and HF dataset (used in SFT) are not gated.
• However, decoding with HF tokenizer still needs access token despite that the model is not gated. The tokenizer is built from JetStream utils, which handles things differently from MaxText.

Checklist

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

• I have performed a self-review of my code.
• 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.

[RissyRan]

LGTM! Just a few minor comments

[RissyRan]

Could you also add announcement to the Main page to claim the support for gpt-oss models?

[shuningjin]

Could you also add announcement to the Main page to claim the support for gpt-oss models?

I just updated "README.md" and "supported_models_and_architectures.md".

[RissyRan]

Thanks!

[github-actions]

🤖 Hi @shuningjin, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

📋 Review Summary

This pull request introduces support for the GPT-OSS family of models, including comprehensive user guides, test scripts, and necessary configuration updates. The changes are well-structured and the documentation is thorough.

🔍 General Feedback

• The addition of end-to-end test scripts for gpt-oss-20b and gpt-oss-120b is excellent for ensuring model correctness and providing clear usage examples.
• Several minor typos and path corrections in existing documentation have been addressed, improving overall quality.
• The validation logic added to pyconfig.py is a good proactive measure to prevent misconfiguration of the new models.
• There were critical formatting errors in the new test scripts that would have caused them to fail. These have been highlighted in the inline comments.

[parambole]

Thank you for adding this. I have left a few comments.

[parambole]

Awesome work!

[github-actions]

🤖 Hi @shuningjin, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.