Llama Architecture: RMSNorm, SwiGLU, RoPE, GQA, and PyTorch Implementation

Evolution of Open-Weight Decoder-Only Language Models

Modern autoregressive language modeling builds upon the decoder-only Transformer topology. Early foundation models (e.g., GPT-3, PaLM, Chinchilla) established key scaling laws relating parameter count $N$ and training tokens $D$. However, initial implementations relied on post-LayerNorm, standard Multi-Head Attention (MHA), ReLU/GeLU activations, and absolute positional embeddings.

Llama (Touvron et al., Meta AI, 2023–2024) redefined open-weight architecture by introducing four synergistic hardware- and compute-efficient design choices:

1. Pre-Normalization with RMSNorm: Replaces standard LayerNorm by eliminating mean-centering, reducing memory access overhead while maintaining training stability.
2. SwiGLU Feed-Forward Networks: Employs gated linear activations with $\mathrm{SiLU}$ non-linearities, providing higher expressivity per parameter.
3. Rotary Position Embeddings (RoPE): Replaces absolute position vectors with complex coordinate rotations, naturally encoding relative token distances and enabling zero-shot context length extrapolation.
4. Grouped-Query Attention (GQA): Shares key-value projection heads across multiple query heads, cutting KV-cache memory consumption by $4\text{–}8\times$ during long-context autoregressive decoding.

Architectural Comparison

Mathematical Foundations

Figure 1: Llama decoder-only Transformer block topology featuring pre-RMSNorm, RoPE-conditioned GQA, and SwiGLU FFN.

1. Root Mean Square Normalization (RMSNorm)

Standard LayerNorm centers input activations by subtracting the mean $\mu$ before scaling by variance ${\sigma }^2$:

$\mathrm{LayerNorm}(\mathbf{x})=\mathbf{\gamma }\odot \frac{\mathbf{x}-\mu }{\sqrt{{\sigma }^2+ϵ}}+\mathbf{\beta },\mu =\frac{1}{d}{\sum }_{i=1}^d{x}_i,{\sigma }^2=\frac{1}{d}{\sum }_{i=1}^d(x_i-\mu {)}^2$

RMSNorm (Zhang & Sennrich, 2019) hypothesizes that the scaling invariance of LayerNorm is the primary stabilizing factor, rendering mean-centering redundant. For activation vector $\mathbf{x}\in {\mathbb{R}}^d$:

$\mathrm{RMSNorm}(\mathbf{x})=\mathbf{\gamma }\odot \frac{\mathbf{x}}{\mathrm{RMS}(\mathbf{x})+ϵ},\mathrm{RMS}(\mathbf{x})=\sqrt{\frac{1}{d}{\sum }_{i=1}^d{x}_i^2}$

where $\mathbf{\gamma }\in {\mathbb{R}}^d$ is a learnable gain parameter and $ϵ=10^{-5}$ prevents division by zero. Eliminating mean calculation reduces GPU global memory synchronization steps and saves $7\text{–}10\%$ kernel latency under half-precision (bfloat16) execution.

2. SwiGLU Gated Feed-Forward Networks

Standard Transformer Multi-Layer Perceptrons compute:

${\mathrm{FFN}}_{\text{Standard}}(\mathbf{x})={\mathbf{W}}_2\cdot \mathrm{GELU}({\mathbf{W}}_1\mathbf{x}+{\mathbf{b}}_1)+{\mathbf{b}}_2$

Llama replaces standard activations with SwiGLU (Shazeer, 2020), utilizing three weight matrices without bias terms:

${\mathrm{FFN}}_{\text{SwiGLU}}(\mathbf{x})={\mathbf{W}}_{\text{down}}\left(\mathrm{SiLU}({\mathbf{W}}_{\text{gate}}\mathbf{x})\odot ({\mathbf{W}}_{\text{up}}\mathbf{x})\right)$

where $\mathrm{SiLU}(z)=z\cdot \sigma (z)=\frac{z}{1+e^{-z}}$, ${\mathbf{W}}_{\text{gate}},{\mathbf{W}}_{\text{up}}\in {\mathbb{R}}^{d_{\text{ff}}\times d}$, and ${\mathbf{W}}_{\text{down}}\in {\mathbb{R}}^{d\times d_{\text{ff}}}$.

To maintain parameter parity with a standard $4d$ FFN expansion ratio, the intermediate dimension $d_{\text{ff}}$ is set to:

$d_{\text{ff}}=\lfloor \frac{2}{3}\cdot 4d\cdot \frac{1}{\text{multiple\_of}}\rfloor \cdot \text{multiple\_of}=\lfloor \frac{8d}{3\cdot 256}\rfloor \cdot 256$

For Llama 3 8B ($d=4096$), $d_{\text{ff}}=14336$.

3. Rotary Position Embeddings (RoPE)

