# Introducing DBRX: The Future of Language Models

Davis Blalock, Bandish Shah, Abhi Venigalla & Ajay Saini, Databricks | MLOps Coffee Sessions | 48:36

Source: https://www.youtube.com/watch?v=J26DlKXJ7nQ
Channel: MLOps Community, now AAIF Live (https://www.youtube.com/@AAIFLive-x1r). Summarised by MLOps Talks.
Page: https://mlopstalks.com/talks/introducing-dbrx-the-future-of-language-models
Published: 2024-04-12
Tags: data-pipelines, gpus, model-serving, training-pipelines

## TL;DR
- DBRX uses a mixture-of-experts architecture with 16 feed-forward experts, routing each token to four of them.
- The team trained DBRX across more than 3,000 GPUs and built monitoring and on-call systems around frequent hardware and network failures.
- DBRX has 132 billion total parameters and was pretrained on 12 trillion tokens, with later training focused more heavily on high-quality data and code.

## Summary
This roundtable explains how Databricks built and trained DBRX. The team moved from dense models toward a mixture-of-experts design, giving the model more parameters while activating only part of the network for each token. Davis Blalock explains why DBRX uses 16 experts and routes each token to four. The engineering discussion covers training on more than 3,000 GPUs, frequent GPU failures, network checks, checkpoint handling, and a dedicated on-call rotation. The team also describes its use of PyTorch FSDP, Composer, streaming data sets, Kubernetes, and custom inference serving. DBRX was trained with 132 billion total parameters and 12 trillion tokens. Toward the end of pretraining, the team shifted the data mix toward higher-quality text and code. The panel also discusses synthetic data, fine-tuning, smaller models, and running quantized DBRX on high-end laptops.

## Key ideas
### DBRX grew out of a push to make higher-quality models without raising training costs as quickly
[05:40](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=340s)
The team began with the goal of helping people build and train their own models. After showing that models near GPT-3 quality could be trained at relatively low cost, they wanted to move toward GPT-3.5 quality and beyond. That required larger models and a more efficient design. The team explored mixture-of-experts architectures, worked on both research and engineering problems, and decided around Christmas to build a model more powerful than the earlier MPT series. The effort then ran around the clock for roughly three months.

### Mixture-of-experts gives DBRX more capacity while each token uses only part of the model
[12:23](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=743s)
Davis Blalock describes DBRX as one model with 16 separate feed-forward networks in each Transformer block. A routing function sends each token to four of those experts. This gives the model many more parameters and possible combinations than a single dense network, while avoiding the cost of running every expert for every token. Using more than one expert gives the model a larger expressive space than routing each token to only one. The tradeoff is that each additional expert increases the computation for every token.

### Large-scale training became a systems and operations problem
[18:57](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=1137s)
The team had worked with thousands of GPUs across many smaller jobs, but DBRX was its first training run at this scale. On a cluster of more than 3,000 GPUs, the team observed roughly three to four GPU failures per day during the run. Monitoring covered the GPUs, Kubernetes, and the network links between machines. When monitoring missed a problem, a dedicated 24/7 on-call rotation handled hardware and training failures. Each incident added to a logbook of failure modes that the team planned to automate before customers trained models at the same scale.

### PyTorch FSDP provided a flexible training stack for different model types
[22:54](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=1374s)
The team used a PyTorch-based stack with Composer on top. Composer was open source and supported training LLMs, diffusion models, and other model types. The team chose PyTorch FSDP because it did not impose strong requirements on model architecture and made it easier to distribute models across large clusters. They reported getting performance close to three-dimensional parallel strategies after optimizing FSDP. High-speed interconnects helped, although the panel noted that other parallelism strategies can still matter on different hardware setups.

### Inference measurement required testing the whole serving system
[25:16](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=1516s)
For model and inference analysis, the team used the PyTorch profiler and memory profiler, along with some internal logging for distributed calls. Inference performance depended on workload shape, batch size, concurrent requests, and burstiness, so simple hardware measurements were not enough. Databricks built an optimized inference web server for DBRX and tested it by sending large numbers of requests through the deployed service. The team measured time to first token, time per output token, generation time, throughput, and memory bandwidth utilization. It also worked with vendor teams and used NVIDIA tooling to inspect kernel performance.

### DBRX used a high token-to-parameter ratio based on smaller-scale experiments
[29:29](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=1769s)
The panel reviews how scaling guidance changed from early large models toward training smaller models on much more data. The team tested different model sizes and token counts at smaller scale and looked for the best tradeoff for customer use cases. DBRX has 132 billion total parameters and was trained on 12 trillion tokens, which is close to a 100-to-one token-to-parameter ratio. The team wanted to use the smallest parameter count it could while preserving quality, rather than simply increasing model size.

### Data processing and streaming had to work across thousands of GPUs
[32:06](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=1926s)
The team used Databricks tools and Spark to process very large text datasets, with Abhi Venigalla describing a roughly 100-fold improvement in data processing speed after the acquisition. Streaming data sets helped split, shuffle, and distribute training data across many GPUs. Object storage created separate problems at scale. Listing files and having every worker download the same metadata could overload storage services, so the team added fixes such as sharing metadata instead of requesting it from every GPU. The panel said the streaming data loader was one part of the infrastructure that held up well as the system scaled.

### The final part of pretraining emphasized high-quality data and code
[41:32](https://www.youtube.com/watch?v=J26DlKXJ7nQ&t=2492s)
The team used curriculum learning during the last 10% or 20% of pretraining. Instead of keeping the data mix focused mainly on web-scale text, it increased the share of high-quality data and code. The goal was to improve the general-purpose model's coding and reasoning performance. Abhi said smaller-scale experiments showed that this change improved the evaluations the team cared about. The data pipeline made it possible to change stream ratios during training. The panel also discussed synthetic data as a way to expand small enterprise datasets and support future fine-tuning products.

## Notable quotes
- Ajay Saini: "When you hit that scale, all sorts of things break." (19:15)
- Davis Blalock: "In a mixture of experts model, what you do is instead of having completely separate models you're going to have one model but all of the feed forward networks pairs of linear layers in a Transformer are going to be duplicated." (13:23)
- Bandish Shah: "We train DBRX on over 3,000 different GPUs." (10:43)
- Abhi Venigalla: "We wanted a general purpose LLM that was also particularly really good at coding and reasoning." (41:50)
- Ajay Saini: "Kubernetes is significantly easier to build multicloud." (43:53)

## Tools & references mentioned
- DBRX
- Databricks
- MPT
- MosaicML
- Mosaic AI
- Composer
- LM Foundry
- PyTorch
- PyTorch FSDP
- DeepSpeed
- Megatron
- Kubernetes
- NVIDIA
- TensorRT-LLM
- Llama.cpp
- Spark
- Chinchilla
- GPT-3
- GPT-3.5
- Llama
- Mistral
- Groq
- Rancher
- AWS
- Azure
- Oracle Cloud Infrastructure
- Google Cloud Platform

## Who should watch
- You are deciding whether a mixture-of-experts model is worth the added training and serving complexity.
- You run large GPU training jobs and want practical details about failures, network monitoring, checkpoint handling, and on-call work.
- You are building data or inference infrastructure for an LLM and want to understand how preprocessing, streaming, benchmarking, and serving fit together.

## Related talks

- [DevTools for Language Models: Unlocking the Future of AI-Driven Applications](https://mlopstalks.com/talks/devtools-for-language-models-unlocking-the-future-of-ai-driven-applications) (Diego Oppenheimer, Factory, 29:55)
- [Let's Continue Bundling into the Database](https://mlopstalks.com/talks/lets-continue-bundling-into-the-database) (Ethan Rosenthal, Square, 51:56)
- [Efficiently Scaling and Deploying LLMs](https://mlopstalks.com/talks/efficiently-scaling-and-deploying-llms) (Hanlin Tang, MosaicML, 25:14)
- [Large Language Models in Production Round-table Conversation](https://mlopstalks.com/talks/large-language-models-in-production-round-table-conversation) (Diego Oppenheimer, Factory HQ & David Hershey, Unusual Ventures & Hannes Hapke, Digits & James Richards, Bountiful & Rebecca Qian, Facebook AI Research, 57:21)
- [The Next Revolution in AI: LLMs and Beyond](https://mlopstalks.com/talks/the-next-revolution-in-ai-llms-and-beyond) (, 13:47)