core/nn/pythia library

GPT-NeoX / Pythia causal language model.

The main structural differences vs. our GPT (GPT-2):

  • No learned positional embedding table. Positions are injected via rotary embeddings (RopeCache) applied to Q and K inside each attention layer.
  • Parallel residual block — instead of the sequential x -> ln1 -> attn -> +x -> ln2 -> mlp -> +x GPT-2 uses, each block computes attn and mlp in parallel from independent LayerNorms of the same input: y = x + attn(input_layernorm(x)) + mlp(post_attention_layernorm(x)) (this is HF's use_parallel_residual: True mode).
  • Untied output head. Pythia keeps a separate embed_out [vocab, hidden] rather than tying to embed_in.

Everything else (biased Q/K/V/output projections, biased 4x MLP, causal mask, KV cache, sampling loop) matches our GPT-2 stack.

The activation used here is the tanh approximation of GELU (Activation.geluTanh). Pythia originally uses exact GELU (erf-based), but the two agree to within a few parts in 1e4 — perfectly fine for inference next-token argmax.