fake_quant: add a more memory efficient backward

[vkuzo]

Stack from ghstack:

• fake_quant: more memory efficient per-channel backward #51255 fake_quant: more memory efficient per-channel backward
• memory efficient fq: use it everywhere, delete the old version #51159 memory efficient fq: use it everywhere, delete the old version
• fake_quant: add a more memory efficient backward #50561 fake_quant: add a more memory efficient backward

Summary:

tl;dr; add an alternative implementation of fake_quantize which saves
a mask of whether the input was clamped during the forward pass and uses it to calculate the backward. The math:

# before - forward (pseudocode)
def fq_forward(x, scale, zp, qmin, qmax):
    q_val = clamp(nearby_int(x / scale) + zp, qmin, qmax)
    fq_val = (q_val - zp) * scale
    return fq_val

# before - backward (pseudocode)
def fq_backward(dy, x, scale, zp, qmin, qmax):
    q_val_unclamped = nearby_int(x / scale) + zp
    mask = qmin <= q_val_unclamped and q_val_unclamped <= qmax
    return dy * mask

# after - forward (pseudocode)
def fq_forward(x, scale, zp, qmin, qmax):
    q_val_unclamped = nearby_int(x / scale) + zp
    mask = qmin <= q_val_unclamped and q_val_unclamped <= qmax
    q_val = clamp(q_val_unclamped, qmin, qmax)
    fq_val = (q_val - zp) * scale
    return fq_val, mask

# after - backward (pseudocode)
def fq_backward(dy, mask):
    return dy * mask

This way the backward function no longer needs the input Tensor, and it can be
gc'ed earlier by autograd. Instead of passing x: FloatTensor, we pass a mask: BoolTensor
with the same number of elements. BoolTensor uses 1 byte per element,
so we expect an upper bound of a 75% memory overhead reduction. We observe a 73% memory
overhead reduction on torchvision's MobileNetV2 in real world tests. Packing the bools
into a custom storage format to take 1 bit per element is an optimization left for the future.

Performance impact of this seems negligible, I observed a 1% to 5% regression on MobileNetV2 but it's unclear if it's real.

Adding this as a new function (as opposed to replacing the old implementation) for easy testing, but
might be worth deleting the old fake_quant backward in a future PR. We can adjust the signature
of this function to take model.training as an additional parameter, and skip the mask computation for eval.

Test Plan:

QAT on MobileNetV2 on FB infra, with opt build flags, batch_size = 32. Results for fbgemm settings, qnnpack results are similar.

# qat_fp32: model with fake_quants turned off (baseline)
# qat_1: step 2 of qat, with observers disabled and fake_quants enabled (all of the overhead is the fake_quants)

# before: fbgemm - qat_fp32 -> qat_1
max memory usage (mib): 3299 -> 4170 (overhead: 26.4%)
latency (ms):  147 -> 181

# after: fbgemm - qat_fp32 -> qat_1
max memory usage (mib): 3302 -> 3528 (overhead: 7.1%)
latency (ms):  147 -> 183

Note: similar metrics are observed in an OSS / torchvision / MobileNetV2 setup, with this command:

python references/classification/train_quantization.py
  --print-freq 1
  --data-path /data/local/packages/ai-group.imagenet-256-smallest-side/prod/
  --output-dir ~/nfs/pytorch_vision_tests/
  --backend qnnpack
  --epochs 5

All CI tests here: #50849

PyTorch microbenchmarks (CUDA performance about the same:

cd benchmarks/operator_benchmark
python -m pt.quantization_test

results: https://gist.github.com/vkuzo/11a7bed73fe60e340862d37e7975e9cd)

Unit tests:

python test/test_quantization.py TestFakeQuantize

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: D25918519

[ngimel]

The kernels look good, I left minor comments.

[ngimel]

You can log sizes here to make error message more informative

[ngimel]

It's unusual to see kernels in pytorch that are hardcoded for float only. Could people want to quantize half or bfloat16 models?

[codecov]

Codecov Report

Merging #50561 (194be25) into gh/vkuzo/209/base (5ec2e26) will decrease coverage by 0.00%.
The diff coverage is 100.00%.

@@                  Coverage Diff                  @@
##           gh/vkuzo/209/base   #50561      +/-   ##
=====================================================
- Coverage              80.88%   80.88%   -0.01%     
=====================================================
  Files                   1931     1931              
  Lines                 210588   210604      +16     
=====================================================
+ Hits                  170339   170343       +4     
- Misses                 40249    40261      +12     

[facebook-github-bot]

This pull request has been merged in 983b8e6.