Architectures / NLP/LLM
🔀 Mixtral MoE Block
Mixtral decoder block — GQA + Sparse MoE (8 experts, top-2) + RMSNorm + RoPE
Layers
9
Parameters
1.45B
Input
1 × 4096 × 4096
Output
1 × 4096 × 4096
Verifier
Clean
Every number on this page is computed from the graph by the same functions the app runs, not written by hand.
Open Mixtral MoE Block on the canvas
Free, no account needed
When to pick it
Pick when you have MoE training infra and want best quality per active param. Routing instability and memory cost (still scales with total params) are the trade-off.
Structure
9 layers. Output shapes are propagated from the input shape, batch dimension excluded.
| Layer | Type | Parameters | Output shape | |
|---|---|---|---|---|
| 1 | hidden_states | Input | shape=[1, 4096, 4096] | 1 × 4096 × 4096 |
| 2 | input_norm | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 3 | self_attn | Grouped Query Attn | embedDim=4096, numHeads=32, numKVHeads=8 | 1 × 4096 × 4096 |
| 4 | rotary_emb | RoPE | ||
| 5 | attn_residual | Add | 1 × 4096 × 4096 | |
| 6 | post_attn_norm | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 7 | block_sparse_moe | MoE Layer | embedDim=4096, numExperts=8, topK=2 | 1 × 4096 × 4096 |
| 8 | moe_residual | Add | 1 × 4096 × 4096 | |
| 9 | hidden_out | Output | 1 × 4096 × 4096 |
What the verifier says
The same 41 structural checks that run on every edit in the app, on this graph.
infoMoE layers require an auxiliary router z-loss + load-balance loss during training to prevent expert collapse. This is not visible in the architecture diagram but must be in the training loop. Fix: Add a note on this layer. Typical aux_loss coefficient: 1e-2 (Mixtral/Switch Transformer). (block_sparse_moe)
moe-no-aux-loss
moe-no-aux-loss
The PyTorch it exports
Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision (usually pre-installed)
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
class MixtralMoEBlock(nn.Module):
def __init__(self):
super().__init__()
self.rmsNorm_1 = nn.RMSNorm(4096)
self.groupedQueryAttention_1 = nn.ModuleDict({
'q_proj': nn.Linear(4096, 4096, bias=False), # 32 heads × 128
'k_proj': nn.Linear(4096, 1024, bias=False), # 8 KV heads × 128
'v_proj': nn.Linear(4096, 1024, bias=False),
'o_proj': nn.Linear(4096, 4096, bias=False),
}) # GQA: 32Q / 8KV heads (requires F.scaled_dot_product_attention)
self.rmsNorm_2 = nn.RMSNorm(4096)
self.moeLayer_1 = nn.ModuleDict({
'router': nn.Linear(4096, 8, bias=False),
'experts': nn.ModuleList([
nn.Sequential(
nn.Linear(4096, 14336, bias=False), nn.SiLU(),
nn.Linear(14336, 4096, bias=False),
) for _ in range(8)
]),
}) # MoE top-2
@staticmethod
def _moe_forward(moe, x, top_k=2):
"""Top-k routed MoE forward over a {'router', 'experts'} ModuleDict."""
scores = moe['router'](x).softmax(dim=-1)
top_w, top_i = scores.topk(top_k, dim=-1)
top_w = top_w / top_w.sum(dim=-1, keepdim=True)
flat_x = x.reshape(-1, x.size(-1))
flat_i = top_i.reshape(-1, top_k)
flat_w = top_w.reshape(-1, top_k)
out = torch.zeros_like(flat_x)
for e_idx in flat_i.unique():
hit = flat_i == e_idx
rows = hit.any(dim=-1)
w = (flat_w * hit).sum(dim=-1)[rows].unsqueeze(-1)
out[rows] += w * moe['experts'][int(e_idx)](flat_x[rows])For agents
This architecture is machine-readable end to end. An agent can list the set, fetch this graph, edit it, and have the edit verified before any GPU time is spent.