Mixtral 8×7B — The Evolving Transformer

Mixtral 8×7B · 32× blocks, SWA W=4096, GQA 32/8, MoE 8 experts top-2

This chapter evolves chapter 04 Mistral 7B into Mixtral 8×7B — same decoder-only stack (wte-only embeddings, sliding-window GQA, RoPE, RMSNorm, weight tying), with one key innovation: sparse mixture-of-experts (MoE) in every block. Each token is routed to top-2 of 8 SwiGLU experts instead of one dense FFN. Hyperparameters follow Mixtral 8×7B from Mistral AI. Same teaching pattern: one config object up front, then each layer in order until the full model runs.

Each section below is one building block. You get a short explanation, the code for that module, a delta from chapter 04 where it helps, and how it connects to what came before. By the end you have a runnable Mixtral class, not just a diagram. The 8×7B recipe stores ~46.7B parameters total while activating ~12.9B per forward pass (top-2 routing across eight SwiGLU experts per layer).

Delta from Mistral 7B (04)

Dense SwiGLU FFN → sparse MoE (8 SwiGLU experts, top-2 router per token)
SwiGLUFeedForward kept as each expert's building block
Attention unchanged: sliding window GQA, RoPE, RMSNorm
~12.9B active params per forward, ~46.7B total stored (diagram + HF Mixtral-8×7B-v0.1)

Full side-by-side with chapter 04: §08 Comparison.

00 · CONFIG

One config object for the whole model

Before any module, we define a single MixtralConfig dataclass for Mixtral 8×7B. Every class below takes config instead of a long list of constructor arguments — same pattern as chapter 04, plus MoE fields: num_experts and top_k.

Mixtral 8×7B uses d_model=4096, N=32 blocks, H=32 query heads, n_kv_heads=8 (4:1 GQA), vocab=32,000, context length 32,768, d_ff=14,336 per expert, sliding window W=4096, 8 experts, top-2 routing — matching the architecture figure on this page (~46.7B total, ~12.9B active). Attention stack matches Mistral 7B; MoE is the architectural delta.

dk property: dmodel // H → 4096 // 32 = 128. Multi-head attention splits d_model across H query heads.

make_sliding_window_mask: same helper as Mistral 7B — causal ∩ sliding window band (used by SlidingWindowGQA, unchanged from chapter 04):

 36def make_sliding_window_mask(max_seq_len, sliding_window):
 37    """Causal mask ∩ sliding window band: token i attends to [max(0, i-W+1) .. i]."""
 38    causal = torch.tril(torch.ones(max_seq_len, max_seq_len))
 39    window = torch.triu(torch.ones(max_seq_len, max_seq_len), diagonal=-(sliding_window - 1))
 40    return causal * window  # (max_seq_len, max_seq_len)

max_seq_len: 32,768 for Mixtral 8×7B (vs 8,192 in Mistral 7B) — longer RoPE tables and mask buffer.

All values below are Mixtral 8×7B — verified against Meta’s mistralai/Mixtral-8x7B-v0.1 config.json (hidden_size, num_local_experts, num_experts_per_tok, etc.).

dmodel4096 — hidden_size

N32 — num_hidden_layers

H32 — num_attention_heads

n_kv_heads8 — num_key_value_heads (GQA)

vocabsize32,000 — vocab_size

max_seq_len32,768 — max_position_embeddings

dff14,336 — intermediate_size per expert

sliding_window4096 — SWA band (diagram; HF v0.1 has null)

num_experts8 — num_local_experts

top_k2 — num_experts_per_tok

norm_eps1e-5 — RMSNorm epsilon

B2 — batch size for demos

dk128 — dmodel // H

