core/nn/llama library

Llama-style causal language model.

Architectural pieces vs. GPT (GPT-2) / PythiaModel (GPT-NeoX):

  • Pre-RMSNorm on both attention and FFN sub-blocks (no bias, no mean-subtract) — see RMSNorm.
  • RoPE full-rotation applied per KV/Q head, base is 500_000.0 for Llama 3 (10_000.0 for Llama 1/2).
  • Grouped-Query Attention — numKvHeads is separate from numHeads (see MultiHeadAttention). For Llama 3.2 1B/3B and Llama 3 8B this is 8 vs. numHeads == 32.
  • SwiGLU FFN — three bias-free Linears (gate_proj, up_proj, down_proj), see SwiGluFfn.
  • Sequential residual — x + attn(rmsNorm(x)) then x + swiglu(rmsNorm(x)) (Pythia is parallel).
  • Weight-tied lm_head by default (Llama 3.2 1B/3B tie embed_tokens ↔ lm_head; Llama 3 8B does not — pass tieWeights: false).
  • No biases anywhere.

Same 2D single-sequence convention as GPT.generate / PythiaModel.