This is an
HQQ all 4-bit (group-size=64) quantized
Llama3.1-8B-Instruct model.
We provide two versions:
pip install git+https://github.com/mobiusml/hqq.git #master branch fix
pip install bitblas #if you use the bitblas backend
1import torch
2from transformers import AutoTokenizer
3from hqq.models.hf.base import AutoHQQHFModel
4from hqq.utils.patching import *
5from hqq.core.quantize import *
6from hqq.utils.generation_hf import HFGenerator
7
8#Settings
9###################################################
10backend = "torchao_int4" #'torchao_int4' #"torchao_int4" (4-bit only) or "bitblas" (4-bit + 2-bit) or "gemlite" (8-bit, 4-bit, 2-bit, 1-bit)
11compute_dtype = torch.bfloat16 if backend=="torchao_int4" else torch.float16
12device = 'cuda:0'
13cache_dir = '.'
14
15#Load the model
16###################################################
17#model_id = 'mobiuslabsgmbh/Llama-3.1-8b-instruct_4bitgs64_hqq' #no calib version
18model_id = 'mobiuslabsgmbh/Llama-3.1-8b-instruct_4bitgs64_hqq_calib' #calibrated version
19
20model = AutoHQQHFModel.from_quantized(model_id, cache_dir=cache_dir, compute_dtype=compute_dtype, device=device).eval()
21tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
22
23#Use optimized inference kernels
24###################################################
25prepare_for_inference(model, backend=backend)
26
27#Generate
28###################################################
29#For longer context, make sure to allocate enough cache via the cache_size= parameter
30gen = HFGenerator(model, tokenizer, max_new_tokens=1000, do_sample=True, compile="partial").warmup() #Warm-up takes a while
31
32gen.generate("Write an essay about large language models", print_tokens=True)
33gen.generate("Tell me a funny joke!", print_tokens=True)
34gen.generate("How to make a yummy chocolate cake?", print_tokens=True)
35