DBRX: Databricks' Enterprise MoE Architecture

DBRX uses 16 experts and activates 4 per token, giving $\binom{16}{4} = 1,820$ expert combinations versus Mixtral’s $\binom{8}{2} = 28$. The result: DBRX achieves better specialization at the same memory footprint (36B active parameters for both models). Databricks’ bet is that fine-grained expert routing matters more for enterprise workloads — SQL generation, structured data analysis, and domain-specific reasoning — where task diversity exceeds what 8 coarse experts can capture.

Architecture Details

import torch
import torch.nn as nn

class DBRXConfig:
    """DBRX architecture configuration — exact values from release."""
    vocab_size = 100352          # Large vocabulary for enterprise text
    hidden_size = 6144
    num_hidden_layers = 40
    max_position_embeddings = 32768

    # Attention
    num_attention_heads = 48
    num_key_value_heads = 8      # GQA: 6 query heads per KV head
    head_dim = 128
    attention_bias = False       # No bias in attention projections
    clip_qkv = 8.0             # Clip QKV values for stability

    # MoE
    num_experts = 16
    num_experts_per_tok = 4      # Top-4 routing (vs Mixtral's top-2)
    expert_intermediate_size = 10752  # Per-expert FFN intermediate

    # RoPE
    rope_theta = 500000.0

    # Derived
    total_params = 132e9         # 132B total
    active_params = 36e9         # 36B active per token

    # Each expert: 3 * 6144 * 10752 = 198M params
    # 16 experts per layer: 3.17B expert params per layer
    # 40 layers: 126.8B expert params
    # Remaining ~5.2B: attention, embeddings, norms, router

class DBRXMoELayer(nn.Module):
    """
    DBRX MoE: 16 experts, top-4 routing.
    Key difference from Mixtral: more experts, higher top-k.
    """

    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_tok

        # Router with learned gating
        self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False)

        # 16 SwiGLU experts
        self.experts = nn.ModuleList([
            SwiGLUExpert(config.hidden_size, config.expert_intermediate_size)
            for _ in range(config.num_experts)
        ])

    def forward(self, hidden_states):
        batch_size, hidden_dim = hidden_states.shape

        # Routing with softmax
        router_logits = self.router(hidden_states)
        routing_weights = torch.softmax(router_logits, dim=-1)

        # Top-4 selection
        top_k_weights, top_k_indices = torch.topk(
            routing_weights, self.top_k, dim=-1
        )
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)

        # Expert computation
        output = torch.zeros_like(hidden_states)
        for k in range(self.top_k):
            for e in range(self.num_experts):
                mask = (top_k_indices[:, k] == e)
                if mask.any():
                    expert_out = self.experts[e](hidden_states[mask])
                    output[mask] += top_k_weights[mask, k:k+1] * expert_out

        return output, router_logits

Fine-Grained MoE: 16 Experts vs 8 Experts

from math import comb, log2

def fine_grained_analysis():
    """
    DBRX's fine-grained approach: more experts, each smaller,
    with higher top-k selection.

    Mixtral: 8 experts, top-2, expert_dim=14336
    DBRX:   16 experts, top-4, expert_dim=10752
    """
    configs = {
        'Mixtral (8E top-2)': {
            'experts': 8, 'top_k': 2,
            'expert_dim': 14336, 'hidden': 4096,
            'combos': comb(8, 2),
        },
        'DBRX (16E top-4)': {
            'experts': 16, 'top_k': 4,
            'expert_dim': 10752, 'hidden': 6144,
            'combos': comb(16, 4),
        },
    }

    for name, cfg in configs.items():
        expert_params = 3 * cfg['hidden'] * cfg['expert_dim']
        active_params = cfg['top_k'] * expert_params
        total_params = cfg['experts'] * expert_params

        bits = log2(cfg['combos'])

        print(f"{name}:")
        print(f"  Combinations: C({cfg['experts']},{cfg['top_k']}) = {cfg['combos']}")
        print(f"  Routing bits: {bits:.1f}")
        print(f"  Expert size: {expert_params/1e6:.0f}M")
        print(f"  Active per token: {active_params/1e6:.0f}M")
        print(f"  Total: {total_params/1e6:.0f}M per layer")

QKV Clipping: A Stability Technique

class DBRXAttention(nn.Module):
    """
    DBRX attention with QKV clipping.
    Clips query, key, value projections to [-clip_qkv, clip_qkv]
    to prevent attention score explosion during training.
    """

    def __init__(self, config):
        super().__init__()
        self.clip_qkv = config.clip_qkv  # 8.0

        self.q_proj = nn.Linear(config.hidden_size,
                                 config.num_attention_heads * config.head_dim,
                                 bias=False)
        self.k_proj = nn.Linear(config.hidden_size,
                                 config.num_key_value_heads * config.head_dim,
                                 bias=False)
        self.v_proj = nn.Linear(config.hidden_size,
                                 config.num_key_value_heads * config.head_dim,
                                 bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * config.head_dim,
                                 config.hidden_size,
                                 bias=False)

    def forward(self, hidden_states, position_ids):
        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)

        # QKV clipping — prevents attention score explosion
        if self.clip_qkv is not None:
            q = q.clamp(-self.clip_qkv, self.clip_qkv)
            k = k.clamp(-self.clip_qkv, self.clip_qkv)
            v = v.clamp(-self.clip_qkv, self.clip_qkv)

        # Standard GQA attention from here
        # (reshape, apply RoPE, compute attention, project output)
        # ...
        return output

