[quant] QuantizedCUDA

[alnfedorov]

🚀 Feature

1. Introduction of the QuantizedCUDA backend(similar to the QuantizedCPU backend).
2. Implementation of the PER_TENSOR_AFFINE QScheme for QuantizedCUDA tensors. (others can be implemented later on)

End users should be able to quantize(qint8, quint8, qint32) and dequantize CUDA tensors with float[16|32|64] dtype.

Motivation

General necessity:
The last generation of the GPU produced by Nvidia (Turing) includes support of the integer Tensor Cores (TC). There are several ways that we can make use of them besides cheap and fast inference:

1. In perspective, we could replace Fake-Quantization during training and work with a fully quantized forward pass to reduce possible training-inference biases and, perhaps, speed up things.
2. This could stimulate researchers to investigate the limits of the cheap and dirty, fully quantized training (backward + forward).

Anyway, the first step to make use of the Tensor Cores is to introduce quantized CUDA tensors.

My user-case:
Up to date, I implemented TVM modules for common layers in the nn (conv, pooling) with use of the tensor cores, both forward and backward pass. Unfortunately, it is not simple to use them with PyTorch framework:

1. Autograd makes a lot of assumptions(enforcing float dtype, etc) about gradients, so I had to remove several of them.
2. I have to store quantization parameters for integer tensors somewhere.

The first problem should be considered elsewhere (out of the PR scope).
Any hacks that came to my mind to solve the second problem didn't work (see alternatives), and it seems that I need a new backend. With a simple functionality to pass tensors around and quantize/dequantize them. We could add more support later.

Pitch

Here is the code:

import torch

t = torch.rand(10)
print(t)
# tensor([0.6088, 0.3496, 0.3973, 0.0884, 0.5340, 0.9819, 0.5057, 0.2072, 0.6677,
#         0.6197], device='cuda:0')
t = torch.quantize_per_tensor(t, 0.01, 0, torch.qint8)
print(t)
# tensor([0.6100, 0.3500, 0.4000, 0.0900, 0.5300, 0.9800, 0.5100, 0.2100, 0.6700,
#         0.6200], size=(10,), dtype=torch.qint8, device='cuda:0',
#        quantization_scheme=torch.per_tensor_affine, scale=0.01, zero_point=0)
t = t.dequantize()
print(t)
# tensor([0.6100, 0.3500, 0.4000, 0.0900, 0.5300, 0.9800, 0.5100, 0.2100, 0.6700,
#         0.6200], device='cuda:0')

Should work without errors.

Alternatives

If end-users wish to simulate QuantizedGPU backend, they should store tensor's quantization parameters somewhere. And there is no reliable way to do so, especially given autograd backward strategy. Tensor's hash is not reliable because different views on the same memory have different hashes. One can use data_ptr to simulate hash, yet it leads to errors when PyTorch uses cached memory allocation. I didn't come up with any other simple hack, looks like adding QuantizedGPU backend is more fruitful.

Additional context

I have started to dig into QuantizedCPU PRs; it looks like I will need to go the same way. However, I would highly appreciate any comments and suggestions. Ideally just a list of files to update to get things to work :).

Edits: more clarity.

cc @jerryzh168 @jianyuh @dzhulgakov @raghuramank100 @jamesr66a

[jerryzh168]

Thanks for the proposal! Yeah I think we are planning to add QuantizedCUDA backend, it's just there is no one working on this right now. I'll be happy to review the PR.

One thing to keep in mind is that we are only quantizing Tensor of dtype float32 right now, so it might take some effort to support quant/dequant of float16 and float64. I think it would be better to start with float32 and once we have the Quantizer API ready, we can think about extension to other dtypes, otherwise you'll need to add extra argument to quantize/dequantize function which would require touching all the code that uses these functions.

[alnfedorov]

Thanks for the feedback @jerryzh168! Sure, you are right. It is better to stick to the float32 right now and keep float16|64 for later work.

I familiarized with PyTorch internals a bit. Here is a WIP pull request. Can create dummy quantized tensors directly on the CUDA device. Given my current progress, I struggle with the Quantizers. I believe, it is a more or less general obstacle for the implementation of the new quantized backends and needs a good discussion.

Right now quantization is done in a virtual call, like quantizer->quantize(self). And it works nicely if we stick to a single backend, like CPU. However, here we generalize and need several levels of the 'dispatching': quantizer type -> tensor backend -> dtype. So, the problem is: How to organize this dispatching? I see several options:

• Create native functions like _per_tensor_affine_quantize(_per_channel_affine_quantize, _per_tensor_symmetric_quantize etc) and write specialized implementations for different backends. In the quantizer, we will simply call the corresponding native function. The main drawback here is that we will have to define quantize/dequantize functions for every quantizer. Users that implement custom Quantizer class with specialized implementations will have to do dispatching by hand.
• We might introduce dispatching into the base Quantizer class. Very simple, via map<DeviceType, function> or virtual methods. Implementations might live in different aten/quantized folders cpu/cuda/etc, like native functions. It's a bit ugly, as we reinvent the wheel but shout work. In this case, users should be able to write highly specialized quantizers if needed.

I am not aware of the internal discussions on the topic and will highly appreciate suggestions and other ideas. Neither of my thoughts seems good enough for me right now, but I will come back later with a fresh head.

[jerryzh168]

I think maybe we can use DispatchStub(https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/DispatchStub.h)
see examples of Copy:
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Copy.h
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Copy.cpp
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Copy.cu

cc @gchanan @smessmer is this common practice in aten?

[alnfedorov]

Also, to make sure that I use the right feature. Am I right to say that CUDA_tensor_applyX is going to be replaced by gpu_kernel in the aten?

[jerryzh168]

cc @VitalyFedyunin do you know?

[alnfedorov]

I have already managed to implement and test all the stuff locally. Waiting for the code style guidelines and will update PR.

[ggeor84]

Is there any progress on this?

[jerryzh168]

@alnfedorov can you link the PR?

[alnfedorov]

@ggeor84 @jerryzh168 Sure, will link and update PR at the end of the next week.

[alnfedorov]

@ggeor84 @jerryzh168. Sorry that it took this much time. I was unexpectedly busy with another matter. Please, see the linked PR for the current implementation, it compiles and successfully passes all tests.

@jerryzh168 Please, could you review it in the nearest time? I have unexpected free 1-2 weeks due to the global quarantine and will be able to fix any suggestions quickly. I had to make several decisions during development(that should be managed by internal team agreements that I am not aware of), so it might require some adjustments.