Package: aqtp
AQT : Accurate Quantized TrainingAQT is a software library designed for easy tensor operation quantization in JAX.FeaturesExcellent quantized int8 model quality with no hand-tuning at all.Excellent training speedup in production using contemporary ML accelerators.Simple and flexible APIs suitable for both production and research.AQT is a tensor op (matmul, einsum, conv, ...) quantization library.
Does not make assumptions about the computations to be Neural Net.AQT can be injected into any Jax computation.
Well tested in:Flax,Pax,MaxTextand other frameworks at Google.AQT is designed for production workloads.
It is well exemplified by these blog entries:Cloud ML AQT introduction50000+ TPU Cloud Ml training with AQTAQT is designed for research.
It is well exemplified by these papers:Pareto-Optimal Quantized ResNet Is Mostly 4-bitPokeBNN: A Binary Pursuit of Lightweight Accuracy4-bit Conformer with Native Quantization Aware Training for Speech RecognitionBinarized Neural Machine TranslationWYTIWYS (What you train is what you serve).
AQT quantized models are bit-exact the same during training and serving.
This side-steps the conventional issue of quantization-induced training-serving bias
that typically happens for Post Training Quantization (PTQ).Let us know if you have any problem with aqt applications by filing an issue on Github.UsageTensor contraction operations in JAX-based neural network libraries, i.e., any form of (high-order) matrix multiplications, including but not limited tojax.numpy.einsumandflax.linen.DenseGeneral, calllax.dot_generalas its core computation. Quantizing a neural network in JAX simply requires substitutinglax.dot_generalwith a quantized variant and keeping other parts as-is, which we call "quantization injection". JAX-based NN libraries, such asFlaxandPax, provide an API for this substitution when creating layers.In this section, we show how AQT produces a quantizeddot_generaland inject it into a neural network defined in JAX. The toy example below can be found in theexample colab.First, install the AQT package named asaqtpin PyPI.# install the AQT library!pipinstallaqtpNext, importaqt.jax.v2.
OtherAQT versionsare obsolete.# necessary importsimportaqt.jax.v2.flax.aqt_flaxasaqtimportaqt.jax.v2.configasaqt_configimportflax.linenasnnAsample neural networkdefined in Flax looks like the following (as a toy example we use a simple MLP, but it can be any model):classMlpBlock(nn.Module):config:aqt_config.DotGeneral|None@nn.compactdef__call__(self,inputs):dot_general=aqt.AqtDotGeneral(self.config)x=nn.Dense(dot_general=dot_general,features=inputs.shape[-1]*4)(inputs)x=nn.relu(x)x=nn.Dense(dot_general=dot_general,features=inputs.shape[-1])(x)returnxAQT can quantize the model by simply replacing thedot_generalinnn.Densewith a quantized dot_general created by the aqt configuration. The example specifies an AQT configuration that quantizes both forward and backward passes to int8. Now let's test it.importjaximportjax.numpyasjnpimportnumpyasnp# Generate some random matrices as inputsdefgen_matrix(rows,columns,seed=0):np.random.seed(seed)returnnp.random.normal(size=(rows,columns)).reshape((rows,columns))inputs=gen_matrix(3,4)# test function that initializes the model and compute the forward passdefinit_and_eval(name,mlp_block,init_seed=0,eval_seed=0):model=mlp_block.init(jax.random.PRNGKey(init_seed),inputs)out=mlp_block.apply(model,inputs,rngs={'params':jax.random.key(eval_seed)})print(f"{name}:\n",out)# create a config that quantizes both forward and backward passes to int8int8_config=aqt_config.fully_quantized(fwd_bits=8,bwd_bits=8)# run and print resultsmlp_fp16=MlpBlock(config=None)mlp_int8=MlpBlock(config=int8_config)init_and_eval('mlp_fp16',mlp_fp16)init_and_eval('mlp_int8',mlp_int8)Results will be the following:mlp_fp16:
 [[ 0.720744    1.5375545  -2.6456933  -1.7605033 ]
 [-0.01541612  0.09728499 -1.5742414  -0.3737522 ]
 [ 0.4071759   1.1941448  -0.6982092  -0.48336366]]
