DBRX

DBRX is a large language model trained by Databricks, and made available under an open license. This repository contains the minimal code and examples to run inference, as well as a collection of resources and links for using DBRX.

• Founder's Blog, DBRX Technical Blog
• Hugging Face: https://huggingface.co/collections/databricks/
• LLM Foundry: https://github.com/mosaicml/llm-foundry

A reference model code can be found in this repository at modeling_dbrx.py.

Note: this model code is supplied for references purposes only, please see the Hugging Face repository for the official supported version.

Model details

DBRX is a Mixture-of-Experts (MoE) model with 132B total parameters and 36B live parameters. We use 16 experts, of which 4 are active during training or inference. DBRX was pre-trained for 12T tokens of text. DBRX has a context length of 32K tokens.

The following models are open-sourced:

Model	Description
DBRX Base	Pre-trained base model
DBRX Instruct	Finetuned model for instruction following

The model was trained using optimized versions of our open source libraries Composer, LLM Foundry, MegaBlocks and Streaming.

For the instruct model, we used the ChatML format. Please see the DBRX Instruct model card for more information on this.

Quick start

To download the weights and tokenizer, please first visit the DBRX Hugging Face page and accept the license. Note: access to the Base model requires manual approval.

We recommend having at least 320GB of memory to run the model.

Then, run:

pip install -r requirements.txt # Or requirements-gpu.txt to use flash attention on GPU(s)
huggingface-cli login           # Add your Hugging Face token in order to access the model
python generate.py              # See generate.py to change the prompt and other settings

For more advanced usage, please see LLM Foundry (chat script, batch generation script)

If you have any package installation issues, we recommend using our Docker image: mosaicml/llm-foundry:2.2.1_cu121_flash2-latest

Inference

Both TensorRT-LLM and vLLM can be used to run optimized inference with DBRX. We have tested both libraries on NVIDIA A100 and H100 systems. To run inference with 16-bit precision, a minimum of 4 x 80GB multi-GPU system is required.

TensorRT-LLM

DBRX support is being added to TensorRT-LLM library: Pending PR

After merging, instructions to build and run DBRX TensorRT engines will be found at: README

vLLM

Please see the vLLM docs for instructions on how to run DBRX with the vLLM engine.

MLX

If you have an Apple laptop with a sufficiently powerful M-series chip, quantized version of DBRX can be run with MLX. See instructions for running DBRX on MLX here.

LLama.cpp

If you have an Apple M-series chip laptop with atleast 64GB RAM, you can run a quantized version of DBRX using llama.cpp.

1. Compile llama.cpp
2. Download a quantized ggml version of dbrx-instruct such as dranger003/dbrx-instruct-iMat.GGUF
3. From llama.cpp folder, run:

./main -ngl 41 -m ./models/ggml-dbrx-instruct-16x12b-iq1_s.gguf -n 256 --repeat_penalty 1.0 --color -i -r "User:" -f prompts/chat-with-bob.txt

Finetune

To finetune DBRX with our open source library LLM Foundry, please see the instructions in our training script (found here). We have finetuning support for both:

• Full parameter finetuning, see the yaml config dbrx-full-ft.yaml
• LoRA finetuning, see the yaml config dbrx-lora-ft.yaml

Note: LoRA support currently cannot finetune the experts, since the experts are fused. Stay tuned for more.

Model card

The model cards can be found at:

• DBRX Base
• DBRX Instruct

Integrations

DBRX is available on the Databricks platform through:

• Mosaic AI Model Serving
• Mosaic AI Playground

Other providers have recently added support for DBRX:

• You.com
• Perplexity Labs
• LlamaIndex (starter example gist)

The same tools used to train high quality MoE models such as DBRX are available for Databricks customers. Please reach out to us at https://www.databricks.com/company/contact if you are interested in pre-training, finetuning, or deploying your own DBRX models!

Issues

For issues with model output, or community discussion, please use the Hugging Face community forum (instruct, base)

For issues with LLM Foundry, or any of the underlying training libraries, please open an issue on the relevant GitHub repository.

License

Our model weights and code are licensed for both researchers and commercial entities. The Databricks Open Source License can be found at LICENSE, and our Acceptable Use Policy can be found here.