32 independent transformer blocks › Attention
Scale similarities
×no parameters/ math.sqrt(self.head_dim)
Control the magnitude of similarities entering softmax.
- in
- [B batch, 32 q heads, T queries, T keys]
- out
- [B batch, 32 q heads, T queries, T keys]
Explanation
For independent, zero-mean, unit-variance query and key features, a 128-term dot product has variance 128; dividing by returns it to 1, so the scores keep the same scale at any head width. Learned features need not meet these assumptions.
Codemodel.py 184 · sku_list.py 19, 231–248
scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)LLAMA3_VOCAB_SIZE = 128256 Model(
core_model_id=CoreModelId.llama3_1_8b,
description="Llama 3.1 8b model",
huggingface_repo="meta-llama/Llama-3.1-8B",
arch_args={
"dim": 4096,
"n_layers": 32,
"n_heads": 32,
"n_kv_heads": 8,
"vocab_size": LLAMA3_VOCAB_SIZE,
"ffn_dim_multiplier": 1.3,
"multiple_of": 1024,
"norm_eps": 1e-05,
"rope_theta": 500000.0,
"use_scaled_rope": True,
},
pth_file_count=1,
),Sources2
Direct evidence
Raw query–key similarities are divided by the square root of head width.
attn-scale · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 184–184
Read code excerpt
scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)Direct evidence
Meta's Llama 3.1 8B base catalogue entry supplies dimensions, head counts, the vocabulary (LLAMA3_VOCAB_SIZE, line 19), MLP rounding, epsilon, rope_theta and use_scaled_rope: True.
config · meta-llama/llama-models / models/sku_list.py
Pinned revision 0e0b8c519242 ↗ · lines 19–248
Read code excerpt
LLAMA3_VOCAB_SIZE = 128256… lines 20–230 omitted …
Model(
core_model_id=CoreModelId.llama3_1_8b,
description="Llama 3.1 8b model",
huggingface_repo="meta-llama/Llama-3.1-8B",
arch_args={
"dim": 4096,
"n_layers": 32,
"n_heads": 32,
"n_kv_heads": 8,
"vocab_size": LLAMA3_VOCAB_SIZE,
"ffn_dim_multiplier": 1.3,
"multiple_of": 1024,
"norm_eps": 1e-05,
"rope_theta": 500000.0,
"use_scaled_rope": True,
},
pth_file_count=1,
),