Llama 3.1

All models, with their blocks →

Size
Published
Meta · Jul 2024
Line
Dense decoder · follows Llama 3
View
Source

meta‑llama/llama3 @ a0940f9, which computes as Llama 3.1 but for its scaled RoPE (meta‑llama/llama‑models @ 0e0b8c5, model.py 45–69); configuration meta‑llama/llama‑models @ 0e0b8c5 · checkpoint, tokenizer and tensors not verifiedHow to read

Final RMSNorm

×

learnedself.norm(h)

Rescale each position’s features using their root mean square, then apply learned gain.

m2(x)=1d∑j=1dxj2y=γfinal⊙xm2(x)+εε=10−5\begin{aligned}m_2(x)&=\frac{1}{d}\sum_{j=1}^d x_j^2\\y&=\gamma_{\mathrm{final}}\odot\frac{x}{\sqrt{m_2(x)+\varepsilon}}\\\varepsilon&=10^{-5}\end{aligned}
in
x32x_{32} [B batch, T positions, 4096 d]
out
zz [B batch, T positions, 4096 d]
parameters
4,096

Normalized in float32 and cast back to the input type before the gain multiplies (model.py 45–46).

Parameters

4,096 parameters per instance.

Exact per-instance accounting
ComponentCalculationParameters
Learned gain4,0964{,}0964,096
Per-instance total4,096

Each normalization instance owns its learned gain vector.

Derived from the configured feature width; checkpoint tensors have not been independently counted.

Codemodel.py 35–303 · sku_list.py 19, 231–248
model.py 35–46
class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight
model.py 236–237
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
model.py 266
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
model.py 266–269
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = ColumnParallelLinear(
            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
        )
model.py 298–303
        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        h = self.norm(h)
        output = self.output(h).float()
        return output
sku_list.py 19
LLAMA3_VOCAB_SIZE = 128256
sku_list.py 231–248
        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,
        ),
Sources3

Direct evidence

RMSNorm rescales the final feature axis and applies one learned gain vector; block and final callers supply configured epsilon.

rms · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 35–266

Read code excerpt
class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

… lines 47–235 omitted …

        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

… lines 238–265 omitted …

        self.norm = RMSNorm(params.dim, eps=params.norm_eps)

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,
        ),

Direct evidence

Final RMSNorm precedes a separate bias-free vocabulary projection; the returned logits are converted to float32.

model-readout · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 266–303

Read code excerpt
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = ColumnParallelLinear(
            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
        )

… lines 270–297 omitted …

        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        h = self.norm(h)
        output = self.output(h).float()
        return output