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

Grouped-query attention

Full sequencePrefillDecode

32 independent transformer blocks › Attention

Rotated queries

×

tensorxq · 160

Query feature pairs after their position-dependent rotations.

shape
QrQ_{\mathrm r} [B batch, 32 q heads, 1 position, 128 d_h]

RoPE runs in float32 and the result is cast back to the input type (model.py 70–75).

Codemodel.py 65–75, 156–160
model.py 65–75
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)
model.py 156–160
        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
Sources1

Contextual evidence

Complex-pair rotations apply to queries and keys, with their original dtype restored.

rope · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 65–160

Read code excerpt
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

… lines 76–155 omitted …

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)