torch.ao.quantization.observer work on CPU (very slow!) due to using registered buffers on CPU
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.6k
- PR merge metrics
- PR metrics pending
Description
### 🚀 The feature, motivation and pitch
This is almost a BUG. The observers for torch.ao.quantization do their work on CPU, for an add reason.
* `UniformQuantizationObserverBase` registers `self.eps` on CPU. Then, code in `UniformQuantizationObserverBase._calculate_qparams` fails if its args are on GPU
* `PerChannelMinMaxObserver` registers `self.min_val`, `self.max_val` on CPU. Then, argument to `forward` is moved to CPU
I noted this by code becoming very slow whenever args are on GPU. I need quantization as part of activation checkpointing, why would I move everything to CPU?
### Alternatives
Don't register `self.eps` or `self.min_val`, `self.max_val` as buffers. I don't understand why this is needed, since quantization is something I do on the fly. If you insist, allow for a device argument, or (better) move the registered buffer to GPU with the first `forward` call, when it is still empty.
### Additional context
_No response_
cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel
Contributor guide
Assessment
This issue has not been assessed yet.