add MaxText GRPO / RL post-training sample on GKE TPU v7x

[pmotgi]

@## Description

This PR introduces a comprehensive new sample for running MaxText GRPO (Reinforcement Learning) post-training on Google Kubernetes Engine (GKE) using Cloud TPUs.

It provides a production-tested foundation for aligning Large Language Models (like Llama-3.1-8B-Instruct) using Group Relative Policy Optimization on TPU v7x infrastructure.

Key Additions:

• Kubernetes Manifests: Includes both DWS-Flex and Spot JobSet manifests (2x2x1 TPU v7x topology) and Kueue cluster queue configuration.
• Architecture & Docker Deployment: A dedicated runner image leveraging MaxText and vLLM (for rollout sampling) concurrently on the same TPU slice.
• TRL to MaxText Mapping Guide: A critical migration guide for users coming from HuggingFace TRL on GPUs. It explicitly details hyperparameter translation gotchas, such as how batch_size vs completions are counted, and sequence length padding behaviors to prevent HBM OOMs.
• Customization Guides: Includes a custom rewards template (custom_rewards_template.py) and a custom dataset integration guide.
• MFU Calculation: Includes a post-training MFU script (calculate_mfu.py) to measure hardware utilization.

co-contributor: @rick-c-goog

[moficodes]

Few file paths are pointing to your local file system.

There are a lot of variable in yaml and commands that needs to be replaces manually by user. Look into using envsubst or let users know how and when to replace them.

[pmotgi]

Few file paths are pointing to your local file system.

There are a lot of variable in yaml and commands that needs to be replaces manually by user. Look into using envsubst or let users know how and when to replace them.

@moficodes - Thanks for your feedbacks. I have updated the PR with all the recommended changes. Let me know if there are any further changes needed. Thanks!

[moficodes]

Couple more clarifications.

[pmotgi]

Couple more clarifications.

no worries @moficodes - I just made few changes. PTAL and let me know if any changes are needed further. Thanks!

[moficodes]

lgtm