Llama 3

All models, with their blocks →

Size
Published
Meta · Apr 2024
Line
Dense decoder
View
Architecture
Source

meta‑llama/llama3 @ a0940f9, configuration meta‑llama/llama‑models @ 0e0b8c5 · checkpoint, tokenizer and tensors not verifiedHow to read

Grouped-query attention

Full sequencePrefillDecode

32 independent transformer blocks › Attention

Row-wise softmax

×

no parametersF.softmax(scores.float(), dim=-1).type_as(xq)

Convert each query row into normalized attention weights.

A[b,i,t,s]=exp⁡(Ss[b,i,t,s])∑u=01000exp⁡(Ss[b,i,t,u]){\htmlData{light=weights softmax}{A}}[b,i,t,s]=\frac{\exp({\htmlData{light=scale}{S_{\mathrm s}}}[b,i,t,s])}{\sum_{u=0}^{1000}\exp({\htmlData{light=scale}{S_{\mathrm s}}}[b,i,t,u])}
in
SsS_{\mathrm s} [B batch, 32 q heads, 1 query, 1001 keys]
out
AA [B batch, 32 q heads, 1 query, 1001 keys]

In float32, then cast back to the query type.

Explanation

bb, ii, tt and ss index batch, head, query and key; uu runs over the keys.

Codemodel.py 187
model.py 187
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
Sources1

Direct evidence

Softmax normalizes the last score axis in float32 and restores the query dtype.

attn-softmax · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 187–187

Read code excerpt
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)