Lightning-AI / Lightning-AI/pytorch-lightning

ignore_modules in Quantization via Bitsandbytes

Open
#19,175 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature help wanted precision: bnb
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.