Direct FBGEMM integraton into ATen (#13777)
authorJames Reed <jamesreed@fb.com>
Fri, 21 Dec 2018 18:32:57 +0000 (10:32 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Fri, 21 Dec 2018 18:35:51 +0000 (10:35 -0800)
commitacbd9c49b0f3414b2a161731b0b46aa0712d78f7
tree289055c244893e4eee3f5d8ca9e9241ab4565585
parent614121c1ef6a34e7653f43cb9be89c28d6766b19
Direct FBGEMM integraton into ATen (#13777)

Summary:
This PR implements infrastructure for post-processing a model to apply int8 quantization to its `nn.Linear` modules. Highlights of the implementation:

1) Inputs and outputs are `float` (quantized and packed internally), but the weight is quantized and packed ahead of time for efficiency. This implementation performs well in small-batch size GEMM calls. It should not be considered a general-purpose quantized GEMM kernel.
2) Weight packing is dependent on machine architecture (e.g. vector register width), so it is done just-in-time. Concretely, it is done on model load for the weights and it is done during operator execution for the input value.
3) Biases are unquantized
4) We fail loudly if we are attempting to run this on a machine that does not support FBGEMM. This is because we do not want a model's numerics to differ based on which machine it is run on. A model containing these FBGEMM ops *must* be run with FBGEMM

The API can be seen in the added test case. Highlights are:
1) `torch.jit.quantized.quantize_linear_modules` walks the module hierarchy of the passed-in Module and replaces all `nn.Linear` modules with a new `QuantizedLinear` module, which encapsulates the behavior described above.
2) `_pack()` and `_unpack()` script methods are present on `QuantizedLinear` modules. These methods should be called before serialization and after deserialization, respectively. This ensures that the weight matrix is properly packed for the running machine's architecture. Note that in the long term, we would like to move toward a more Pickle-style serialization technique, rather than having these explicit methods that mutate member values. This is blocked on being able to assign attributes in a ScriptMethod, among other things.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/13777

Differential Revision: D13383276

Pulled By: jamesr66a

fbshipit-source-id: 00f29c9f34544add2b90107e3cf55a287802c344
CMakeLists.txt
aten/src/ATen/native/QuantizedLinear.cpp [new file with mode: 0644]
aten/src/ATen/native/native_functions.yaml
cmake/Dependencies.cmake
test/test_jit.py
tools/autograd/gen_python_functions.py
torch/csrc/autograd/utils/wrap_outputs.h
torch/jit/quantized.py [new file with mode: 0644]
ubsan.supp