Lightning-AI / Lightning-AI/pytorch-lightning
ignore_modules in Quantization via Bitsandbytes
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Currently, the strings provided in `ignore_modules` will be checked in as a [prefix](https://github.com/Lightning-AI/pytorch-lightning/blob/119039b932c79f1e5daef67345628e675ab3c0bf/src/lightning/fabric/plugins/precision/bitsandbytes.py#L281). This will cause problems when someone wants to quantise specific module and ignore some of the modules. Like in the following module list:
```
llama.visual_backbone.visual.transformer.resblocks.7.mlp.c_fc
llama.visual_backbone.visual.transformer.resblocks.7.mlp.gelu
llama.visual_backbone.visual.transformer.resblocks.7.mlp.c_proj
llama.visual_backbone.visual.transformer.resblocks.7.ln_2
llama.visual_backbone.visual.transformer.resblocks.7.adapter_attn
llama.visual_backbone.visual.transformer.resblocks.7.adapter_attn.conv_A
llama.visual_backbone.visual.transformer.resblocks.7.adapter_attn.conv_B
llama.visual_backbone.visual.transformer.resblocks.8
llama.visual_backbone.visual.transformer.resblocks.8.attn
llama.visual_backbone.visual.transformer.resblocks.8.attn.out_proj
llama.visual_backbone.visual.transformer.resblocks.8.ln_1
llama.visual_backbone.visual.transformer.resblocks.8.mlp
llama.visual_backbone.visual.transformer.resblocks.8.mlp.c_fc
llama.visual_backbone.visual.transformer.resblocks.8.mlp.gelu
llama.visual_backbone.visual.transformer.resblocks.8.mlp.c_proj
llama.visual_backbone.visual.transformer.resblocks.8.ln_2
llama.visual_backbone.visual.transformer.resblocks.8.adapter_attn
llama.visual_backbone.visual.transformer.resblocks.8.adapter_attn.conv_A
llama.visual_backbone.visual.transformer.resblocks.8.adapter_attn.conv_B
```
there are many layers of `resblocks`, when I want to quantize all layers other than `conv_A` and `conv_B`, I can't do that because in the `ignore_modules` I've to provide complete path to `resblocks` and since there are many layers, it's simply not possible.
### Pitch
What I suggest is to instead of checking prefix, one should search for module while ignoring, like:
```
if isinstance(child, torch.nn.Linear) and not any(s in fullname for s in ignore_modules):
```
### Alternatives
_No response_
### Additional context
_No response_
cc @borda @carmocca @awaelchli
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/lightning/fabric/plugins/precision/bitsandbytes.py around the ignore_modules check at line 281, and inspect how module full names are matched during quantization. Verify the intended matching behavior against the module paths in the issue and add or update relevant coverage so ignored modules are left unquantized.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 38/100