mlp_int8:
 [[ 0.7030779   1.5099456  -2.6334763  -1.7550919 ]
 [-0.00901393  0.08774488 -1.5644912  -0.3728472 ]
 [ 0.40121436  1.189411   -0.6939187  -0.48000643]]We can see that the quantized MLP produces similar outputs as the unquantized one.Flexible Quantization ConfigsThe example inusageuses the default configuration that quantizes both forward and backward passes to 8-bit, but AQT provides a much more flexible configuration system. TheDotGeneralclass can configure forward and backward tensor contraction operations separately.@dataclasses.dataclassclassDotGeneral:"""Configuration of quantization of dot_general and its gradients."""fwd:DotGeneralRawdlhs:DotGeneralRawdrhs:DotGeneralRawIn eachDotGeneral.DotGeneralRaw, we can configure quantization of each input tensor of those ops separately and the hardware dtype to use (eg. jnp.bfloat16, jnp.float16, jnp.float8_e4m3fn, jnp.float8_e5m2, jnp.int8, jnp.int4).@dataclasses.dataclassclassDotGeneralRaw:"""Configuration of quantization of one dot_general without gradient."""lhs:Tensor# left hand siderhs:Tensor# right hand sidedg_in_dtype:Optional[DType]dg_accumulator_dtype:Optional[DType]local_aqt:Optional[LocalAqt]# sharded quantizationInside config.Tensor we can configure the numerics used for each tensor, which includes number of bits, calibration algorithm, stochastic rounding, and many other quantization parameters.@dataclasses.dataclassclassTensor:"""Configuration of quantization of one tensor or one side of tensor op."""numerics:Numericscalib_shared_axes:Optional[list[int]]scale_stop_grad:boolcalibration:calibration.Calibration# calibration algorithmpo2_scale:bool# round calibration to power of 2use_fake_quant:booluse_fwd_quant:Optional[bool]# use quantized fwd in the bwd passAQT VersionsAs of today there are several independent AQT implementations in this package:JAX Legacy AQTObsolete version of AQT still used by some customers.JAX AQTv1Version of AQT that was developed with acceleration of NN inference in mind.TF AQTv1Tensorflow counterpart of JAX AQTv1.JAX AQTv2AQT implementing universal matmul quantization.AQTv2 is the recommended library.
