DBRX on 3x 3090 GPUs

test.py

# conda create -n dbrx python=3.10 -y && conda activate dbrx
# pip install torch transformers tiktoken flash_attn bitsandbytes

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("SinclairSchneider/dbrx-instruct-quantization-fixed", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("SinclairSchneider/dbrx-instruct-quantization-fixed", device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True, load_in_4bit=True)

input_text = "What does it take to build a great LLM?"
messages = [{"role": "user", "content": input_text}]
input_ids = tokenizer.apply_chat_template(messages, return_dict=True, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=200)
print(tokenizer.decode(outputs[0]))