pytorch / pytorch/pytorch

torch.ao.quantization.observer work on CPU (very slow!) due to using registered buffers on CPU

Open
#172,082 3 comments 0 reactions 0 assignees View on GitHub
oncall: quantization
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.