RoPE (Su et al., 2021) injects positional information by rotating query and key vectors in complex 2D vector subspaces.

For a 2D vector coordinate pair $(x_{2k},x_{2k+1})$ at sequence position $m\in [0,T-1]$, the rotation matrix is defined as:

${\mathbf{R}}_{{\theta }_k,m}=\left(\begin{array}{cc}\cos (m{\theta }_k)& -\sin (m{\theta }_k)\\ \sin (m{\theta }_k)& \cos (m{\theta }_k)\end{array}\right),{\theta }_k={\theta }_{\text{base}}^{-2k/d},k\in \left[0,\frac{d}{2}-1\right]$

The full block-diagonal transformation ${\mathbf{R}}_{\Theta ,m}\in {\mathbb{R}}^{d\times d}$ acts on query ${\mathbf{q}}_m$ and key ${\mathbf{k}}_n$:

${\tilde{\mathbf{q}}}_m={\mathbf{R}}_{\Theta ,m}{\mathbf{q}}_m,{\tilde{\mathbf{k}}}_n={\mathbf{R}}_{\Theta ,n}{\mathbf{k}}_n$

Relative Coordinate Invariance Property:

Evaluating the attention dot-product:

$⟨{\tilde{\mathbf{q}}}_m,{\tilde{\mathbf{k}}}_n⟩={\mathbf{q}}_m^{\top }{\mathbf{R}}_{\Theta ,m}^{\top }{\mathbf{R}}_{\Theta ,n}{\mathbf{k}}_n={\mathbf{q}}_m^{\top }{\mathbf{R}}_{\Theta ,n-m}{\mathbf{k}}_n=g({\mathbf{q}}_m,{\mathbf{k}}_n,m-n)$

The inner product is strictly a function of the relative distance $(m-n)$, preserving orthogonal norm invariance $\Vert {\tilde{\mathbf{q}}}_m\Vert =\Vert {\mathbf{q}}_m\Vert$. In Llama 3.1, ${\theta }_{\text{base}}$ was scaled from $10^4$ to $5\times 10^5$, pushing maximum wavelength limits to support $128\text{K}$ context windows without frequency collision.

4. Grouped-Query Attention (GQA) and KV Cache Scaling

In standard Multi-Head Attention (MHA), the number of query heads $h$ equals the number of key/value heads $g$ ($h=g$). During autoregressive generation, storing past keys and values requires significant GPU VRAM:

${\text{Memory}}_{\text{MHA}}(\text{bytes})=2\times b\times L\times T\times h\times d_k\times \text{bytes\_per\_elem}$

Grouped-Query Attention (Ainslie et al., 2023) groups $h$ query heads into $g$ shared key-value groups ($r=h/g$ queries per KV head).

$\mathbf{Q}\in {\mathbb{R}}^{B\times h\times T\times d_k},\mathbf{K}\in {\mathbb{R}}^{B\times g\times T\times d_k},\mathbf{V}\in {\mathbb{R}}^{B\times g\times T\times d_k}$

Before computing attention, $\mathbf{K}$ and $\mathbf{V}$ are repeated $r$ times along the head dimension:

$\tilde{\mathbf{K}}=repeat\_interleave(\mathbf{K},r,\text{dim}=1),\tilde{\mathbf{V}}=repeat\_interleave(\mathbf{V},r,\text{dim}=1)$

$\mathrm{GQA}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}\left(\frac{\mathbf{Q}{\tilde{\mathbf{K}}}^{\top }}{\sqrt{d_k}}+{\mathbf{M}}_{\text{causal}}\right)\tilde{\mathbf{V}}$

KV Cache Memory Reduction:

$\frac{{\text{Memory}}_{\text{GQA}}}{{\text{Memory}}_{\text{MHA}}}=\frac{g}{h}=\frac{1}{r}$

For Llama 3 8B ($h=32,g=8$), memory footprint drops by $4\times$. For 70B ($h=64,g=8$), memory consumption drops by $8\times$, enabling $128\text{K}$ sequences on single-GPU instances.

Production PyTorch Implementation

Below is a complete, modular PyTorch implementation of the RMSNorm, SwiGLU FFN, RoPE Embedding Generator, GroupedQueryAttention, and complete LlamaTransformerBlock.

Step 1: Core Modular Layers

