GPT-2

All models, with previews →

OpenAI · Feb–Nov 2019Dense decoder

Prepare Q, K, V

Reference code openai/gpt-2 @ 9b63575, huggingface/transformers @ 71c6f69, configuration openai-community/gpt2 @ 607a30d · checkpoint tensors not independently counted · How to read

Full sequenceTrainingPrefillDecode

Transformer block › Attention

Q, K, V split

×

no parameterstf.split(c, 3, axis=2)

Split the fused 2,304 features into three 768-feature tensors: queries, keys, and values.

Q,K,V=split⁡(qkv,3){\htmlData{light=q qkv-fork}{{\htmlClass{vl-stream-q}{Q}}}},{\htmlData{light=k qkv-fork}{{\htmlClass{vl-stream-k}{K}}}},{\htmlData{light=v qkv-fork}{{\htmlClass{vl-stream-v}{V}}}}=\operatorname{split}(\mathrm{qkv},3)
in
qkv\mathrm{qkv} [B batch, T positions, 2304 q, k, v]
out
Q/K/VQ / K / V [B batch, T positions, 768 d]
Codemodel.py 75–103
model.py 75
    def split_heads(x):
model.py 77
        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])
model.py 103
        q, k, v = map(split_heads, tf.split(c, 3, axis=2))
Sources1

Direct evidence

The fused field is split three ways and each result is reshaped into the head axis.

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

Read code excerpt
    def split_heads(x):

… lines 76–76 omitted …

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

… lines 78–102 omitted …

        q, k, v = map(split_heads, tf.split(c, 3, axis=2))