This is an
HQQ all 4-bit (group-size=64) quantized
Llama3.1-70B-Instruct model.
pip install git+https://github.com/mobiusml/hqq.git #master branch fix
pip install bitblas
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###################################################
17model_id = 'mobiuslabsgmbh/Llama-3.1-70b-instruct_4bitgs64_hqq'
18model = AutoHQQHFModel.from_quantized(model_id, cache_dir=cache_dir, compute_dtype=compute_dtype, device=device)
19tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
20
21#Use optimized inference kernels
22###################################################
23prepare_for_inference(model, backend=backend) #takes a while to init...
24
25#Generate
26###################################################
27#For longer context, make sure to allocate enough cache via the cache_size= parameter
28gen = HFGenerator(model, tokenizer, max_new_tokens=1000, do_sample=True, compile="partial").warmup() #Warm-up takes a while
29
30gen.generate("Write an essay about large language models", print_tokens=True)
31gen.generate("Tell me a funny joke!", print_tokens=True)
32gen.generate("How to make a yummy chocolate cake?", print_tokens=True)
33