fix quant grad function calculation error #3160

eedalong · 2020-12-06T08:13:02Z

Quant grad calculation is not correct in current nni's implementation according to google's whitepaper, it should be

Cjkkkk · 2020-12-07T08:35:33Z

nni/compression/pytorch/compressor.py

@@ -608,24 +626,35 @@ def quant_backward(tensor, grad_output, quant_type):
        tensor
            gradient of the input of quantization operation
        """
-        return grad_output
+        tensor_q = QuantGrad._quantize(tensor, scale, zero_point)
+        mask = (tensor_q >= qmin) * (tensor_q <= qmax)


mask = (tensor_q < qmin) | (tensor_q > qmax) grad_output[mask] = 0 return grad_output

Cjkkkk · 2020-12-07T08:42:07Z

nni/compression/pytorch/compressor.py

        else:
            raise ValueError("unrecognized QuantType.")

+        bits = QuantGrad.get_bits_length(wrapper.config, quant_type)
+        qmin, qmax = torch.Tensor([0], device=tensor.device), torch.Tensor([(1 << bits) - 1], device=tensor.device)
+        ctx.save_for_backward(tensor, wrapper.module.scale, wrapper.module.zero_point, qmin, qmax)


you only need save mask here if you want to do clip gradient, which is a bytetensor and therefore 4 times smaller than saving a floattensor.

I think the reason why we used to use a pass-through-estimator is that we always set rmin, rmax to be the min(tensor), max(tensor). But your propose is correct if rmin is larger than the min(tensor) and rmax is smaller than the max(tensor).

yes， for weight， rmin and rmax is always the min and max of the tensor, but think about output. Because we use ema for updating rmin and rmax, so not all item in output is in [rmin, rmax]

Yes, your propose is necessary in quantizing ouput, thanks for submitting this PR!

@eedalong @Cjkkkk , thanks for the discussion. The design of this pr looks reasonable. let's keep it in this pr. We will do some refactor later to put forward logic and backward logic together within the same class, which should be much easier for users to customize new quantization algorithms.

Cjkkkk · 2020-12-07T08:49:53Z

nni/compression/pytorch/compressor.py

@@ -608,24 +626,35 @@ def quant_backward(tensor, grad_output, quant_type):
        tensor
            gradient of the input of quantization operation
        """
-        return grad_output
+        tensor_q = QuantGrad._quantize(tensor, scale, zero_point)


No need to calculate tensor_q here, return a mask from wrapper.quantizer.quantize_input for example.

No, we need tensor_q because we need unclampped quantized data, wrapper.quantizer.quantize_input is clampped to [qmin, qmax], so the mask will always be ones everywhere if we just use wrapper.quantizer.quantize_input

I mean you can do this inside the wrapper.quantizer.quantize_input and return the mask which can be used in backward.

Therefore you don't do quantization twice and tensor saved will be smaller by 4 times.

Yeah, u are right. I did write the code like that at first.. For doing this ,we need to modify _quantize() to return mask. But i gave it up because i dont think it's a good design. Just let _quantize() to return quantized tensor is a better design i think.

True, but keeping this design will actually make many quantization algorithms with special backward function harder and less efficient to implement. The actual problem is quant backward is inside the autograd while the forward is inside the quantizer module. I think it will be much better if forward and backward logic is moved into autograd function together and quantizer only maintain some states.

What I said is actually what Pytorch implementation does.

Yes u r right. And this calls for a HUGE design change on NNI's QAT implementation.

for now, i think i will just push the correct STE code and what u mentioned is left for next pr~

eedalong · 2020-12-07T08:56:44Z

check pytorch's implmentation guys:

back_grad for fake_quantize 's implementation in cuda in pytorch: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/quantized/cuda/fake_quantize_core.cu
back_forward for fake_quantize's implementation in python in pytorch
https://github.com/pytorch/pytorch/blob/master/torch/quantization/_learnable_fake_quantize.py

the key point is this picture

J-shang · 2020-12-07T09:24:53Z

nni/compression/pytorch/compressor.py