from __future__ import annotations

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""

    def __init__(self, dim: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._norm(x.float()).type_as(x) * self.weight


class SwiGLUFFN(nn.Module):
    """SwiGLU Feed-Forward Network: down_proj( SiLU(gate_proj(x)) * up_proj(x) )."""

    def __init__(self, dim: int, hidden_dim: int) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


def precompute_rope_freqs_cis(dim: int, end: int, theta: float = 500000.0) -> torch.Tensor:
    """Precompute complex frequencies for Rotary Position Embeddings."""
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64 [end, dim // 2]
    return freqs_cis


def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply RoPE rotation in complex space."""
    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 = freqs_cis[None, None, :xq.shape[2], :]
    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)

Step 2: Grouped-Query Attention & Transformer Block

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from llama_layers import RMSNorm, SwiGLUFFN, apply_rotary_emb


class GroupedQueryAttention(nn.Module):
    """Grouped-Query Attention (GQA) module with RoPE integration."""

    def __init__(
        self,
        dim: int,
        num_heads: int = 32,
        num_kv_heads: int = 8,
        max_seq_len: int = 131072,
    ) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = dim // num_heads
        self.num_queries_per_kv = num_heads // num_kv_heads

        self.q_proj = nn.Linear(dim, num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(dim, num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(dim, num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(num_heads * self.head_dim, dim, bias=False)

    def forward(
        self,
        x: torch.Tensor,
        freqs_cis: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        b, seq_len, _ = x.shape

        # Linear projections
        xq = self.q_proj(x).view(b, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        xk = self.k_proj(x).view(b, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        xv = self.v_proj(x).view(b, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # Apply RoPE
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        # GQA Repeat Interleave: [B, num_kv_heads, T, d_k] -> [B, num_heads, T, d_k]
        xk = xk.repeat_interleave(self.num_queries_per_kv, dim=1)
        xv = xv.repeat_interleave(self.num_queries_per_kv, dim=1)

        # Scaled Dot-Product Flash Attention
        output = F.scaled_dot_product_attention(
            xq, xk, xv, attn_mask=mask, is_causal=(mask is None and seq_len > 1)
        )

        output = output.transpose(1, 2).contiguous().view(b, seq_len, -1)
        return self.o_proj(output)


class LlamaTransformerBlock(nn.Module):
    """Complete Pre-RMSNorm Llama Decoder Block."""

    def __init__(
        self,
        dim: int = 4096,
        hidden_dim: int = 14336,
        num_heads: int = 32,
        num_kv_heads: int = 8,
        eps: float = 1e-5,
    ) -> None:
        super().__init__()
        self.attention_norm = RMSNorm(dim, eps=eps)
        self.attention = GroupedQueryAttention(dim, num_heads=num_heads, num_kv_heads=num_kv_heads)
        self.ffn_norm = RMSNorm(dim, eps=eps)
        self.feed_forward = SwiGLUFFN(dim, hidden_dim=hidden_dim)

    def forward(
        self,
        x: torch.Tensor,
        freqs_cis: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        h = x + self.attention(self.attention_norm(x), freqs_cis, mask=mask)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

Step 3: Production PEFT QLoRA Fine-Tuning Setup

from __future__ import annotations

import torch
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig


def setup_llama_qlora(model_id: str = "meta-llama/Meta-Llama-3.1-8B-Instruct") -> tuple[nn.Module, AutoTokenizer]:
    # 1. 4-bit NormalFloat Quantization Configuration
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )

    # 2. Load Model & Tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=bnb_config,
        device_map="auto",
        attn_implementation="flash_attention_2",
    )
    model = prepare_model_for_kbit_training(model)

    # 3. Target All Linear Projections
    peft_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )

    peft_model = get_peft_model(model, peft_config)
    peft_model.print_trainable_parameters()
    return peft_model, tokenizer

Empirical Benchmark Evaluation

Quantitative evaluation comparing model parameters, pretraining token count, and performance across standard language benchmarks:

Troubleshooting Common Deployment Faults

1. Numerical Overflow & NaN Gradients in FP16

• Symptom: Training loss immediately diverges or evaluates to NaN within the first 100 steps under standard FP16.
• Root Cause: RoPE frequency rotations and large intermediate activation scales in deep layers exceed FP16 maximum dynamic range ($65,504$).
• Remedy: Enforce strict torch.bfloat16 precision across training and inference.

2. End-of-Sequence Leaking & Run-On Generation

• Symptom: Model ignores conversation completion and continues generating repetitive text.
• Root Cause: Tokenizer EOS mismatch; Llama 3 defines <|eot_id|> (ID: 128009) as the conversational turn termination token, distinct from <|end_of_text|> (ID: 128001).
• Remedy: Explicitly configure eos_token_id=[128001, 128009] during model.generate().

3. VRAM OOM During Long-Context Inference ($>32\text{K}$)

• Symptom: Out-of-memory errors occur despite static model weights fitting into GPU memory.
• Root Cause: KV-cache allocation scales linearly with context length $T$; naive contiguous preallocation fragments memory.
• Remedy: Deploy through vLLM utilizing PagedAttention, enable 8-bit KV-cache quantization (--kv-cache-dtype fp8), or set --max-model-len 32768 to constrain maximum allocated page buffers.

References

1. Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971.
2. Touvron, H., et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288.
3. Dubey, A., et al. (2024). The Llama 3 Herd of Models. arXiv:2407.21783.
4. Su, J., et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.
5. Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv:2002.05202.
6. Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023.