update mustdrop #414
update mustdrop #414
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 @chengtao-lv, 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 refines the mustdrop token reduction mechanism by introducing greater configurability and control over token retention, especially for vision models. It enables dynamic specification of the pruning layer and allows for a fixed number of tokens to be retained during spatial merging, enhancing the precision and flexibility of model compression.
Highlights
- Configurable Pruning Location: The
MustDropclass now initializesself.pruning_locfromself.special_config, allowing the specific layer where token reduction is applied to be dynamically configured. - Fixed Token Retention in Pooling: The
conditional_poolingfunction has been updated to accept afix_rparameter. Whenfix_ris non-zero, it overrides the dynamically calculatedrvalue, enabling a precise, predefined number of tokens to be retained during the pooling process. - Dynamic Token Retention Calculation: The
spatial_merge_hooknow calculates thefix_rvalue based onretained_tokensandvision_token_lengthprovided inpruning_paras. This allows the system to programmatically determine and enforce the exact number of tokens to retain, particularly useful for vision models. - Hook Mechanism Update: The token reduction hook registration has been changed from a
register_forward_pre_hookto aregister_forward_hook. This means the hook now processes the output of the specified block after its forward pass. Additionally, the target block for the hook is now dynamically determined byself.pruning_loc, making the hook application more flexible.
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 pull request updates the MustDrop token reduction module, making the pruning location configurable, introducing an option to set a fixed number of tokens to merge (fix_r), and changing the reduction mechanism to a forward_hook. The review identified a critical issue with the forward_hook modifying the module's output, a high severity issue with a potential ZeroDivisionError, and a medium severity issue with configuration parameter handling.
| merge = conditional_pooling(hidden_states, spatial_threshold, window_size, fix_r) | ||
| hidden_states, size = merge_wavg(merge, hidden_states, None) | ||
| return (hidden_states,) + args[1:], kwargs | ||
| return (hidden_states,) |
There was a problem hiding this comment.
The spatial_merge_hook is registered as a forward_hook. Returning only (hidden_states,) replaces the block's entire output, potentially discarding other outputs like attention weights, which can cause errors in subsequent layers. Ensure the hook preserves other elements of the output tuple.
return (hidden_states,) + layer_outs[1:]| fix_r = (pruning_paras['vision_token_length'] - retained_tokens) \ | ||
| // (window_size[0] * window_size[1] - 1) |
There was a problem hiding this comment.
The denominator window_size[0] * window_size[1] - 1 in the fix_r calculation can be zero if window_size is (1, 1), leading to a ZeroDivisionError. Add a check to prevent this.
if (window_size[0] * window_size[1] - 1) > 0:
fix_r = (pruning_paras['vision_token_length'] - retained_tokens) // (window_size[0] * window_size[1] - 1)
else:
fix_r = 0| self.register_reduction_modules() | ||
|
|
||
| def add_sparse_config(self): | ||
| self.pruning_loc = self.special_config['pruning_loc'] |
There was a problem hiding this comment.
Directly accessing self.special_config['pruning_loc'] can raise a KeyError if the key is missing. Consider adding a check to ensure the key exists before accessing it to prevent potential crashes.
No description provided.