QAT and quant postprocessing for torch.nn.Embedding

[bfineran]

Embedding layers are a large storage requirement for transformer models. This PR introduces a pathway to quantize embedding values. By default this behavior will be enabled by the QuantizationModifier

Python Example

>>> import torch
>>> from sparseml.pytorch.optim import QuantizationModifier
>>> module = torch.nn.Embedding(100, 100)
>>> QuantizationModifier().apply(module)
>>> module
Embedding(
  100, 100
  (activation_post_process): FakeQuantize(
    fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
    (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)
  )
  (weight_fake_quant): FakeQuantize(
    fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            quant_min=-128, quant_max=127, dtype=torch.qint8, qscheme=torch.per_tensor_affine, ch_axis=-1,         scale=tensor([1.]), zero_point=tensor([0])
    (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)
  )
)

FP32 ONNX Graph

QAT ONNX Graph

Quant ONNX Graph

[markurtz]

Overall looks good, a few small comments

[rahul-tuli]