Support Int4 ONNX Export

[Satrat]

Int4 QAT was working out of the box, but the onnx export step was failing because onnx only supports 8-bit integer quantization. The fix for this was to trick onnx into thinking we were exporting 8-bit quantized weights by manually adjusting the range of the FakeQuantize modules. The weights themselves will still be between [-8, 7], and the original range boundaries are still accessible via the FakeQuantize.quant_min and FakeQuantize.quant_max attributes (they are deprecated by PyTorch, replaced by FakeQuantize.activation_post_process.quant_[min\max])

Testing

Run tests/sparseml/pytorch/test_torch_to_onnx_exporter/test_export_4bit_model for example usage, or just run TorchToONNX.export() on any Int4 QAT model:

model = ...
sample_batch = ...
exporter = TorchToONNX(sample_batch)
exporter.export(model, "model.onnx")

Old Behavior

Export fails with error:

"name": "SymbolicValueError",
"message": "For (quant_min, quant_max), ONNX allows only (0, 127), (0, 255) and (-128, 127). Got (-8, 7)

New Behavior

Export completes without errors

[rahul-tuli]

Nice diff! LGTM pending comments