We plan to port remaining features from AQTv1 to AQTv2 and
delete AQTv1 in early Q1 2024. Below we describe details about that.Inference AccelerationThe most important AQTv2 (to be ported from AQTv1) missing features are:https://github.com/google/aqt/issues/282https://github.com/google/aqt/issues/280Lack of these features prevents AQTv2 from accelerating inference with small batch.
The only option today is dynamic quantization where
each tensor op is quantized independently and quantization scales are found just-in-time.Backpropagation AccelerationAQTv2 speeds up training and fine-tuning.
We verified 1.2x to 1.4x reduction in step time on 1B to 16B large Transformer
models to a given quality on TPUs.Today in order to do it correctly one needs to understand that for each
two-argument tensor op (matmul, einsum, conv) in the forward pass,
there are two in the backward pass.
One has to understand how to configure them.We will be updating config file with current best practices.AQT ServingQuantization is applied to both activations and weights right before a matmul.
During training, gradients flow "through" the operation and update the latent
floating-point weights.
During model serving, however, re-quantizing weights is unnecessary and has large overhead because:Weights are stationary during serving. Recomputing weight quantization is a waste.Recomputing quantization requires loading the latent BF16 weights, which
consumes 2x more memory bandwidth than loading INT8 weights.AQT provides a solution to the above problem, which is called "serving conversion".
Serving conversion is a constant folding process that creates new variables in
the checkpoint to store the folded INT8 weights.
It requires running one dummy inference.Note that:Weights can be either lhs or rhs inputs to a matmul. AQT supports storing
both in checkpoints, configured bylhs_quant_mode and rhs_quant_mode.
Users need to set the correct configuration themselves.In some cases where checkpoint variables require metadata such as sharding axis,
users can configurevariable initializersto fit their needs.At serving mode, AQT will look for the INT8 variables in a checkpoint.
If found, it skips the weight quantization and returns INT8 variables
as-is as inputs to matmul, thus saving the memory bandwidth and avoids recalculation.Applying the conversion to an unquantized floating-point model is
equivalent to post-training quantization (PTQ) serving.
Applying the conversion to a forward-only quantized AQT model is
quantization-aware training (QAT) serving.
In that case it is important to use the same AQT config during training and serving to maintain WYTIWYS.Theflax end-to-end exampleprovides a code snippet on how to perform serving conversion and model serving in AQT.Other Weight TransformationsConsidermatmul(a, w)as an activation-weight matmul.
Sometimes there is another transformationTon weights before passed into the matmul,
i.e.,matmul(..., T(w))is computed.
In this case, quantizing weights directly, i.e.,matmul(..., T(Q(w))), can reduce the
checkpoint size and save memory bandwidth, but it will not accelerate the matmul
becauseTwill likely return floats.
In order to get matmul acceleration, the quantization functionQshould be
inserted just before matmul, i.e.,matmul(..., Q(T(w))).AQT pursues the goal of both compressing the checkpoint AND accelerating the matmul.
This requires storing the entirew_q = Q(T(w))in the checkpoint and
using it in serving directly, i.e.,matmul(..., w_q).Note that AQT provides a quantizedmatmul_aqtas a whole such thatmatmul_aqt(..., T(w)) = matmul(..., Q(T(w))).Q(T(w))is not visible outside ofmatmul_aqt.
The main reason is thatmatmul_aqthas custom gradient defined for it.How AQT Works InternallyIn this section we:show how to get quantization acceleration in JAX,explain what AQT INT8 does under-the-hood (using the simplest INT8 configuration),run the code on a simple example.Code in this section can be found and executable in theexample colab.
Note that this section mainly explains how AQT works and why it can achieve a good quality. For AQT tutorial, user can refer to theusagesection.Thematmul_true_int8takes real INT8 as inputs, returns int32. The matmul computationjnp.matmulcallslax.dot_generalin its source, which is a JAX wrapper forXLA DotGeneralop that implements all MXU ops (this is where we have int8 acceleration on TPUs) except convolution. This is how one can get hardware acceleration of quantized matmul in JAX.importjax.numpyasjnpdefmatmul_true_int8(lhs,rhs):assertlhs.dtype==jnp.int8assertrhs.dtype==jnp.int8result=jnp.matmul(lhs,rhs,preferred_element_type=jnp.int32)assertresult.dtype==jnp.int32returnresultGenerate some random data:batch_size=3channels_in=4channels_out=5a=gen_matrix(batch_size,channels_in)# Activationsw=gen_matrix(channels_in,channels_out)# WeightsBelow is how AQT works internally using the simplest INT8 configuration.
Even though names such as "batch" and "channels" are used, "w" and "a", which are evocative of neural networks, one may note thataqt_matmul_int8algorithm is not DNN specific.defaqt_matmul_int8(a,w):max_int8=127# This function is customizable and injectable, i.e:# users can inject custom quant code into an AQT config.defquant_int8(x):returnjnp.clip(jnp.round(x),-max_int8,max_int8).astype(jnp.int8)# Calibration. Calibration function is also customizable and injectable.a_s=max_int8/jnp.max(jnp.abs(a),axis=1,keepdims=True)w_s=max_int8/jnp.max(jnp.abs(w),axis=0,keepdims=True)asserta_s.shape==(batch_size,1)# shapes checked for illustrationassertw_s.shape==(1,channels_out)# int8 matmul with int32 accumulatorresult=matmul_true_int8(quant_int8(a*a_s),quant_int8(w*w_s))/(a_s*w_s)assertresult.shape==(batch_size,channels_out)returnresultNote that each example in a batch and each output channel will have their own separate scale. This reduces the effect of outliers in "w" and "a" to just one row or column, making a tighter calibration and much better quality of quantization. Comparing aqt_matmul_int8 to float matmul, their outputs are close.print(f"jnp.matmul(a, w):\n",jnp.matmul(a,w))print(f"aqt_matmul_int8(a, w):\n",aqt_matmul_int8(a,w))# should expect the following outputs
jnp.matmul(a, w):
 [[ 3.6095254   5.8575077   1.9510972   4.732388    1.9792626 ]
 [ 4.335892    0.9743651   2.7298734   4.3540883   3.637487  ]
 [-0.07735002  2.7310796  -0.3519049   0.19912864 -1.2023292 ]]
aqt_matmul_int8(a, w):
 [[ 3.5998788   5.8562713   1.9385538   4.7426414   1.9792401 ]
 [ 4.321886    0.99681264  2.737299    4.3591022   3.6352503 ]
 [-0.07714217  2.7415617  -0.35343346  0.20568734 -1.1974115 ]]Citing AQTWe will be publishing AQT whitepaper soon.