32 independent transformer blocks › Attention
Normalized attention input
×tensorself.attention_norm(x)
Normalized residual features feeding the three attention projections.
- shape
- [B batch, 1 positions, 4096 d]
Codemodel.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 outSources1
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