core/nn/gptj library
GPT-J causal language model (EleutherAI, 2021).
GPT-J was the first widely-used open replica of a GPT-3-class model. Architecturally it sits between GPT-2 and GPT-NeoX:
-
Rotary positional embedding, applied only to the first
rotaryDimdimensions of each head (64 out of 256 for the 6B model). Uses the interleaved-pair convention:q_rot[..., 2i] = q[..., 2i]*cos - q[..., 2i+1]*sinq_rot[..., 2i+1] = q[..., 2i]*sin + q[..., 2i+1]*cos. This is different from GPT-NeoX / LLaMA / Pythia which use "half-split" rotary. The GPTJHFLoader handles this by permuting the Q/K weight rows at load time so our shared half-split RopeCache produces the mathematically identical result inside attention (see loader for the derivation). -
Parallel residual block with a single shared LayerNorm. Where Pythia has two LNs (
input_layernormandpost_attention_layernorm), GPT-J uses one:h = ln(x) y = x + attn(h) + mlp(h)Both attn and mlp read from the same normalised input, and both are added into the residual — an important departure from GPT-2's sequential
x -> attn -> +x -> mlp -> +x. -
No bias on attention Q/K/V/out projections. The MLP retains biases, and (unlike Pythia) the untied
lm_headalso has a bias. -
Untied output head (
lm_head.weightV, D+lm_head.bias). -
Uses the standard GPT-2 byte-level BPE tokenizer, but with the vocabulary padded to 50400 for GPU-friendly matmul shapes.
The activation is the tanh approximation of GELU (gelu_new),
same as Pythia and GPT-2.
Classes
- GPTJBlock
- GPT-J transformer block — parallel residual with a single shared LN.
- GPTJConfig
- GPTJModel