model: adapt mllama4 to VisionAttention #8512
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @wenchen76, I'm Gemini Code Assist1! 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!
I've implemented changes to the mllama4 model to introduce comprehensive support for Vision attention. This work directly addresses issue #8487, aiming to enable the model to process and understand visual inputs by integrating a dedicated vision pipeline and associated components.
Highlights
- Vision Model Integration: I've integrated a comprehensive vision model pipeline, including components for patch embedding, positional embeddings, and a dedicated vision encoder, to enable the model to process visual inputs.
- Custom Multi-Modal Projector: I've re-implemented the multi-modal projector within the codebase to specifically handle the projection of vision features into the text model's hidden space, ensuring compatibility and efficient data flow.
- Modular Vision Components: I've introduced several new modular classes, such as
Llama4VisionMLP,Llama4VisionEncoderLayer,Llama4VisionEncoder,Llama4UnfoldConvolution, andLlama4VisionPixelShuffleMLP, to build out the necessary vision processing capabilities. - Quantization and Parallelism Support: The newly added vision components are designed to support quantization configurations and various data parallelism strategies (e.g.,
ColumnParallelLinear,RowParallelLinear,ReplicatedLinear) for optimized and efficient execution.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 or fill out our survey to provide feedback.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request adapts mllama4 to support Vision attention by reimplementing vision components from transformers within sglang. The changes are extensive and introduce new classes for the vision pipeline. My review focuses on ensuring the new components are correctly integrated and configured. I've identified a critical issue with an incompatible API change that will cause a runtime error, a high-severity bug related to missing arguments during instantiation, and a couple of medium-severity issues regarding code clarity and potential bugs with data parallelism. The PR is marked as WIP, so these findings should help in finalizing the implementation.
d443cf8 to
a138126
Compare
|
server command: peak mem is tested with a request containing 3 images
|
0f9ec92 to
4fd9a4e
Compare
Co-authored-by: root <mickjagger19@icloud.com>
Co-authored-by: root <mickjagger19@icloud.com>
Co-authored-by: root <mickjagger19@icloud.com>
Co-authored-by: root <mickjagger19@icloud.com>
Co-authored-by: root <mickjagger19@icloud.com>
Co-authored-by: root <mickjagger19@icloud.com>
Motivation
To resolve #8487 and close #8468, adapt mllama4 to support Vision attention
Modifications
Accuracy Test
WIP
Benchmark & Profiling
WIP
Checklist