mixtral.pyimports + MixtralConfig

  1from dataclasses import dataclass
  2import torch
  3import torch.nn as nn
  4import torch.nn.functional as F
  5import math
  6
  7
  8# ============================================================
  9# Config — Mixtral 8×7B from Jiang et al., 2024 (Mistral AI)
 10# Same stack as Mistral 7B + sparse MoE FFN (8 experts, top-2 per token)
 11# ============================================================
 12@dataclass
 13class MixtralConfig:
 14    dmodel: int = 4096           # Mixtral 8×7B
 15    N: int = 32                  # transformer blocks
 16    H: int = 32                  # query heads
 17    n_kv_heads: int = 8          # GQA — 4:1 ratio (32 Q / 8 KV)
 18    vocabsize: int = 32000       # BPE vocabulary size
 19    max_seq_len: int = 32768     # max context length (32k)
 20    dff: int = 14336             # SwiGLU inner dim per expert
 21    norm_eps: float = 1e-5       # RMSNorm epsilon
 22    sliding_window: int = 4096   # attention window W
 23    num_experts: int = 8         # SwiGLU experts per layer
 24    top_k: int = 2               # experts activated per token
 25    B: int = 2                   # batch size (for demo runs)
 26
 27    @property
 28    def dk(self):
 29        # per-head dimension — dmodel must be divisible by H
 30        return self.dmodel // self.H
 31
 32
 33config = MixtralConfig()  # Mixtral 8×7B — shared by every module below

01 · EMBEDDINGS

Token embedding only — unchanged from Mistral

Mixtral keeps the same wte-only input as chapter 04: a token embedding table mapping IDs to d_model vectors, no learned position table, no dropout. Position is handled by RoPE inside attention.

The forward pass is simply wte(x) — one lookup, straight into the block stack. No delta from Mistral 7B here.

mixtral.pywte only

  1# inside Mixtral.__init__ — token embedding only (no wpe):
186        self.wte = nn.Embedding(config.vocabsize, config.dmodel)  # wte(x): (B, S, dmodel)
  3
  4# inside Mixtral.forward — no positional sum, no dropout:
198        _, S = x.shape  # x: (B, S)
199        x = self.wte(x)  # (B, S, dmodel)

▼ Show original version

Mistral 7B in chapter 04 — identical wte pattern (excerpt from mistral.py).

mistral.pywte only

  1# inside Mistral.__init__ — token embedding only (no wpe):
  2        self.wte = nn.Embedding(config.vocabsize, config.dmodel)  # wte(x): (B, S, dmodel)
  3
  4# inside Mistral.forward — no positional sum, no dropout:
  5        _, S = x.shape  # x: (B, S)
  6        x = self.wte(x)  # (B, S, dmodel) — no wpe; RoPE handles position in attention

02 · RoPE

Rotary positional embeddings — unchanged from Mistral

RoPE encodes position by rotating Q and K vectors in each head's 2D subspaces. Cos/sin tables are precomputed up to max_seq_len and registered as buffers — same implementation as chapter 04 (now sized for 32k context).

No delta from Mistral 7B. Sliding-window GQA limits which keys each query sees; RoPE still rotates Q and K the same way.

mixtral.pyRotaryPositionalEmbedding

107class RotaryPositionalEmbedding(nn.Module):
108    def __init__(self, config, base: float = 10000.0) -> None:
109        super().__init__()
110        self.dk = config.dk
111        freqs = 1.0 / (
112            base ** (torch.arange(0, self.dk, 2).float() / self.dk)
113        )  # (dk//2,)
114        pos = torch.arange(0, config.max_seq_len).float()  # (max_seq_len,)
115        angles = torch.outer(pos, freqs)  # (max_seq_len, dk//2)
116        self.register_buffer("cos_cached", torch.cos(angles))  # (max_seq_len, dk//2)
117        self.register_buffer("sin_cached", torch.sin(angles))  # (max_seq_len, dk//2)
118
119    def forward(self, x):
120        B, H, S, dk = x.shape  # x: (B, H, S, dk)
121        x_even = x[..., 0::2]  # (B, H, S, dk//2)
122        x_odd = x[..., 1::2]  # (B, H, S, dk//2)
123        cos = self.cos_cached[:S].unsqueeze(0).unsqueeze(0)  # (1, 1, S, dk//2)
124        sin = self.sin_cached[:S].unsqueeze(0).unsqueeze(0)  # (1, 1, S, dk//2)
125        rotated_even = x_even * cos - x_odd * sin  # (B, H, S, dk//2)
126        rotated_odd = x_even * sin + x_odd * cos  # (B, H, S, dk//2)
127        rotated = torch.stack([rotated_even, rotated_odd], dim=-1)  # (B, H, S, dk//2, 2)
128        return rotated.view(B, H, S, dk)  # (B, H, S, dk)

