GPT-2

All models, with their blocks →

Size
Published
OpenAI · Feb–Nov 2019
Line
Dense decoder
View
Architecture
Source

openai/gpt‑2 @ 9b63575, huggingface/transformers @ 71c6f69, configuration openai‑community/gpt2 @ 607a30d · checkpoint tensors not independently countedHow to read

Prepare Q, K, V

Full sequenceTrainingPrefillDecode

Transformer block › Attention

Reshape into attention heads

×

no parametersmap(split_heads, tf.split(c, 3, axis=2))

Reshape and permute Q, K, and V independently into 12 heads of width 64.

[B,T,768][B,T,768]
→  [B,T,12,64]\to\;[B,T,12,64]
→  [B,12,T,64]\to\;[B,12,T,64]
in
QQ [B batch, T positions, 768 d]KK [B batch, T positions, 768 d]VV [B batch, T positions, 768 d]
out
Q/K/VQ / K / V [B batch, 12 heads, T positions, 64 d_h]
Codemodel.py 75, 77
model.py 75
    def split_heads(x):
model.py 77
        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])
Sources1

Direct evidence

The released split_heads function introduces a head axis and moves it before the position axis.

openai · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 75–79

Read code excerpt
    def split_heads(x):

… lines 76–76 omitted …

        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])