feat: add grpo trainer and simplify gsm8k grpo example

[dhh1995]

Description

As per #542 , wrap training scripts into trainers.

This PR extracts the common parts in the training examples into a trainer that supports both FSDP and Megatron backends into a GRPOTrainer. It also provides a much simpler example script that uses the trainer.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not
  work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with
  jb build docs
• No critical issues raised by AI reviewers (/gemini review)

Breaking Change Details (if applicable):

Additional Context

---

Need help? Check the Contributing Guide or ask in
GitHub Discussions!

[gemini-code-assist]

Summary of Changes

Hello @dhh1995, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the training infrastructure by introducing a dedicated GRPOTrainer class. This trainer centralizes the common components required for GRPO training, offering support for both FSDP and Megatron distributed training backends. The change aims to streamline the development of new training examples and improve code maintainability, as demonstrated by a simplified GSM8K example.

Highlights

• New GRPOTrainer Class: Introduced a new GRPOTrainer class in areal/trainer/rl.py to encapsulate common training logic for GRPO, supporting both FSDP and Megatron backends. This refactors existing training scripts into a reusable component.
• Backend Support: The GRPOTrainer is designed to work with both FSDP (Fully Sharded Data Parallel) and Megatron distributed training backends, providing a unified interface for different parallel strategies.
• Simplified GSM8K Example: A new example script, examples/math/gsm8k_grpo_trainer.py, has been added. This script demonstrates a much simpler way to train a GRPO model for the GSM8K task by leveraging the newly created GRPOTrainer.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a GRPOTrainer to encapsulate the training logic for GRPO, which is a great step towards simplifying the training examples. The new gsm8k_grpo_trainer.py example demonstrates this simplification effectively. The overall structure of the GRPOTrainer is well-designed, with clear separation of concerns for model creation, data loading, and the training loop.

My review includes a few suggestions:

• A high-severity fix to prevent a crash when no validation dataset is provided.
• A medium-severity improvement for more robust handling of environment variables.
• A medium-severity style suggestion to follow PEP 8 import conventions in the example script.

Overall, this is a solid contribution that improves the usability and structure of the training code.

[garrett4wade]

Please respect gemini's comments as well. Overall LGTM.