[WIP] 8-bit quantization for inference

[kpuatamazon]

Add support for 8-bit quantized matrix multiplication in inference.

This code depends on the intgemm branch in my fork of MXNet: https://github.com/kpuatamazon/incubator-mxnet/tree/intgemm . This will turn into a pull request against MXNet.

Quantization on one thread runs 2.95x as fast as the baseline on one thread. Quantization on one thread is 1.28x as fast as the baseline on four threads. Results on an AWS c5.x12large.

Baseline 1 thread:
export MXNET_ENGINE_TYPE=NaiveEngine
export OMP_NUM_THREADS=1
[INFO:__main__] Processed 2969 lines. Total time: 934.8082, sec/sent: 0.3149, sent/sec: 3.1761
real    15m38.468s
user    15m37.986s
sys    0m5.215s

Baseline 4 threads:
export MXNET_ENGINE_TYPE=NaiveEngine
export OMP_NUM_THREADS=4
[INFO:__main__] Processed 2969 lines. Total time: 406.0341, sec/sent: 0.1368, sent/sec: 7.3122
real	6m48.662s
user	26m46.845s
sys	0m28.490s

Quantized 1 thread:
export MXNET_ENGINE_TYPE=NaiveEngine
export OMP_NUM_THREADS=1
[INFO:__main__] Processed 2969 lines. Total time: 314.9344, sec/sent: 0.1061, sent/sec: 9.4274
real    5m17.408s
user    5m17.297s
sys    0m4.476s

BLEU: 42.6 quantized, 42.5 baseline float32. No significant change.

Note that the on-disk format of the int8 file is dependent on the CPU architecture. A fix for this is pending a change to intgemm to separate the quantization and rearrangement steps.

The model is converted to 8-bit offline using a program. Here's a program to convert a model from fp32 to int8. You should also change the config file's dtype to int8. I'm soliciting suggestions on how to make this cleanly, probably another command-line program.

#!/usr/bin/env python3
import mxnet as mx
model = mx.nd.load("params.best")
#Find all weight tensors except for the source embeddings.  This includes output_layer.weight.
dense = [k[0:-7] for k in model.keys() if k.endswith('.weight') and not k.startswith("embedding_source.")]
#The positional embeddings are not quantized yet.
dense.remove("encoder.pos_embedding")
dense.remove("decoder.pos_embedding")
for param in dense:
  name = param + ".weight"
  b = model[name]
  b_scale = 127.0 / mx.nd.contrib.intgemm_maxabsolute(b)
  b_prepared = mx.nd.contrib.intgemm_prepare_weight(b, multiplier = b_scale.asscalar())
  model[name] = b_prepared
  model[param + ".scaling"] = 1.0 / b_scale
mx.nd.save("params.best.quant", model)

Pull Request Checklist

• Changes are complete (if posting work-in-progress code, prefix your pull request title with '[WIP]'
  until you can check this box.
• Unit tests pass (pytest)
• Were system tests modified? If so did you run these at least 5 times to account for the variation across runs? (Not modified)
• System tests pass (pytest test/system)
• Passed code style checking (./style-check.sh)
• You have considered writing a test
• Updated major/minor version in sockeye/__init__.py. Major version bump if this is a backwards incompatible change.
• Updated CHANGELOG.md

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[kpuatamazon]

Updated:

• CPU-independent disk format.
• If you've got a float32 model, just add --dtype int8 on the command line and it will quantize on the fly.
• To quantize and save to disk, do this:

import sockeye.model
model = sockeye.model.load_model("/path/to/float32_model", dtype = 'int8', for_disk_saving = True)
a = model[0]
a.save_parameters("/path/to/int8_model/params.best")
a.save_config("/path/to/int8_model")

You'll also need to ln -s /path/to/float32_model/{version,*.json} /path/to/int8_model/. Opinions welcome on a nice command-line version of this. Should it just copy the vocabs?

[fhieber]

Really looking forward to the corresponding mxnet change to get this merged!
Left a few comments, mostly minor style comments.

I think it would be nice to test in8 quantization in the system tests. This would entail quantizing the model in the test suite and have another decoding pass which allows you to assert on output similarity and/or BLEU. This would also clarify the workflow with int8 quantization.

[fhieber]

Couldn't you inherit from mx.gluon.nn.basic_layers.Dense directly and only overwrite cast() and hybrid_forward?

[kpuatamazon]

I agree. Tried to do this, will need a consultation with a gluon expert.

[tdomhan]

what was the issue?

[fhieber]

I guess we need to carefully set the prefix for the inheriting class to make sure the parameter names match.

[davvil]

We could potentially quantize from FP16, right? Or everything on disk is FP32?

[kpuatamazon]

There isn't a kernel to quantize from FP16 to INT8. CPUs aren't so great at FP16 anyway; they only have instructions to convert to/from FP32 then do all the math in FP32.

[fhieber]

So this means that being able to quantize to int8 for inference requires having trained an FP32 model?

[kpuatamazon]

Do you have stable training in FP16? I guess I could add a code path to convert FP16 -> FP32 -> int8 which, sadly, is how the CPU would do it anyway.

[kpuatamazon]

Now supports three disk formats:

1. Regular sockeye float32. Can be used for float32 inference (obviously) and this is the default. Run int8 with --dtype int8
2. float32 + scaling factors. Can be used for float32 inference and this is the default. Default: float32. Can run int8 with --dtype int8
3. int8 + scaling factors. Cannot be used for float32 inference (it won't wastefully reconstitute the matrices). Default: int8.

Adding scaling factors (transition 1 -> 2):

import sockeye.model
model = sockeye.model.load_model("model", for_disk_saving='float32', dtype='int8')
model[0].save_parameters("model.annotated/params.best")
model[0].save_config("model.annotated")
#Warning: do not use the loaded model for inference.  Load from disk.

Adding scaling factors and quantizing (transition 1 -> 3):

import sockeye.model
model = sockeye.model.load_model("model", for_disk_saving='int8', dtype='int8')
model[0].save_parameters("model.annotated/params.best")
model[0].save_config("model.annotated")
#Warning: do not use the loaded model for inference.  Load from disk.

In both cases you'll need the *.json and version copied manually.

[mjdenkowski]

Approved for merge into an intermediate branch for final cleanup.