Grouped-query attention

Reference code meta-llama/llama3 @ a0940f9, configuration meta-llama/llama-models @ 0e0b8c5 · checkpoint, tokenizer and tensors not verified · How to read

32 independent transformer blocks › Attention

Normalized attention input

×

tensorself.attention_norm(x)

Normalized residual features feeding the three attention projections.

shape
xx [B batch, T positions, 4096 d]
Codemodel.py 236–248
model.py 236–248
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out
Sources1

Contextual evidence

Attention receives attention_norm(x) and its update is added to x; the MLP receives ffn_norm(h) and its update is added to h.

handoff · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 236–248

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

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out