Llama 3 8B · Scale similarities

32 independent transformer blocks › Attention

Scale similarities

no parameters/ math.sqrt(self.head_dim)

Control the magnitude of similarities entering softmax.

{\htmlData{light=scale}{S_{\mathrm s}}}={\htmlData{light=similarities similarity}{S}}/\sqrt{128}

in: $S$ [B batch, 32 q heads, T queries, T keys]
out: $S_{\mathrm s}$ [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 $\sqrt{128}$ 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, 188–207

model.py 184

        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)

sku_list.py 19

LLAMA3_VOCAB_SIZE = 128256

sku_list.py 188–207

def llama3_base_models() -> list[Model]:
    return [
        Model(
            core_model_id=CoreModelId.llama3_8b,
            description="Llama 3 8b model",
            huggingface_repo="meta-llama/Llama-3-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": False,
            },
            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 original Llama 3 8B base catalogue entry supplies dimensions, head counts, vocabulary, MLP rounding settings, epsilon and unscaled RoPE settings.

config · meta-llama/llama-models / models/sku_list.py
Pinned revision 0e0b8c519242 ↗ · lines 19–207

Read code excerpt

LLAMA3_VOCAB_SIZE = 128256

… lines 20–187 omitted …

def llama3_base_models() -> list[Model]:
    return [
        Model(
            core_model_id=CoreModelId.llama3_8b,
            description="Llama 3 8b model",
            huggingface_repo="meta-llama/Llama-3-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": False,
            },
            pth_file_count=1,
        ),