[WIP] lmms-engine RL infra - #186
Open
pufanyi wants to merge 27 commits into
Open
Conversation
- Introduced a new shell script `run_vizdoom_qwen3_5_4b.sh` for launching the RL training with VizDoom. - Created a configuration file `vizdoom_qwen3_5_4b.yaml` for specifying training parameters and dataset configurations. - Implemented `lmms_eval` paths and task loader to facilitate the integration of LMMS evaluation with RL training. - Developed `FSDP2GRPORLTrainer` class for handling the training process, including rollout management and checkpointing. - Enhanced the training step to compute GRPO loss and manage policy weight synchronization. - Added utility functions for managing RL configurations and data loading.
…onfiguration updates
- Implemented RL training loop in `src/lmms_engine/launch/rl.py`. - Introduced `TransformersChatModelServer` for model serving in `src/lmms_engine/rl/model_server/hf.py`. - Created GRPO training adapters in `src/lmms_engine/train/rl/grpo/` including configuration and payload handling. - Developed `RLTrainRunner` for managing online RL training in `src/lmms_engine/train/rl/runner.py`. - Added necessary utility functions and data processing for RL training.
- Created a new configuration file `_vizdoom.ini` for ViZDoom with paths for IWADs, auto-execution scripts, and global settings. - Introduced a new module `lmms_engine.rl.ray` with classes for managing Ray cluster specifications, worker descriptors, and resource planning. - Implemented `RayNodeScheduler` for managing node roles and starting Ray clusters. - Added `RayRLMultinodeRuntime` to facilitate multi-node reinforcement learning training. - Developed training logic in `train.py` to integrate Ray Train with the RL training framework.
…proved batch processing
- Updated InMemoryDataBuffer to resolve global train batch size and minimum trajectories from config. - Enforced explicit model server requirement in build_rollout_episode_specs for RL rollout tasks. - Enhanced Ray model server initialization with runtime environment settings and validation for GPU resources. - Modified VLLMChatModelServer to accept new parameters for chat template content format and multimodal processor kwargs. - Improved RayClusterSpec and RayResourcePlan to derive configurations from a provided dictionary instead of environment variables. - Implemented sharding for RL train batches to ensure proper distribution across multiple workers. - Added validation for vLLM model server configurations to ensure compliance with expected structure. - Updated training engine configuration to remove deprecated parameters and enforce new structure.
…nd logging - Added support for policy weight synchronization in the RLTrainRunner class, including configuration options for initial synchronization and draining inflight rollouts. - Implemented logging for training rank placement and shard balance during training. - Introduced a new method to build a weight synchronization client for Ray actors. - Enhanced the sharding of training batches to balance trajectory steps across GPUs. - Added new utility functions for managing weight synchronization directories and validating configurations. - Created a smoke test for weight synchronization across multiple GPUs, ensuring proper actor placement and checkpoint validation. - Developed unit tests for orchestrator backpressure handling and policy weight synchronization functionality.
- Added disk delta functionality for efficient weight synchronization, including publishing and applying deltas. - Introduced ModelServerManager to manage multiple model servers by role, supporting ray_actor_pool and external OpenAI-compatible servers. - Created ReferenceLogprobAnnotator for annotating trajectories with log probabilities from a reference model. - Enhanced unit tests for policy model server validation, trajectory annotation, and resource planning in Ray. - Refactored existing code to support new features and improve maintainability.
kcz358
reviewed
Jul 7, 2026
Collaborator
There was a problem hiding this comment.
Can split this into a different config in examples and pass with --config-name and --config-path to split with original default config
kcz358
reviewed
Jul 7, 2026
Collaborator
There was a problem hiding this comment.
Maybe can put this into examples rl or in docs.
kcz358
reviewed
Jul 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.