Enterprise Training Data

def dbrx_data_composition():
    """
    DBRX was trained with an enterprise focus.
    """
    # Estimated data composition based on Databricks' publications
    data_mix = {
        'web_text_curated': {
            'fraction': 0.35,
            'description': 'High-quality web text, heavily filtered',
        },
        'code': {
            'fraction': 0.20,
            'description': 'GitHub code, enterprise codebases',
        },
        'academic_papers': {
            'fraction': 0.10,
            'description': 'ArXiv, PubMed, academic publications',
        },
        'books': {
            'fraction': 0.08,
            'description': 'Books3, Gutenberg, educational texts',
        },
        'structured_data': {
            'fraction': 0.10,
            'description': 'SQL, JSON, CSV, database schemas',
        },
        'business_documents': {
            'fraction': 0.07,
            'description': 'Reports, contracts, business communications',
        },
        'math_reasoning': {
            'fraction': 0.05,
            'description': 'Mathematical proofs, problem solving',
        },
        'conversational': {
            'fraction': 0.05,
            'description': 'Dialogue, QA pairs, instructions',
        },
    }

    # Key difference from general-purpose models:
    # 10% structured data + 7% business documents = 17% enterprise focus
    # Most open models allocate < 5% to structured/enterprise data
    return data_mix

Serving Performance

def dbrx_serving_analysis():
    """
    DBRX serving characteristics on different hardware.
    """
    total_params = 132e9
    active_params = 36e9

    # Memory requirements
    fp16_mem = total_params * 2 / 1e9  # 264 GB
    int8_mem = total_params * 1 / 1e9  # 132 GB
    int4_mem = total_params * 0.5 / 1e9  # 66 GB

    # KV cache (32K context, GQA with 8 KV heads)
    kv_per_token = 2 * 8 * 128 * 2 * 40  # K+V * heads * dim * FP16 * layers
    kv_32k = kv_per_token * 32768 / 1e9  # GB

    configs = {
        '2x H100-80G (FP16)': {
            'mem': 160,
            'fits': fp16_mem < 160,
            'tps_bs1': 45,
        },
        '4x A100-80G (FP16)': {
            'mem': 320,
            'fits': True,
            'tps_bs1': 38,
        },
        '2x A100-80G (INT4)': {
            'mem': 160,
            'fits': int4_mem + kv_32k < 160,
            'tps_bs1': 42,
        },
        '1x H100-80G (INT4)': {
            'mem': 80,
            'fits': int4_mem + kv_32k < 80,
            'tps_bs1': 32,
        },
    }

    return configs

DBRX vs Competitors

DBRX on Enterprise Tasks

def dbrx_enterprise_benchmarks():
    """
    DBRX was specifically evaluated on enterprise-relevant tasks
    beyond standard academic benchmarks.
    """
    enterprise_results = {
        'sql_generation': {
            'benchmark': 'Spider (SQL)',
            'dbrx': 72.1,
            'mixtral': 65.4,
            'llama2_70b': 62.8,
            'note': 'Structured data training pays off',
        },
        'json_extraction': {
            'benchmark': 'Custom JSON extraction',
            'dbrx': 88.4,
            'mixtral': 81.2,
            'llama2_70b': 78.5,
        },
        'document_qa': {
            'benchmark': 'DocQA enterprise subset',
            'dbrx': 74.8,
            'mixtral': 70.1,
            'llama2_70b': 68.9,
        },
        'code_completion': {
            'benchmark': 'Enterprise Python (internal)',
            'dbrx': 68.2,
            'mixtral': 61.5,
            'llama2_70b': 55.3,
        },
    }
    return enterprise_results

Lessons and Legacy

def dbrx_lessons():
    """
    DBRX's impact on the MoE landscape:

    1. Fine-grained MoE validation: 16 experts with top-4 routing
       outperforms 8 experts with top-2 at similar active compute.
       This finding influenced DeepSeek V2/V3's design.

    2. Enterprise data matters: allocating 17% of training data to
       enterprise-relevant content (SQL, JSON, business docs)
       yields measurable gains on enterprise tasks.

    3. Memory footprint is the bottleneck: DBRX's 132B total params
       require multi-GPU serving, which limits its deployment advantage
       over dense models of similar quality.

    4. Open-source MoE ecosystem: DBRX was the second major open MoE
       after Mixtral, expanding the available model diversity.
    """
    pass

DBRX validated two important ideas: first, that fine-grained MoE (more experts, higher top-k) outperforms coarse-grained MoE at matched active compute; second, that enterprise-focused training data produces meaningfully better performance on business-relevant tasks (SQL, structured data, document understanding). The model’s main limitation was its memory footprint — at 132B total parameters, it required 2-4 GPUs for serving, which undermined the efficiency advantage of MoE for many deployment scenarios. DeepSeek V2 and V3 later demonstrated that fine-grained MoE could be pushed much further (160-256 experts) with better efficiency through architectural innovations like MLA.