Lightning-AI / Lightning-AI/lightning-thunder
quantization: process tensors on meta device directly, maybe implement CPU quantization (if it is easy)
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Currently the BitsAndBytesLinearQuant4bit for submodule always calls `bitsandbytes.functional.quantize_4bit`. This is somewhat touchy for CPU tensors because `quantize_4bit` only works on GPU tensors but it is outright not so nice for meta tensors, where we only would need to get the right shapes.
https://github.com/Lightning-AI/lightning-thunder/blob/e64d347def39bb47101efafe4177adf9f77a63ec/thunder/transforms/quantization.py#L93-L103
Contributor guide
No contributing guide indexed for this repository
Research direction
Inspect thunder/transforms/quantization.py around lines 93-103 and trace how BitsAndBytesLinearQuant4bit handles meta and CPU tensors. Confirm the expected tensor shapes and the behavior of bitsandbytes.functional.quantize_4bit; done means meta tensors no longer require GPU quantization, with CPU behavior addressed if supported.
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
- Mostly clear
- Newbie friendliness
- 45/100