SpecForge icon indicating copy to clipboard operation
SpecForge copied to clipboard

Add Eagle3 training for more MLLM model

Open KerwinKai opened this issue 1 month ago • 9 comments

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

  1. QwenVLOnlineEagle3Model was 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 by OnlineEagle3Model.
  2. The functionality for loading MLLM models is also placed in the HFEagle3TargetModel class.
  3. Supports the process_info function in preprocessing.py for 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.

KerwinKai avatar Nov 15 '25 09:11 KerwinKai

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 QwenVLOnlineEagle3Model has been removed, and the generic OnlineEagle3Model now 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 HFEagle3TargetModel class 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.py module now dynamically supports different MLLM processors, such as Qwen2_5_VLProcessor and Qwen3OmniMoeProcessor, allowing for flexible handling of diverse multimodal input features including images, videos, and audio.
  • New Dependency Added: The qwen-omni-utils package 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.

gemini-code-assist[bot] avatar Nov 15 '25 09:11 gemini-code-assist[bot]

Great job!!!!! This enables us to support all MLLM models. @FrankLeeeee

sleepcoo avatar Nov 17 '25 07:11 sleepcoo

What is the current status of this PR now?

FrankLeeeee avatar Nov 21 '25 06:11 FrankLeeeee

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.

KerwinKai avatar Nov 23 '25 12:11 KerwinKai

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,
            )

justadogistaken avatar Nov 23 '25 17:11 justadogistaken

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, 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.

KerwinKai avatar Nov 23 '25 17:11 KerwinKai

Does this work for other VLMs which are not Qwen?

FrankLeeeee avatar Nov 24 '25 05:11 FrankLeeeee

there is conflict with the main branch.

FrankLeeeee avatar Nov 24 '25 13:11 FrankLeeeee

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.

KerwinKai avatar Nov 24 '25 17:11 KerwinKai