▼ Show original version

Mistral 7B in chapter 04 — same RotaryPositionalEmbedding (excerpt from mistral.py).

mistral.pyRotaryPositionalEmbedding

 65class RotaryPositionalEmbedding(nn.Module):
 66    def __init__(self, config, base: float = 10000.0) -> None:
 67        super().__init__()
 68        self.dk = config.dk
 69        freqs = 1.0 / (
 70            base ** (torch.arange(0, self.dk, 2).float() / self.dk)
 71        )  # (dk//2,)
 72        pos = torch.arange(0, config.max_seq_len).float()  # (max_seq_len,)
 73        angles = torch.outer(pos, freqs)  # (max_seq_len, dk//2)
 74        self.register_buffer("cos_cached", torch.cos(angles))  # (max_seq_len, dk//2)
 75        self.register_buffer("sin_cached", torch.sin(angles))  # (max_seq_len, dk//2)
 76
 77    def forward(self, x):
 78        B, H, S, dk = x.shape  # x: (B, H, S, dk)
 79        x_even = x[..., 0::2]  # (B, H, S, dk//2)
 80        x_odd = x[..., 1::2]  # (B, H, S, dk//2)
 81        cos = self.cos_cached[:S].unsqueeze(0).unsqueeze(0)  # (1, 1, S, dk//2)
 82        sin = self.sin_cached[:S].unsqueeze(0).unsqueeze(0)  # (1, 1, S, dk//2)
 83        rotated_even = x_even * cos - x_odd * sin  # (B, H, S, dk//2)
 84        rotated_odd = x_even * sin + x_odd * cos  # (B, H, S, dk//2)
 85        rotated = torch.stack([rotated_even, rotated_odd], dim=-1)  # (B, H, S, dk//2, 2)
 86        return rotated.view(B, H, S, dk)  # (B, H, S, dk)

03 · RMSNORM

Root-mean-square normalization — unchanged from Mistral

RMSNorm scales each token vector by the reciprocal root-mean-square of its features, then multiplies by a learned gamma. No mean centering, no bias — same as chapter 04.

No delta from Mistral 7B.

mixtral.pyRMSNorm

 43class RMSNorm(nn.Module):
 44    def __init__(self, config) -> None:
 45        super().__init__()
 46        self.norm_eps = config.norm_eps
 47        self.gamma = nn.Parameter(torch.ones(config.dmodel))  # (dmodel,)
 48
 49    def forward(self, x):
 50        # x: (B, S, dmodel)
 51        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.norm_eps)  # (B, S, 1)
 52        return x * rms * self.gamma  # (B, S, dmodel)

▼ Show original version

Mistral 7B in chapter 04 — same RMSNorm (excerpt from mistral.py).

mistral.pyRMSNorm

 41class RMSNorm(nn.Module):
 42    def __init__(self, config) -> None:
 43        super().__init__()
 44        self.norm_eps = config.norm_eps
 45        self.gamma = nn.Parameter(torch.ones(config.dmodel))  # (dmodel,)
 46
 47    def forward(self, x):
 48        # x: (B, S, dmodel)
 49        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.norm_eps)  # (B, S, 1)
 50        return x * rms * self.gamma  # (B, S, dmodel)

04 · SWIGLU

SwiGLU feed-forward — each MoE expert uses this block

Mixtral reuses the same SwiGLU FFN as Mistral 7B: three bias-free matrices (w1 gate, w2 up, w3 down) with activation silu(w1(x)) * w2(x). Inner dimension dff=14336 per expert comes from config.

The module is unchanged from chapter 04 — the delta is how many copies exist: eight experts in MoE (section 05) instead of one dense self.ff per block.

