Add Eagle3 training for more MLLM model
Motivation
In addition to the previously supported qwen2.5 vl, more multimodal language models have emerged, such as omni, dsvl, internvl, qwen3vl, etc. An abstract solution needs to be designed to easily support the training of draft models of these new models.
Modifications
QwenVLOnlineEagle3Modelwas removed because, for the draft model, only hidden states and the target are needed for SFT fine-tuning;pixel_values,image_grid_thw, etc., only need to be provided to the target model during the prepare_data stage. Therefore, for the draft model, whether it's LLM or MLLM, it has the same input and output, a structure that can be represented byOnlineEagle3Model.- The functionality for loading MLLM models is also placed in the
HFEagle3TargetModelclass. - Supports the
process_infofunction inpreprocessing.pyfor different MLLM models to obtain input features for image, video, and audio.
Testing of this design is still ongoing. I will try to add more support for the aforementioned MLLM models in this PR, and I very much hope to get guidance on whether the changes to this class abstraction are concise enough.
Checklist
- [ ] Format your code according to the Code Formatting with Pre-Commit.
- [ ] Add unit tests as outlined in the Running Unit Tests.
- [ ] Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
- [ ] Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
- [ ] For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
- [ ] Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.
Summary of Changes
Hello @KerwinKai, I'm Gemini Code Assist[^1]! 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 enhances the Eagle3 training framework by introducing a more abstract and generalized approach to supporting Multimodal Large Language Models (MLLMs). The changes consolidate MLLM-specific logic into common components, streamline the loading of various MLLM architectures, and improve data preprocessing capabilities to handle a broader range of multimodal inputs. This refactoring allows for easier integration and training of new MLLM models, moving beyond the previously limited support for Qwen2.5-VL.
Highlights
- Abstracted MLLM Training: The specialized
QwenVLOnlineEagle3Modelhas been removed, and the genericOnlineEagle3Modelnow uniformly handles the training of both Large Language Model (LLM) and Multimodal Large Language Model (MLLM) draft models by abstracting their input and output structures. - Centralized MLLM Loading: The
HFEagle3TargetModelclass has been enhanced to manage the loading of various MLLM architectures, including Qwen2.5-VL and Qwen3-Omni-Moe, providing a more unified approach to integrating different target models. - Enhanced Multimodal Data Preprocessing: The
preprocessing.pymodule now dynamically supports different MLLM processors, such asQwen2_5_VLProcessorandQwen3OmniMoeProcessor, allowing for flexible handling of diverse multimodal input features including images, videos, and audio. - New Dependency Added: The
qwen-omni-utilspackage has been added to the project's requirements, enabling support for the Qwen3-Omni-Moe model.
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 :thumbsup: and :thumbsdown: 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.
[^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.
Great job!!!!! This enables us to support all MLLM models. @FrankLeeeee
What is the current status of this PR now?
What is the current status of this PR now?
Hi, this PR is now ready for review. Meanwhile, under @sleepcoo 's guidance, we are also experimenting with training the draft model for Omni.
Hi. Thanks to your work. I've tried training draft for qwen2_5-vl with your work. I found that there is no places to get position_ids from target_model. get_rope_index will give 3D of position_ids, and it works well while trainning qwen vl model. If we don't do it, I guess there will be 2D of position_ids. However LlamaMutiRotaryEmbedding needs the 3D of position_ids.
position_ids, rope_deltas = self.target_model.model.get_rope_index(
input_ids,
image_grid_thw,
None,
second_per_grid_ts=None,
attention_mask=attention_mask_tensor,
)
Hi. Thanks to your work. I've tried training draft for qwen2_5-vl with your work. I found that there is no places to get position_ids from target_model. get_rope_index will give 3D of position_ids, and it works well while trainning qwen vl model. If we don't do it, I guess there will be 2D of position_ids. However
LlamaMutiRotaryEmbeddingneeds the 3D of position_ids.position_ids, rope_deltas = self.target_model.model.get_rope_index( input_ids, image_grid_thw, None, second_per_grid_ts=None, attention_mask=attention_mask_tensor, )
Hi, I haven't modified the training logic for Qwen2.5-VL in this PR, but I’ll run tests as soon as possible to identify the issue.
Does this work for other VLMs which are not Qwen?
there is conflict with the main branch.
there is conflict with the main branch.
I’ll resolve the merge conflicts as soon as possible and add support for DS-VL and InternVL based on the current architecture.
From the design of the Omni model (see modeling_qwen3_omni_moe.py#L3976 ), our Eagle-based acceleration can only apply to the first-stage thinker model’s generate step, as the talker model already incorporates its own MTP module.
Therefore, to effectively integrate Eagle into Omni’s inference pipeline within SGLang, we likely need to decouple the two-stage forward pass (thinker → talker). Given this, I recommend prioritizing the review of https://github.com/sgl-project/SpecForge/pull/251 first.