Details the forward pass of a MicroLM Transformer model, showing the flow from input_ids through Embedding, Transformer Blocks, MultiHeadSelfAttention, Swi
%%{init: {"theme": "base", "themeVariables": {"background": "#ffffff", "primaryColor": "#f8fafc", "primaryBorderColor": "#94a3b8", "primaryTextColor": "#0f172a", "lineColor": "#64748b"}}}%%
flowchart TB
I["input_ids(B, T)"] --> E["Embedding<br>(B, T, d_model)"]
E --> B["TransformerBlock × 8"]
B --> N1["RMSNorm"]
N1 --> A["MultiHeadSelfAttention<br>Q/K/V · RoPE · causal mask"]
A --> R1["Residual Add"]
R1 --> N2["RMSNorm"]
N2 --> F["SwiGLU FFN"]
F --> R2["Residual Add"]
R2 --> FN["Final RMSNorm"]
FN --> H["lm_head"]
H --> L["logits<br>(B, T, vocab_size)"]
classDef input fill:#eff6ff,stroke:#60a5fa,color:#0f172a;
classDef core fill:#f8fafc,stroke:#94a3b8,color:#0f172a;
classDef attn fill:#f0fdf4,stroke:#22c55e,color:#0f172a;
classDef ffn fill:#fff7ed,stroke:#fb923c,color:#0f172a;
classDef out fill:#fdf4ff,stroke:#d946ef,color:#0f172a;
class I,E input;
class B,N1,N2,R1,R2,FN core;
class A attn;
class F ffn;
class H,L out;This flowchart illustrates the forward pass of a MicroLM Transformer model. It begins with input_ids, which are processed by an Embedding layer. The embedded tokens then pass through 8 Transformer Blocks, each featuring RMSNorm, MultiHeadSelfAttention with RoPE and causal masking, and a SwiGLU Feed-Forward Network, all connected with residual additions. A final RMSNorm and lm_head produce the output logits.
Use this diagram to understand the sequential data flow and key components within a modern, simplified Transformer architecture, such as MicroLM. It's ideal for learning about the inference or training path of an LLM, specifically how input tokens are transformed into output logits through attention, FFNs, and normalization layers.
To adapt this diagram, one could modify the number of Transformer Blocks, specify different attention mechanisms, or detail the internal operations of RoPE or SwiGLU. It can also be extended to show specific data shapes at each stage, or to integrate with other components like KV Cache for inference optimization.