Views
No views yet
pip install git+https://github.com/mobiusml/hqq.git;
pip install git+https://github.com/mobiusml/gemlite.git; #to use the gemlite backend
pip install bitblas #to use the bitblas backend1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from hqq.utils.patching import *
4from hqq.core.quantize import *
5from hqq.utils.generation_hf import patch_model_for_compiled_runtime
6
7#Settings
8###################################################
9backend = "gemlite" #"torchao_int4" (4-bit only) or "bitblas" (4-bit + 2-bit) or "gemlite" (8-bit, 4-bit, 2-bit, 1-bit)
10compute_dtype = torch.bfloat16 if backend=="torchao_int4" else torch.float16
11device = 'cuda:0'
12cache_dir = '.'
13model_id = "mobiuslabsgmbh/Mixtral-8x7B-Instruct-v0.1_4bitgs64_hqq_hf"
14
15model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=compute_dtype, cache_dir=cache_dir, device_map=device, attn_implementation="sdpa")
16tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
17
18#Use optimized inference kernels
19########################################################################
20prepare_for_inference(model, backend=backend, verbose=True)
21
22#Load gemlite cache for faster warm-up
23if(backend == 'gemlite'):
24 import gemlite
25 gemlite.core.GemLiteLinear.load_config('gemlite_config.json')
26
27#Generate
28########################################################################
29from hqq.utils.generation_hf import HFGenerator
30#Mixtral doesn't support cuda graphs with HF unfortuantely...
31#gen = HFGenerator(model, tokenizer, max_new_tokens=1000, do_sample=True, compile=None)
32
33gen = HFGenerator(model, tokenizer, max_new_tokens=1000, do_sample=True, compile="partial",
34 compile_options={"mode": "max-autotune-no-cudagraphs"}
35 )#.enable_cuda_graph()
36
37gen.generate("Write an essay about large language models", print_tokens=True)
38
39########################################################################
40# #Inference with model,generate()
41# from hqq.utils.generation_hf import patch_model_for_compiled_runtime
42
43# patch_model_for_compiled_runtime(model, tokenizer, pre_compile=False)
44
45# prompt = "Write an essay about large language models."
46# inputs = tokenizer.apply_chat_template([{"role":"user", "content":prompt}], tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True)
47# outputs = model.generate(**inputs.to(model.device), max_new_tokens=1000, cache_implementation="static", pad_token_id=tokenizer.pad_token_id)
48# #print(tokenizer.decode(outputs[0])
49
50########################################################################
51#Save gemlite cache
52if(backend == 'gemlite'):
53 gemlite.core.GemLiteLinear.cache_config('/tmp/gemlite_config.json') 1##################################################################
2import torch
3import torch.nn as nn
4from typing import Optional
5from vllm.model_executor.layers.linear import RowParallelLinear
6from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
7class MixtralMLPRowParallel(nn.Module):
8
9 def __init__(
10 self,
11 num_experts: int,
12 hidden_size: int,
13 intermediate_size: int,
14 quant_config: Optional[QuantizationConfig] = None,
15 ) -> None:
16 super().__init__()
17 self.num_experts = num_experts
18 self.ffn_dim = intermediate_size
19 self.hidden_dim = hidden_size
20
21 self.w1 = RowParallelLinear(self.hidden_dim,
22 self.ffn_dim,
23 bias=False,
24 quant_config=quant_config)
25 self.w2 = RowParallelLinear(self.ffn_dim,
26 self.hidden_dim,
27 bias=False,
28 quant_config=quant_config)
29 self.w3 = RowParallelLinear(self.hidden_dim,
30 self.ffn_dim,
31 bias=False,
32 quant_config=quant_config)
33
34 # TODO: Use vllm's SiluAndMul
35 self.act_fn = nn.SiLU()
36
37 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
38 w1_out, _ = self.w1(hidden_states)
39 w1_out = self.act_fn(w1_out)
40 w3_out, _ = self.w3(hidden_states)
41 current_hidden_states = w1_out * w3_out
42 current_hidden_states, _ = self.w2(current_hidden_states)
43 return current_hidden_states
44
45import vllm.model_executor.models.mixtral_quant as mixtral_quant
46mixtral_quant.MixtralMLP = MixtralMLPRowParallel
47##################################################################
48
49from vllm import LLM
50from vllm.sampling_params import SamplingParams
51model_id = "mobiuslabsgmbh/Mixtral-8x7B-Instruct-v0.1_4bitgs64_hqq_hf"
52
53llm = LLM(model=model_id, gpu_memory_utilization=0.80)
54sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=1024)
55outputs = llm.generate(["What is the capital of Germany?"], sampling_params)
56print(outputs[0].outputs[0].text)