PyramidDrop and SparseVLM for llava #396
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @SmudgedWings, 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!
This pull request enhances the token reduction capabilities of PyramidDrop and SparseVLM by introducing explicit support for LLaVA models. It includes specialized handling for LLaVA's multimodal inputs through custom hooks and refactors internal parameter management for consistency across various token reduction modules.
Highlights
- LLaVA Model Compatibility: Added specific logic within PyramidDrop and SparseVLM to support different LLaVA model implementations (
LlavaHfandLlava), ensuring proper integration of token reduction techniques across variants. - Multimodal Input Handling: Introduced new
input_hook_llavafunctions in both PyramidDrop and SparseVLM. These hooks are designed to process LLaVA's multimodal inputs, correctly identifying and managing image and text tokens during the model's forward pass, particularly by interceptingprepare_inputs_labels_for_multimodal. - Parameter Management Refactor: Standardized the storage and access of pruning parameters across PyramidDrop, SparseVLM, and ToMe modules by migrating from
self.model.model.parametersto a more descriptiveself.pruning_parasvariable.
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 is currently in preview and 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 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 introduces PyramidDrop and SparseVLM token reduction techniques for Llava models. The changes primarily involve adding model-specific logic to handle differences between LlavaHf and Llava model architectures, especially around accessing the underlying language model and registering hooks. My review focuses on improving code maintainability by addressing repeated conditional logic, an incorrect type hint, and a leftover debug comment.
| output_hidden_states: Optional[bool] = None, | ||
| return_dict: Optional[bool] = None, | ||
| ) -> Union[Tuple, BaseModelOutput]: | ||
| ) -> Union[Tuple]: |
There was a problem hiding this comment.
The return type hint was changed to Union[Tuple], but the function can still return a BaseModelOutput instance (lines 86-90) if return_dict is true. This makes the type hint incorrect.
| ) -> Union[Tuple]: | |
| ) -> Union[Tuple, BaseModelOutput]: |
| if self.model.__class__.__name__ == 'LlavaHf': | ||
| llama_model = self.model.vlm_model.language_model.model | ||
| elif self.model.__class__.__name__ == 'Llava': | ||
| llama_model = self.model.vlm_model.model |
There was a problem hiding this comment.
This if/elif block to get the llama_model based on the model's class name is repeated later in this file inside pruning_hook (lines 223-226). This code duplication makes the code harder to maintain. Consider refactoring this logic into a helper method within the PyramidDrop class.
| if self.model.__class__.__name__ == 'LlavaHf': | ||
| self.model.embed_tokens.register_forward_pre_hook( | ||
| functools.partial(input_hook, pruning_pars=self.pruning_paras) | ||
| ) | ||
| elif self.model.__class__.__name__ == 'Llava': | ||
| from llava.constants import IMAGE_TOKEN_INDEX | ||
| hook_fn = input_hook_llava( | ||
| self.model.vlm_model.prepare_inputs_labels_for_multimodal, | ||
| self.pruning_paras | ||
| ) | ||
| self.model.vlm_model.prepare_inputs_labels_for_multimodal = MethodType( | ||
| hook_fn, self.model.vlm_model | ||
| ) |
There was a problem hiding this comment.
This if/elif block to register hooks based on the model type introduces logic specific to LlavaHf and Llava. A similar pattern is used in sparsevlm.py. To improve maintainability and reduce code duplication across different token reduction modules, consider abstracting this model-specific hook registration logic.
| v_t = hidden_states[:, v_token_start: text_token_start, :] | ||
| t_t = hidden_states[:, text_token_start:, :] | ||
| m_v_t = v_t @ t_t.transpose(1, 2) # [1, 576, 53] | ||
| m_v_t = v_t @ t_t.transpose(1, 2) # [1, 576, 53] # 52? |
There was a problem hiding this comment.
This comment # 52? appears to be a temporary debugging note and should be removed.
| m_v_t = v_t @ t_t.transpose(1, 2) # [1, 576, 53] # 52? | |
| m_v_t = v_t @ t_t.transpose(1, 2) # [1, 576, 53] |
| if self.model.__class__.__name__ == 'LlavaHf': | ||
| self.model.embed_tokens.register_forward_pre_hook( | ||
| functools.partial( | ||
| input_hook, | ||
| pruning_pars=self.pruning_paras | ||
| ) | ||
| ) | ||
| elif self.model.__class__.__name__ == 'Llava': | ||
| from llava.constants import IMAGE_TOKEN_INDEX | ||
| hook_fn = input_hook_llava( | ||
| self.model.vlm_model.prepare_inputs_labels_for_multimodal, | ||
| self.pruning_paras | ||
| ) | ||
| self.model.vlm_model.prepare_inputs_labels_for_multimodal = MethodType( | ||
| hook_fn, self.model.vlm_model | ||
| ) | ||
| ) | ||
|
|
||
| self.model.model.register_forward_pre_hook( | ||
| if self.model.__class__.__name__ == 'LlavaHf': | ||
| llama_model = self.model.model | ||
| elif self.model.__class__.__name__ == 'Llava': | ||
| llama_model = self.model.model.model | ||
| llama_model.register_forward_pre_hook( | ||
| functools.partial( | ||
| register_module_pars, | ||
| pruning_pars=self.model.model.parameters), | ||
| pruning_pars=self.pruning_paras), | ||
| with_kwargs=True | ||
| ) |
There was a problem hiding this comment.
There are two consecutive if/elif blocks here that check self.model.__class__.__name__. This introduces duplicated conditional logic and makes the code harder to read and maintain. This logic is also very similar to what's in llmc/compression/token_reduction/pyramiddrop.py. Consider determining the model type once and storing it as a boolean flag or refactoring the logic to get llama_model into a helper method.
2 participants
No description provided.