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 rotaryDim dimensions 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]*sin q_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_layernorm and post_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_head also has a bias.

  • Untied output head (lm_head.weight V, 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