LLaMA Transformer Implementation in PyTorch

This implementation builds the LLaMA-style Transformer from scratch. It includes RMSNorm, Rotary Position Embeddings (RoPE), GQA-style multi-head attention with masking, SwiGLU MLP, and a stack of pre-norm decoder layers.

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

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

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        # Normalize and scale
        x_norm = x / rms * self.weight
        return x_norm

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_seq_len_cached = max_position_embeddings

        # Initialize cache
        t = torch.arange(self.max_seq_len_cached, device=inv_freq.device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :])
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :])

    def forward(self, x, seq_len=None):
        if seq_len > self.max_seq_len_cached:
            # Extend cache if needed
            self.max_seq_len_cached = seq_len
            t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1)
            self.cos_cached = emb.cos()[None, None, :, :]
            self.sin_cached = emb.sin()[None, None, :, :]
        
        return (
            self.cos_cached[:, :, :seq_len, ...],
            self.sin_cached[:, :, :seq_len, ...]
        )

class LlamaAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config['hidden_size']
        self.num_heads = config['num_attention_heads']
        self.head_dim = self.hidden_size // self.num_heads
        
        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        
        self.rotary_emb = RotaryEmbedding(self.head_dim)

    def forward(self, hidden_states, attention_mask=None):
        batch_size, seq_length, _ = hidden_states.size()
        
        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)
        
        # Reshape for multi-head attention
        q = q.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Apply rotary embeddings
        cos, sin = self.rotary_emb(q, seq_length)
        q, k = self.apply_rotary_pos_emb(q, k, cos, sin)
        
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        
        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
            attention_mask = (1.0 - attention_mask) * torch.finfo(scores.dtype).min
            scores = scores + attention_mask
        
        attn_weights = F.softmax(scores, dim=-1)

        # output
        attn_output = torch.matmul(attn_weights, v)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_length, self.hidden_size)
        
        return self.o_proj(attn_output)

    @staticmethod
    def apply_rotary_pos_emb(q, k, cos, sin):
        q_embed = (q * cos) + (q.roll(shifts=1, dims=-1) * sin)
        k_embed = (k * cos) + (k.roll(shifts=1, dims=-1) * sin)
        return q_embed, k_embed

class LlamaMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_size = config['hidden_size']
        intermediate_size = config['intermediate_size']
        
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x):
        # SwiGLU activation
        gate = F.silu(self.gate_proj(x))
        up = self.up_proj(x)
        return self.down_proj(gate * up)

class LlamaDecoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.self_attn = LlamaAttention(config)
        self.mlp = LlamaMLP(config)
        self.input_norm = RMSNorm(config['hidden_size'])
        self.post_attention_norm = RMSNorm(config['hidden_size'])

    def forward(self, hidden_states, attention_mask=None):
        # Pre-norm
        norm_states = self.input_norm(hidden_states)
        attn_output = self.self_attn(norm_states, attention_mask)
        hidden_states = hidden_states + attn_output
        
        # MLP with pre-norm
        norm_states = self.post_attention_norm(hidden_states)
        mlp_output = self.mlp(norm_states)
        hidden_states = hidden_states + mlp_output
        
        return hidden_states

class LlamaModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        self.embed_tokens = nn.Embedding(config['vocab_size'], config['hidden_size'])
        self.layers = nn.ModuleList([
            LlamaDecoderLayer(config) for _ in range(config['num_hidden_layers'])
        ])
        self.norm = RMSNorm(config['hidden_size'])

        # Initialize weights
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids, attention_mask=None):
        hidden_states = self.embed_tokens(input_ids)
        
        for layer in self.layers:
            hidden_states = layer(hidden_states, attention_mask)
        
        hidden_states = self.norm(hidden_states)
        return hidden_states

def test_llama_model():
    config = {
        'vocab_size': 32000,
        'hidden_size': 512,
        'num_attention_heads': 8,
        'num_hidden_layers': 4,
        'intermediate_size': 1024
    }

    model = LlamaModel(config)
    
    batch_size = 2
    seq_length = 16
    input_ids = torch.randint(0, config['vocab_size'], (batch_size, seq_length))
    
    attention_mask = torch.ones(batch_size, seq_length)
    
    output = model(input_ids, attention_mask)

    # Assertions for output shape
    assert output.shape == (batch_size, seq_length, config['hidden_size']), \
        f"Expected output shape {(batch_size, seq_length, config['hidden_size'])}, got {output.shape}"
    
    # Test attention mechanism
    attention_layer = LlamaAttention(config)
    hidden_states = torch.randn(batch_size, seq_length, config['hidden_size'])
    attn_output = attention_layer(hidden_states, attention_mask)
    
    assert attn_output.shape == hidden_states.shape, \
        f"Expected attention output shape {hidden_states.shape}, got {attn_output.shape}"
    
    print("All tests passed!")

if __name__ == "__main__":
    test_llama_model()