Different quantize methods can be configured by --quantize_type="int8_per_channel"/"int4_per_channel"/"int8_blockwise"/"int4_blockwise" when running run_server.py, run_offline.py and run_interactive.py. (README is also updated accordingly)
Quantize bf16 checkpoint with int4 blockwise quantization for linear layers. (Embedding layers will be kept in int8 per-channel quant if int4 or blockwise quant is enabled)
Quantization configs are stored in QuantizationConfig dataclass, Environment stores a QuantizationConfig instance. Model initiate quantized layers based on the quantization config in environment.
Int4 weight loading workflow:
There is no torch.int4, so convert_checkpoints will store the in4 weights in int8 container
When the jax state_dict is extracted from checkpoint in engine.py, we cast the int8 tensors to int4 JAX tensors.
New quantization support
Added {int8, int4} x {per_channel, blockwise} quantized linear layers.
Added asymmetric quant support to quant/dequant function and quantized layers, but it's not exposed to cmd config, for experimental purpose.
Changes:
Add quantization cmd configs.
Add support of int4 and blockwise quantization in convert_checkpoint.py, move quantization map to model class.
Implementation of blockwise quantized linear layer. Explored different implementations and only the einsum without flatten operands works the best, the others are still kept in case they become useful in the future.
Int4 weight loading in engine.py
Refactor the weight loading logic in engine.py
Update quantize and dequantize function (Enhancement accuracy by adding eps, round, blockwise and asymmetric quant support)
Test:
Correctness
Added unit tests for quantize/dequantize function, quantized linear accuracy, blockwise quant sharding.
run_interactive to verify the int4 blockwise quant generates meaning tokens
Latency
Summarize in an internal doc, it's not always outperforming int8 per-channel, but int4 blockwise matmul kernel is faster.
User journey
Different quantize methods can be configured by
--quantize_type="int8_per_channel"/"int4_per_channel"/"int8_blockwise"/"int4_blockwise"
when runningrun_server.py
,run_offline.py
andrun_interactive.py
. (README
is also updated accordingly)run_interactive
, similar forrun_offline.py
andrun_server.py
Quantization config workflow
Quantization configs are stored in
QuantizationConfig
dataclass, Environment stores aQuantizationConfig
instance. Model initiate quantized layers based on the quantization config in environment.Int4 weight loading workflow:
torch.int4
, soconvert_checkpoints
will store the in4 weights in int8 containerengine.py
, we cast the int8 tensors to int4 JAX tensors.New quantization support
Changes:
convert_checkpoint.py
, move quantization map to model class.einsum without flatten operands
works the best, the others are still kept in case they become useful in the future.engine.py
engine.py
eps
,round
, blockwise and asymmetric quant support)Test: Correctness
run_interactive
to verify the int4 blockwise quant generates meaning tokens Latency