mixtral.pySwiGLUFeedForward

 55class SwiGLUFeedForward(nn.Module):
 56    def __init__(self, config) -> None:
 57        super().__init__()
 58        self.w1 = nn.Linear(config.dmodel, config.dff, bias=False)  # gate
 59        self.w2 = nn.Linear(config.dmodel, config.dff, bias=False)  # up
 60        self.w3 = nn.Linear(config.dff, config.dmodel, bias=False)  # down
 61
 62    def forward(self, x):
 63        # x: (B, S, dmodel) → silu(w1) * w2 → w3 → (B, S, dmodel)
 64        return self.w3(F.silu(self.w1(x)) * self.w2(x))

▼ Show original version

Mistral 7B in chapter 04 — same SwiGLUFeedForward (excerpt from mistral.py).

mistral.pySwiGLUFeedForward

 53class SwiGLUFeedForward(nn.Module):
 54    def __init__(self, config) -> None:
 55        super().__init__()
 56        self.w1 = nn.Linear(config.dmodel, config.dff, bias=False)  # gate
 57        self.w2 = nn.Linear(config.dmodel, config.dff, bias=False)  # up
 58        self.w3 = nn.Linear(config.dff, config.dmodel, bias=False)  # down
 59
 60    def forward(self, x):
 61        # x: (B, S, dmodel) → silu(w1) * w2 → w3 → (B, S, dmodel)
 62        return self.w3(F.silu(self.w1(x)) * self.w2(x))

05 · MOE ROUTER + MOE

Sparse mixture-of-experts — the main delta from Mistral

This is the main delta from chapter 04. MoERouter scores each token against 8 experts, softmaxes, takes top-2, and renormalizes weights on the active set. MoE holds eight SwiGLUFeedForward experts and accumulates weighted outputs for tokens routed to each expert.

Training-style forward here loops over expert slots and experts with boolean masks — clear for teaching. Production stacks batch tokens per expert for efficiency.

y = ∑_{k=1..top_k} w_k · E_{i_k}(x) · i_k ∈ TopK(softmax(G(x)))

mixtral.pyMoERouter + MoE

 67class MoERouter(nn.Module):
 68    def __init__(self, config) -> None:
 69        super().__init__()
 70        self.num_experts = config.num_experts
 71        self.top_k = config.top_k
 72        self.router = nn.Linear(config.dmodel, config.num_experts, bias=False)  # (dmodel,) → (num_experts,)
 73
 74    def forward(self, x):
 75        # x: (B, S, dmodel)
 76        scores = self.router(x)  # (B, S, num_experts)
 77        prob = F.softmax(scores, dim=-1)  # (B, S, num_experts)
 78        weight, indices = torch.topk(prob, self.top_k, dim=-1)  # (B, S, top_k) each
 79        weight = weight / weight.sum(dim=-1, keepdim=True)  # renormalize over active experts
 80        return weight, indices  # (B, S, top_k)
 81
 82
 83class MoE(nn.Module):
 84    def __init__(self, config) -> None:
 85        super().__init__()
 86        self.experts = nn.ModuleList(
 87            [SwiGLUFeedForward(config) for _ in range(config.num_experts)]
 88        )
 89        self.router = MoERouter(config)
 90        self.top_k = config.top_k
 91
 92    def forward(self, x):
 93        # x: (B, S, dmodel)
 94        weights, indices = self.router(x)  # (B, S, top_k)
 95        output = torch.zeros_like(x)  # (B, S, dmodel)
 96
 97        for k in range(self.top_k):
 98            for i, expert in enumerate(self.experts):
 99                mask = indices[:, :, k] == i  # (B, S) — tokens routed to expert i at slot k
100                if mask.any():
101                    expert_out = expert(x[mask])  # (num_selected, dmodel)
102                    selected_weights = weights[:, :, k][mask].unsqueeze(-1)  # (num_selected, 1)
103                    output[mask] += selected_weights * expert_out
104        return output  # (B, S, dmodel)

▼ Show original version

Mistral 7B in chapter 04: one dense self.ff per block — no router (excerpt from mistral.py).

mistral.pyBlock — dense ff

