refine sparsevlm for llava #418
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 focuses on refining the SparseVLM token reduction mechanism, specifically optimizing its integration and performance within the LLaVA model. The changes introduce greater configurability for pruning and merging operations, enhance the handling of attention-related parameters across model layers, and improve the flexibility of token retention strategies. These updates aim to make SparseVLM more robust and adaptable for efficient processing within the LLaVA framework.
Highlights
- Enhanced SparseVLM Configuration: I've updated the
sparsevlm.ymlconfiguration to expand the defaultpruning_locto multiple layers ([2, 6, 15]) and explicitly introducedprune_flagandmerge_flagparameters, allowing for more granular control over SparseVLM's token reduction behavior. - Improved Token Reduction Logic for LLaVA: The core token reduction logic in
sparsevlm.pyhas been refined to correctly propagate and handleattention_maskandposition_embeddingsacross model layers. This ensures better compatibility and performance with LLaVA's attention mechanisms. A new conditional path was added to allow 'merge-only' operation when pruning is disabled. - Dynamic Sparse Token List Management: I've refactored the initialization of sparse token lists into a new
update_list()function. This function now dynamically sets the token retention values based on whether pruning, merging, or both are enabled, providing more adaptive token retention strategies for different scenarios. - Refined Hook Management and Parameter Propagation: The
update_kwargs_hooknow consistently stores critical model parameters (position_ids,attention_mask,cache_position,position_embeddings) for subsequent layers. Additionally, the hook registration forLlavamodels has been adjusted to apply at the block level, ensuring broader applicability of the token reduction logic.
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
The code changes refine the SparseVLM implementation for Llava models. The changes correctly propagate attention_mask and other parameters through the attention layers. However, the heavy reliance on global variables for configuration introduces maintainability and correctness risks and should be refactored to use instance-level state. A potential bug related to tensor dimension squeezing and some magic numbers should be addressed.
| prune_flag = True | ||
| merge_flag = True | ||
| sparse_token_list_192 = [] | ||
| sparse_token_list_128 = [] | ||
| sparse_token_list_64 = [] | ||
| sparse_token_dict = {} |
There was a problem hiding this comment.
The introduction of these module-level global variables can cause unexpected behavior, especially if multiple SparseVLM instances are created. Each instance could overwrite the global configuration, leading to race conditions or incorrect configurations. Consider encapsulating these variables as instance attributes of the SparseVLM class to ensure that each instance manages its own state.
| merge_sparse_token = cluster_and_merge(merge_token_stage2, cluster_num) | ||
|
|
||
| cluster_idx = total_sparse_token_idx.squeeze(0)[merge_token_stage2_idx[index_down]] | ||
| cluster_idx = cluster_idx.squeeze(0) |
There was a problem hiding this comment.
The tensor cluster_idx is 1-dimensional. If its size is 1 (i.e., shape is (1,)), squeeze(0) will convert it into a 0-dimensional scalar. This will cause an error in the torch.cat operation on line 293, which expects a sequence of tensors. Remove this line.
| def update_list(): | ||
| global sparse_token_list_192, sparse_token_list_128, sparse_token_list_64 | ||
| global prune_flag, merge_flag, sparse_token_dict | ||
|
|
||
| if layer_dict == {2: 0, 6: 1, 15: 2}: # 2*576 4*300 10*200 16*110 | ||
| sparse_token_list_192 = [300, 200, 110] | ||
| sparse_token_list_128 = [303, 110, 36] | ||
| sparse_token_list_64 = [66, 30, 17] | ||
| prune_flag, merge_flag = True, True | ||
| elif prune_flag and merge_flag: | ||
| sparse_token_list_192 = [180] | ||
| sparse_token_list_128 = [114] | ||
| sparse_token_list_64 = [48] | ||
| elif prune_flag: | ||
| sparse_token_list_192 = [192] | ||
| sparse_token_list_128 = [128] | ||
| sparse_token_list_64 = [64] | ||
| elif merge_flag: | ||
| sparse_token_list_192 = [149] | ||
| sparse_token_list_128 = [78] | ||
| sparse_token_list_64 = [7] | ||
| else: | ||
| raise RuntimeError( | ||
| 'Both prune_flag and merge_flag are False — sparseVLM is inactive.' | ||
| ) | ||
|
|
||
| sparse_token_dict = { | ||
| 192: sparse_token_list_192, | ||
| 128: sparse_token_list_128, | ||
| 64: sparse_token_list_64 | ||
| } | ||
| sparse_token_dict = { | ||
| 192: sparse_token_list_192, | ||
| 128: sparse_token_list_128, | ||
| 64: sparse_token_list_64 | ||
| } |
There was a problem hiding this comment.
This function heavily relies on and modifies global variables, contributing to maintainability issues. The condition if layer_dict == {2: 0, 6: 1, 15: 2}: is brittle. It hardcodes a specific configuration, making the code less flexible. Refactor this into a private method of SparseVLM that operates on instance attributes.
| else: | ||
| attn_logits = pruning_pars['attn_logits'] | ||
| prune_flag = pruning_pars.get('prune_flag', True) |
There was a problem hiding this comment.
For consistency, merge_flag should also be accessed using .get() to prevent potential KeyError exceptions, similar to how prune_flag is handled.
prune_flag = pruning_pars.get('prune_flag', True)|
|
||
| merge_token_idx_stage1 = torch.where(pred_score_vis == 0)[1] | ||
| merge_token_stage1 = relation_vis_text[0][merge_token_idx_stage1] | ||
| merge_token_num_stage1 = int(merge_token_idx_stage1.shape[0] * 0.3) + 1 # Top 30% | ||
| if prune_flag: |
There was a problem hiding this comment.
The value 0.3 is a magic number. Define it as a named constant to improve readability and maintainability.
| non_topk_idx = all_idx[~torch.isin(all_idx, merge_token_stage2_idx)] | ||
| pred_score_vis[0][non_topk_idx] = 1 | ||
| policy[:, v_token_start:text_token_start] = \ | ||
| pred_score_vis.type(dtype=hidden_states.dtype) | ||
|
|
||
| merge_token_stage2 = total_sparse_token[:, merge_token_stage2_idx, :] | ||
| cluster_num = int(merge_token_stage2.shape[1] / 10) + 1 |
There was a problem hiding this comment.
The value 10 is a magic number used to determine the number of clusters. Define it as a named constant to improve code clarity and maintainability.
| @@ -363,7 +388,7 @@ def read_parameter_hook(module, args, kwargs, pruning_pars): | |||
| with_kwargs=True | |||
| ) | |||
| elif self.model.__class__.__name__ == 'Llava': | |||
| self.blocks[block_idx].self_attn.register_forward_pre_hook( | |||
| self.blocks[block_idx].register_forward_pre_hook( | |||
There was a problem hiding this comment.
Consider registering the forward pre-hook directly to the block instead of the self_attn module. This simplifies the hook registration and ensures that the hook is applied to the entire block.
self.blocks[block_idx].register_forward_pre_hook(
functools.partial(
update_kwargs_hook,
pruning_pars=self.pruning_paras,
layer_idx=block_idx
),
with_kwargs=True
)
No description provided.