[bugfix] fix grpo mllm multi turn

[hjh0119]

1. Fixed chunking of multimodal data when the number of rollout samples is dynamic.
2. overide the multimodal data from rollout_infos.

#5818

[gemini-code-assist]

Summary of Changes

Hello @hjh0119, 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 addresses key issues in handling multimodal data within GRPO multi-turn training. It refactors the input chunking logic to ensure accurate processing of multimodal data, particularly when dealing with dynamic sample sizes. Additionally, it introduces a mechanism to override multimodal content using rollout_infos, significantly improving the adaptability and control over multimodal interactions during training. These enhancements are reflected in both the core trainer logic and the accompanying developer documentation.

Highlights

• Multimodal Data Chunking Fix: The input chunking mechanism in GRPOTrainer has been refactored to correctly handle multimodal data, especially when the number of rollout samples is dynamic. This involved moving the chunking logic into a dedicated get_chunked_inputs method, ensuring proper processing for both LLM and MLLM inputs.
• Rollout Info Override for Multimodal Data: New functionality has been implemented to allow dynamic modification and overriding of multimodal data (specifically images, audios, and videos) using rollout_infos within the merge_output_input_data function. This enhances flexibility in multi-turn multimodal interactions.
• Documentation Update: New sections have been added to both the Chinese and English developer guides. These sections explain how to leverage rollout_infos for dynamic multimodal data modification in multi-turn scenarios, providing clear guidance for developers.

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 addresses a bug in multimodal data chunking for GRPO and introduces a feature to override multimodal data using rollout_infos. The core change involves refactoring the data chunking logic into a new get_chunked_inputs method, which correctly handles re-encoding for multimodal models. The data merging logic is also updated to support overriding from rollout_infos. The documentation has been updated to reflect these changes. The implementation looks solid and effectively fixes the bug while adding useful functionality. I have one minor suggestion to improve code maintainability.

[gemini-code-assist]

For better maintainability and to avoid magic strings, consider defining ['images', 'videos', 'audios'] as a class-level or module-level constant. This makes it easier to manage and reuse these keys if they are needed elsewhere in the class.