Update veRL to 0.7.0

[chenyushuo]

Description

As the title says.

Checklist

Please check the following items before code is ready to be reviewed.

• Code has passed all tests
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review

[gemini-code-assist]

Summary of Changes

Hello @chenyushuo, 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 integrates the latest verl library, bringing a suite of enhancements to the training infrastructure. It introduces more sophisticated LoRA and optimizer configurations, alongside significant improvements to distributed training workers for FSDP and Megatron. The update also incorporates a new rollout correction mechanism to refine policy learning and includes comprehensive profiling and memory management utilities. These changes collectively aim to boost the flexibility, performance, and debuggability of large-scale model training.

Highlights

• veRL Dependency Update: The core verl dependency has been upgraded from version 0.5.0 to 0.7.0, indicating a significant update to the underlying reinforcement learning library.
• Enhanced LoRA Configuration: New configuration options for LoRA (Low-Rank Adaptation) have been introduced, including exclude_modules and lora_adapter_path, allowing for more granular control over LoRA application and the ability to load pre-trained LoRA adapters.
• Optimizer and Scheduler Refinements: The optimizer configuration now defaults to AdamW, and the learning rate scheduler type is more flexible. The optimizer creation process has been refactored to use a dedicated build_optimizer utility, improving modularity.
• FSDP and Megatron Worker Improvements: Both FSDP and Megatron distributed training workers have been updated to include advanced profiling capabilities, configurable NCCL timeouts, and more robust model loading and checkpointing logic, especially for multimodal models and custom architectures.
• Rollout Correction Mechanism: A new RolloutCorrection mechanism has been added to the algorithm configuration, introducing fields like rollout_is and bypass_mode to manage importance sampling weights and off-policy corrections during training.
• Memory Management and Profiling: Several changes introduce explicit calls to aggressive_empty_cache and set_expandable_segments for better GPU memory management, alongside integration with a distributed profiler for performance analysis.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 updates the veRL dependency to version 0.7.0. The changes are extensive, reflecting a major version upgrade. Key modifications include updates to configuration classes in trinity/common/verl_config.py to align with new veRL features, and significant refactoring in worker implementations (fsdp_workers.py, megatron_workers.py) to support new functionalities like advanced profiling, router replay for MoE models, and improved LoRA adapter handling. Checkpointing logic has also been enhanced for better reliability and to support asynchronous operations. Overall, the changes are well-integrated and necessary for the dependency upgrade. I have one minor suggestion to ensure consistency in metric logging.

[gemini-code-assist]

The learning rate lr obtained from get_last_lr() can be a tensor. To ensure it's a scalar value for metric logging and to maintain consistency with how it's handled in update_actor, it's safer to convert it to a float if it's a tensor.

Suggested change

	metrics["critic/lr"] = lr
	metrics["critic/lr"] = lr.item() if torch.is_tensor(lr) else lr