126class Block(nn.Module):
127    def __init__(self, config) -> None:
128        super().__init__()
129        self.gqa = SlidingWindowGQA(config)
130        self.ff = SwiGLUFeedForward(config)
131        self.rms_1 = RMSNorm(config)  # pre-norm before attention
132        self.rms_2 = RMSNorm(config)  # pre-norm before FFN
133
134    def forward(self, x, rope, mask):
135        # x: (B, S, dmodel)
136        x = x + self.gqa(self.rms_1(x), rope, mask)  # (B, S, dmodel)
137        x = x + self.ff(self.rms_2(x))  # (B, S, dmodel)
138        return x  # (B, S, dmodel)

06 · BLOCK

RMSNorm + SlidingWindowGQA + RMSNorm + MoE

Block structure matches Mistral 7B: two RMSNorms, sliding-window GQA attention, pre-norm residuals. The only swap is self.ff → self.moe — attention still takes (x, rope, mask).

mixtral.pyBlock

168class Block(nn.Module):
169    def __init__(self, config) -> None:
170        super().__init__()
171        self.gqa = SlidingWindowGQA(config)
172        self.moe = MoE(config)  # sparse MoE replaces dense SwiGLU FFN
173        self.rms_1 = RMSNorm(config)  # pre-norm before attention
174        self.rms_2 = RMSNorm(config)  # pre-norm before MoE
175
176    def forward(self, x, rope, mask):
177        # x: (B, S, dmodel)
178        x = x + self.gqa(self.rms_1(x), rope, mask)  # (B, S, dmodel)
179        x = x + self.moe(self.rms_2(x))  # (B, S, dmodel)
180        return x  # (B, S, dmodel)

▼ Show original version

Mistral 7B in chapter 04: same block layout with dense self.ff (excerpt from mistral.py).

mistral.pyBlock

126class Block(nn.Module):
127    def __init__(self, config) -> None:
128        super().__init__()
129        self.gqa = SlidingWindowGQA(config)
130        self.ff = SwiGLUFeedForward(config)
131        self.rms_1 = RMSNorm(config)  # pre-norm before attention
132        self.rms_2 = RMSNorm(config)  # pre-norm before FFN
133
134    def forward(self, x, rope, mask):
135        # x: (B, S, dmodel)
136        x = x + self.gqa(self.rms_1(x), rope, mask)  # (B, S, dmodel)
137        x = x + self.ff(self.rms_2(x))  # (B, S, dmodel)
138        return x  # (B, S, dmodel)

07 · ASSEMBLE MIXTRAL

Token embed, RoPE, N blocks, final norm, lm_head

Mixtral owns everything. wte maps token IDs to vectors. A single RotaryPositionalEmbedding is shared across all blocks. N Blocks via nn.ModuleList, each taking (x, rope, mask). Final RMSNorm then lm_head projects to vocab with weight tying.

Attention mask and RoPE match chapter 04; MoE inside every block is the delta. With the §00 defaults, that is the full Mixtral 8×7B stack.

mixtral.pyMixtral

183class Mixtral(nn.Module):
184    def __init__(self, config) -> None:
185        super().__init__()
186        self.wte = nn.Embedding(config.vocabsize, config.dmodel)  # wte(x): (B, S, dmodel)
187        self.rope = RotaryPositionalEmbedding(config)
188        self.blocks = nn.ModuleList([Block(config) for _ in range(config.N)])
189        self.norm = RMSNorm(config)  # (B, S, dmodel) → (B, S, dmodel)
190        self.lm_head = nn.Linear(config.dmodel, config.vocabsize, bias=False)
191        self.lm_head.weight = self.wte.weight  # weight tying
192        self.register_buffer(
193            "mask",
194            make_sliding_window_mask(config.max_seq_len, config.sliding_window),
195        )  # (max_seq_len, max_seq_len)
196
197    def forward(self, x):
198        _, S = x.shape  # x: (B, S)
199        x = self.wte(x)  # (B, S, dmodel)
200        for block in self.blocks:
201            x = block(x, self.rope, self.mask)  # (B, S, dmodel)
202        x = self.norm(x)  # (B, S, dmodel)
203        return self.lm_head(x)  # (B, S, vocabsize)

