DeepSeek Architecture Implementation in PyTorch

This implementation builds a DeepSeek-style network from scratch. It includes RMSNorm-based pre-norm blocks, a multi-head latent attention module, a Mixture-of-Experts (MoE) feed-forward router that combines expert outputs by softmax-weighted routing, and a stack of DeepSeekBlocks with token-level residuals and a final projection head.

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

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)
        x_normalized = x / rms * self.weight
        return x_normalized
    
class MultiHeadLatentAttention(nn.Module):
    def __init__(self, dim, num_heads, num_latents, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scaling = self.head_dim ** -0.5

        self.latents = nn.Parameter(torch.randn(num_latents, dim))
        
        # Projection matrices
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        self.out_proj = nn.Linear(dim, dim)
        
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        B, L, D = x.shape
        H = self.num_heads

        k = self.k_proj(x) 
        v = self.v_proj(x)  

        q = self.q_proj(x)  

        q = q.view(B, L, H, self.head_dim).transpose(1, 2)  
        k = k.view(B, L, H, self.head_dim).transpose(1, 2)  
        v = v.view(B, L, H, self.head_dim).transpose(1, 2)  

        # Scaled dot-product attention
        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scaling  
        
        if mask is not None:
            attn = attn.masked_fill(mask.unsqueeze(1) == 0, float('-inf'))
        
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        out = torch.matmul(attn, v)  
        out = out.transpose(1, 2).contiguous().view(B, L, D)  

        return self.out_proj(out)
    
class ExpertLayer(nn.Module):
    """Single expert layer with feed-forward network"""
    def __init__(self, dim, ff_dim):
        super().__init__()
        self.ff = nn.Sequential(
            nn.Linear(dim, ff_dim),
            nn.GELU(),
            nn.Linear(ff_dim, dim)
        )

    def forward(self, x):
        return self.ff(x)
    
class MixtureOfExperts(nn.Module):
    """Mixture of Experts layer with routing"""
    def __init__(self, dim, num_experts, ff_dim):
        super().__init__()
        self.num_experts = num_experts
        self.router = nn.Linear(dim, num_experts)
        self.experts = nn.ModuleList([
            ExpertLayer(dim, ff_dim) for _ in range(num_experts)
        ])

    def forward(self, x):
        # Calculate routing probabilities
        router_logits = self.router(x)  
        router_probs = F.softmax(router_logits, dim=-1)
        
        final_output = torch.zeros_like(x)

        # Route input to each expert and combine outputs
        for i, expert in enumerate(self.experts):
            expert_output = expert(x)
            expert_weight = router_probs[..., i:i+1]
            final_output += expert_output * expert_weight

        return final_output
    
class DeepSeekBlock(nn.Module):
    def __init__(self, dim, num_heads, num_latents, num_experts, ff_dim):
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.attn = MultiHeadLatentAttention(dim, num_heads, num_latents)
        self.norm2 = RMSNorm(dim)
        self.moe = MixtureOfExperts(dim, num_experts, ff_dim)

    def forward(self, x, mask=None):
        x = x + self.attn(self.norm1(x), mask)
        x = x + self.moe(self.norm2(x))
        return x
    
class DeepSeek(nn.Module):
    def __init__(self, 
                 dim=512, 
                 num_layers=6,
                 num_heads=8,
                 num_latents=64,
                 num_experts=4,
                 ff_dim=2048):
        super().__init__()
        
        self.embedding = nn.Linear(dim, dim)  
        
        self.blocks = nn.ModuleList([
            DeepSeekBlock(dim, num_heads, num_latents, num_experts, ff_dim)
            for _ in range(num_layers)
        ])
        
        self.final_norm = RMSNorm(dim)
        self.output_proj = nn.Linear(dim, dim)

    def forward(self, x, mask=None):
        x = self.embedding(x)
        
        for block in self.blocks:
            x = block(x, mask)
            
        x = self.final_norm(x)
        x = self.output_proj(x)
        
        return x
    
def test_deepseek():
    batch_size = 2
    seq_length = 16
    dim = 512
    
    model = DeepSeek(
        dim=dim,
        num_layers=6,
        num_heads=8,
        num_latents=64,
        num_experts=4,
        ff_dim=2048
    )

    print("Test 1: Basic forward pass")
    x = torch.randn(batch_size, seq_length, dim)
    output = model(x)
    assert output.shape == (batch_size, seq_length, dim)
    print("✓ Output shape is correct:", output.shape)

    print("\nTest 2: Forward pass with attention mask")
    mask = torch.ones(batch_size, seq_length, seq_length)
    mask[:, :, seq_length//2:] = 0 
    output_masked = model(x, mask)
    assert output_masked.shape == (batch_size, seq_length, dim)
    print("✓ Masked output shape is correct:", output_masked.shape)

    print("\nTest 3: RMSNorm")
    norm = RMSNorm(dim)
    norm_output = norm(x)
    assert torch.allclose(torch.mean(norm_output ** 2), torch.tensor(1.0), atol=1e-1)
    print("✓ RMSNorm output has approximately unit variance")

    print("\nTest 4: Mixture of Experts")
    moe = MixtureOfExperts(dim, num_experts=4, ff_dim=2048)
    moe_output = moe(x)
    assert moe_output.shape == x.shape
    print("✓ MoE output shape matches input shape:", moe_output.shape)

    print("\nAll tests passed successfully!")

if __name__ == "__main__":
    test_deepseek()