Package: gptq-triton
GPTQ-tritonThis is my attempt at implementing a Triton kernel for GPTQ inference.  This code is based on theGPTQ-for-LLaMacodebase, which is itself based on theGPTQcodebase.@article{frantar-gptq,
  title={{GPTQ}: Accurate Post-training Compression for Generative Pretrained Transformers}, 
  author={Elias Frantar and Saleh Ashkboos and Torsten Hoefler and Dan Alistarh},
  year={2022},
  journal={arXiv preprint arXiv:2210.17323}
}Installationpip install .MotivationAs of today (2023-03-27) the CUDA kernels in the aforementioned codebases do not scale well with context length, running up to 10x slower when the context is large versus the equivilent FP16 model.  To solve this I'm implementing the inference kernel in Triton, which should allow for much better scaling.The implementation is based around the matmul tutorial from the Triton documentation.  The main difference is decoding the quantized weights before performing each sub-block of the matrix multiplication.Fusing of the FF layers and QKV matrix are also applied.PerformanceThis benchmark was run on a 3090 using thebenchmark_generate.pyscript.Accuracy (PPL)The following results were obtained using theppl.pyscript with a stride of 512 and a context length of 2048.
For the 4bit CUDA results, a custom version ofppl.pywas used, as the current script is dedicated to the Triton kernel convensions.
it/s numbers are from a 3090.LLaMA-7BBitsgroup-sizememory(MiB)it/sWikitext2PTBC4FP1616-173731.645.047.856.99GPTQ CUDA4-188050.115.448.24-GPTQ Triton4-163231.705.448.247.48LLaMA-13BBitsgroup-sizememory(MiB)it/sWikitext2PTBC4FP1616-31633-4.527.196.66GPTQ Triton4-1103250.954.747.497.00GPTQ Triton412895470.924.677.386.99LLaMA-30BBitsgroup-sizememory(MiB)it/sWikitext2PTBC4FP1616-72491-3.616.506.07GPTQ Triton4-1199890.403.896.716.44GPTQ Triton4128200550.393.796.706.34GPTQ Triton4512195470.393.786.626.30RequirementsSeesetup.cfg, but note that a nightlytransformersis preferred right now. v4.28.1 might work.  Known workingtransformerscommit is28f26c107b4a1c5c7e32ed4d9575622da0627a40.Quantizing a modelThequantize.pyscript is used to quantize a HuggingFace model.  Example usage:./quantize.py --model <Path to a HF FP16 model> --dataset c4 --wbits 4 --groupsize -1 --act-order --true-sequential --save <Path to the output folder>Arguments:--model: Path to a HF FP16 model--dataset: Dataset to use for calibration.  Can bewikitext-2,ptb,ptb-neworc4.--seed: Seed for sampling the calibration data.--nsamples: Number of calibration data samples.--percdamp: Percent of the average Hessian diagonal to use for dampening (default 0.01).--wbits: Number of bits to use for quantization.--groupsize: Groupsize to use for quantization; default (-1) uses full row.--save: Save quantized result to this folder.--safetensors: Save using the safetensors format.--act-order: Use activation order quantization.--true-sequential: Use true sequential quantization.NOTE:The Triton kernel is currently only implemented for 4-bits and groupsize -1.Explanation ofgroupsizeThe GPTQ quantization algorithm gets applied tonn.Linear,nn.Conv2d, andtransformers.Conv1dlayers.  (NOTE:quantize.pycurrently only supports LLaMA like models, and thus onlynn.Linearlayers are quantized, andlm_headis skipped.)  Each matrix is quantized into a quantized weight matrix, quantized zeros, and float16 scale (bias is not quantized).  During matmul, the weights are decoded using the formulaw = (w - z - 1) * s.Scales and zeros are per-outfeature, so when there is no grouping, scales and zeros would be1xOutfeatures.  That means that each row of the matrix (i.e. along the infeatures dimension) is quantized using the same scalar scale and zero.  When grouping is used, each row is split intogroupsizevalues, and each group is quantized using its own scalar scale and zero.  This means that the scales and zeros are(Infeatures//groupsize)xOutfeatures.Groupsize provides a tradeoff.  Lower groupsizes offer more granularity to the quantization and thus less loss of accuracy, but decrease the memory savings offered by quantization.Explanation ofnsamplesanddatasetnsamples and dataset effect the calibration data.  This input data is fed through the network during quantization to calibrate the algorithm.  See the GPTQ paper for more detail.Explanation oftrue-sequentialModels are quantized sequentially, one "layer" at a time.  For example, LLaMA 7B has 32 layers, starting after the input embedding, following by the head.  Each of these layers has many different Linear modules that can be quantized.  Without thetrue-sequentialflag, these Linear modules will be quantized in an arbitrary order.  With thetrue-sequentialflag, the Linear modules will be quantized in the order they would be encountered during a forward pass.  This can provide an accuracy boost.Explanation ofact-orderI don't know.  Looking at the code (ingptq.py), it seems to re-order the matrix before quantization based on theargsortof the estimatedH.  The order in which the matrix columns of the matrix are quantized might have an impact on final accuracy.act-orderwas introduced by the GPTQ authors to improve accuracy when quantizing "small" models like LLaMA 7B.Filesbenchmark_generate.py- A script for benchmarking generation speed at different prompt lengths and generation lengths.Benchmark.ipynb- A notebook for benchmarking the Triton kernel against the CUDA kernel and FP16.quantize.py- A script for quantizing a model.generate.py- An example script for generating text from a model.  Example usage:./generate.py --model <Path to your quantized model> --quant --prompt "Write a story about a duck: Once upon a time there was a duck" --temperature 0.6 --top-p 0.6 --repetition-penalty 1.1ppl.py- A script for calculating the perplexity of a model against wikitext2, PTB, and C4.  This is useful for verifying correctness of the Triton kernel, comparing it to the CUDA kernel and the original FP16 model.Verify.ipynb- A notebook for verifying the correctness of the Triton kernel.