08 · COMPARISON

Mistral 7B (04) vs Mixtral 8×7B (05)

Canonical comparison for this chapter — same pair as the delta card above, with code-level detail after you have read the stack.

      
        AxisMistral 7B (04)Mixtral 8×7B (05)

        BackboneDecoder-only · RoPE · RMSNorm · SwiGLU · GQA · SWASame attention stack
FFN per layerOne dense SwiGLUFeedForward (self.ff)8 SwiGLU experts + top-2 MoERouter (self.moe)
Active FFN paramsAll FFN weights every token2 of 8 experts per token (~12.9B active)
Total parameters~7.3B (dense 7B model)~46.7B stored, sparse activation
AttentionSlidingWindowGQA, W=4096Unchanged from Mistral 7B
Position encodingRoPE on Q/KRoPE on Q/K (unchanged)
Embeddingswte only, no wpewte only (unchanged)
NormalizationRMSNorm, bias=FalseRMSNorm (unchanged)
GQA ratio32 Q heads · 8 KV heads (4:1)32 Q heads · 8 KV heads (unchanged)
Output headlm_head weight-tied to wteWeight tying (unchanged)
Context length8,192 in mistral.py32,768 (max_position_embeddings)
Vocabulary32,00032,000
dmodel / dff4096 / 14,3364096 / 14,336 per expert
SwiGLU moduleSwiGLUFeedForwardSame class — one copy per expert

      
    

Axis	Mistral 7B (04)	Mixtral 8×7B (05)
Backbone	Decoder-only · RoPE · RMSNorm · SwiGLU · GQA · SWA	Same attention stack
FFN per layer	One dense `SwiGLUFeedForward` (`self.ff`)	8 SwiGLU experts + top-2 `MoERouter` (`self.moe`)
Active FFN params	All FFN weights every token	2 of 8 experts per token (~12.9B active)
Total parameters	~7.3B (dense 7B model)	~46.7B stored, sparse activation
Attention	`SlidingWindowGQA`, W=4096	Unchanged from Mistral 7B
Position encoding	RoPE on Q/K	RoPE on Q/K (unchanged)
Embeddings	`wte` only, no `wpe`	`wte` only (unchanged)
Normalization	RMSNorm, `bias=False`	RMSNorm (unchanged)
GQA ratio	32 Q heads · 8 KV heads (4:1)	32 Q heads · 8 KV heads (unchanged)
Output head	`lm_head` weight-tied to `wte`	Weight tying (unchanged)
Context length	8,192 in `mistral.py`	32,768 (`max_position_embeddings`)
Vocabulary	32,000	32,000
d_model / d_ff	4096 / 14,336	4096 / 14,336 per expert
SwiGLU module	`SwiGLUFeedForward`	Same class — one copy per expert

HF mistralai/Mixtral-8x7B-v0.1 sets sliding_window to null (full causal in that checkpoint). This chapter keeps W=4096 in mixtral.py to match the diagram and chapter 04’s sliding-window teaching path.

09 · REFERENCES

Papers & technical sources

Primary reports and references for this chapter. Read these for full equations, training details, and official hyperparameters.

Mixtral of Experts (Jiang et al., 2024)Sparse MoE FFN — 8 experts, top-2 routing, same attention stack as Mistral 7B.
Mixtral 8×7B — Mistral AIRelease notes and active vs total parameter counts.
Hugging Face: mistralai/Mixtral-8x7B-v0.1config.json for hidden_size, num_local_experts, sliding_window.

Delta from Mistral 7B (04)

One config object for the whole model

Token embedding only — unchanged from Mistral

Rotary positional embeddings — unchanged from Mistral

Root-mean-square normalization — unchanged from Mistral

SwiGLU feed-forward — each MoE expert uses this block

Sparse mixture-of-experts — the main delta from Mistral

RMSNorm + SlidingWindowGQA + RMSNorm + MoE

Token embed, RoPE, N blocks, final norm, lm_head

Mistral 7B (04) vs Mixtral 8×7B (05)

Papers & technical sources

Comments