@@ -608,24 +626,35 @@ def quant_backward(tensor, grad_output, quant_type):
        tensor
            gradient of the input of quantization operation
        """


Comment of this function need to be updated or delete.

oh yeah, forgot to update the comment

linbinskn · 2020-12-07T10:39:50Z

nni/compression/pytorch/compressor.py

        return grad_output

    @staticmethod
    def forward(ctx, tensor, quant_type, wrapper, **kwargs):
-        ctx.save_for_backward(tensor, torch.Tensor([quant_type]))
        if quant_type == QuantType.QUANT_INPUT:


Suggest changing defination of QuantType to 'input', 'weight' and 'output' and we don't need to modify quant_type in the if-else clause.

Um.., I dont think it is relavant with this PR. May be you need to update all definition for quant type in another pr

could you help update line 583-585 to update the values from 0, 1, 2 to 'input', 'weight', 'output'? we have checked the code, it is safe to make this modification :)

QuanluZhang · 2020-12-07T10:40:30Z

nni/compression/pytorch/compressor.py

+        get bit for quantize config
+        :param config:
+        :param quant_type:
+        :return:


please update docstring format to be consistent with others

linbinskn · 2020-12-07T10:41:06Z

nni/compression/pytorch/compressor.py

@@ -589,8 +589,26 @@ class QuantGrad(torch.autograd.Function):
    """
    Base class for overriding backward function of quantization operation.
    """
+    @classmethod
+    def _quantize(cls, x, scale, zp):


Suggest changing zp to zero_point to keep consensus.

yeah, changed already

fix quant grad function

e8e46c4

QuanluZhang requested review from Cjkkkk, linbinskn and J-shang December 7, 2020 07:49

Cjkkkk reviewed Dec 7, 2020

View reviewed changes

optimize grad computation

39a26ef

eedalong requested a review from Cjkkkk December 7, 2020 09:19

J-shang reviewed Dec 7, 2020

View reviewed changes

linbinskn reviewed Dec 7, 2020

View reviewed changes

QuanluZhang reviewed Dec 7, 2020

View reviewed changes

linbinskn reviewed Dec 7, 2020

View reviewed changes

yuanxiulong added 3 commits December 7, 2020 18:48

update function comment

049d787

lint all comments

f1bf04e

meaningless fix, just for reruning pipeline test

ba40b5b

QuanluZhang approved these changes Dec 8, 2020

View reviewed changes

Cjkkkk approved these changes Dec 8, 2020

View reviewed changes

change QUANT_TYPE definition

6702284

J-shang approved these changes Dec 8, 2020

View reviewed changes

linbinskn approved these changes Dec 8, 2020

View reviewed changes

QuanluZhang merged commit 2f6a74f into microsoft:master Dec 8, 2020

This was referenced Dec 19, 2020

fix QAT ema issue and tensor type error #3211

Closed

fix QAT ema issue and tensor type error #3219

Merged

liuzhe-lz mentioned this pull request Dec 25, 2020

v2.0 Release Plan #2935

Closed

77 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix quant grad function calculation error #3160

fix quant grad function calculation error #3160

eedalong commented Dec 6, 2020

Cjkkkk Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020

Cjkkkk Dec 7, 2020

Cjkkkk Dec 7, 2020

eedalong Dec 7, 2020

Cjkkkk Dec 7, 2020

QuanluZhang Dec 8, 2020

Cjkkkk Dec 7, 2020

eedalong Dec 7, 2020 •

edited

Loading

Cjkkkk Dec 7, 2020

Cjkkkk Dec 7, 2020

eedalong Dec 7, 2020 •

edited

Loading

Cjkkkk Dec 7, 2020

Cjkkkk Dec 7, 2020

eedalong Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020

eedalong commented Dec 7, 2020 •

edited

Loading

J-shang Dec 7, 2020

eedalong Dec 7, 2020 •

edited

Loading

linbinskn Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020

QuanluZhang Dec 8, 2020

eedalong Dec 8, 2020

QuanluZhang Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020

linbinskn Dec 7, 2020

eedalong Dec 7, 2020

fix quant grad function calculation error #3160

fix quant grad function calculation error #3160

Conversation

eedalong commented Dec 6, 2020

Cjkkkk Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eedalong Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eedalong Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eedalong Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eedalong commented Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

eedalong Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

linbinskn Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

QuanluZhang Dec 7, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Cjkkkk Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020 •

edited

Loading

eedalong commented Dec 7, 2020 •

edited

Loading

eedalong Dec 7, 2020 •

edited

Loading

linbinskn Dec 7, 2020 •

edited

Loading

QuanluZhang Dec 7, 2020 •

edited

Loading