Transformer Architecture

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

import torch
import torch.nn as nn
 
class PositionWiseFFN(nn.Module):
    """
    Position-wise Feed-Forward Network (FFN) as used in the original Transformer.
    
    Architecture: Linear(d_model → d_ff) → ReLU → Linear(d_ff → d_model)
    
    The same network is applied independently to each position in the sequence.
    No information flows between positions within the FFN.
    """
    def __init__(
        self,
        d_model: int = 512,
        d_ff: int = 2048,
        dropout: float = 0.1,
        activation: str = "relu"
    ):
        super().__init__()
        
        # First linear projection: expand to intermediate dimension
        self.w1 = nn.Linear(d_model, d_ff)
        
        # Second linear projection: contract back to model dimension
        self.w2 = nn.Linear(d_ff, d_model)
        
        # Activation function
        if activation == "relu":
            self.activation = nn.ReLU()
        elif activation == "gelu":
            self.activation = nn.GELU()
        else:
            raise ValueError(f"Unknown activation: {activation}")
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply position-wise FFN.
        
        Args:
            x: Input tensor of shape [batch_size, seq_len, d_model]
            
        Returns:
            Output tensor of shape [batch_size, seq_len, d_model]
        """
        # Expand: d_model → d_ff
        hidden = self.w1(x)
        
        # Apply non-linearity
        hidden = self.activation(hidden)
        
        # Apply dropout (after activation, before second projection)
        hidden = self.dropout(hidden)
        
        # Contract: d_ff → d_model
        output = self.w2(hidden)
        
        return output
 
 
def visualize_ffn_dimensions():
    """Demonstrate the dimension flow through FFN."""
    d_model, d_ff = 512, 2048
    batch_size, seq_len = 2, 10
    
    ffn = PositionWiseFFN(d_model, d_ff)
    x = torch.randn(batch_size, seq_len, d_model)
    
    # Trace through the network
    print(f"Input shape: {x.shape}")
    print(f"  → {[batch_size, seq_len, d_model]}")
    
    with torch.no_grad():
        h1 = ffn.w1(x)
        print(f"After W1: {h1.shape}")
        print(f"  → Expanded from {d_model} to {d_ff} (4x)")
        
        h2 = ffn.activation(h1)
        print(f"After activation: {h2.shape} (same)")
        
        y = ffn.w2(h2)
        print(f"After W2: {y.shape}")
        print(f"  → Contracted back to {d_model}")
    
    # Parameter count
    w1_params = d_model * d_ff + d_ff  # weights + bias
    w2_params = d_ff * d_model + d_model  # weights + bias
    total = w1_params + w2_params
    print(f"
Parameter count:")
    print(f"  W1: {d_model} × {d_ff} + {d_ff} = {w1_params:,}")
    print(f"  W2: {d_ff} × {d_model} + {d_model} = {w2_params:,}")
    print(f"  Total: {total:,}")
 
visualize_ffn_dimensions()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
 
# Activation function implementations
class ActivationFunctions:
    """Common activation functions used in Transformer FFNs."""
    
    @staticmethod
    def relu(x: torch.Tensor) -> torch.Tensor:
        """ReLU: max(0, x)"""
        return F.relu(x)
    
    @staticmethod
    def gelu(x: torch.Tensor) -> torch.Tensor:
        """GELU: x * Φ(x) where Φ is the CDF of standard normal."""
        return F.gelu(x)
    
    @staticmethod
    def gelu_approx(x: torch.Tensor) -> torch.Tensor:
        """Fast GELU approximation used in many implementations."""
        return 0.5 * x * (1 + torch.tanh(
            np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3)
        ))
    
    @staticmethod
    def swish(x: torch.Tensor) -> torch.Tensor:
        """Swish (SiLU): x * sigmoid(x)"""
        return x * torch.sigmoid(x)
    
    @staticmethod
    def swish_beta(x: torch.Tensor, beta: float = 1.0) -> torch.Tensor:
        """Parameterized Swish: x * sigmoid(beta * x)"""
        return x * torch.sigmoid(beta * x)
 
 
def compare_activations():
    """Compare activation functions and their gradients."""
    x = torch.linspace(-4, 4, 1000, requires_grad=True)
    
    activations = {
        'ReLU': F.relu,
        'GELU': F.gelu,
        'Swish': lambda x: x * torch.sigmoid(x),
    }
    
    print("Activation Comparison at key points:")
    print("-" * 60)
    test_points = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
    
    for name, fn in activations.items():
        values = fn(test_points)
        print(f"{name:10s}: {values.tolist()}")
    
    print("
Gradient at x=0:")
    for name, fn in activations.items():
        x0 = torch.tensor([0.0], requires_grad=True)
        y = fn(x0)
        y.backward()
        print(f"  {name}: {x0.grad.item():.4f}")
    
    print("
Gradient at x=-1:")
    for name, fn in activations.items():
        x0 = torch.tensor([-1.0], requires_grad=True)
        y = fn(x0)
        y.backward()
        print(f"  {name}: {x0.grad.item():.4f}")
 
 
def sparsity_analysis():
    """Analyze sparsity induced by different activations."""
    torch.manual_seed(42)
    
    # Simulate FFN intermediate activations
    batch_size, seq_len, d_ff = 16, 128, 2048
    pre_activation = torch.randn(batch_size, seq_len, d_ff)
    
    relu_out = F.relu(pre_activation)
    gelu_out = F.gelu(pre_activation)
    swish_out = pre_activation * torch.sigmoid(pre_activation)
    
    # Count "zero" or near-zero activations
    threshold = 1e-6
    
    relu_zeros = (relu_out.abs() < threshold).float().mean().item()
    gelu_zeros = (gelu_out.abs() < threshold).float().mean().item()
    swish_zeros = (swish_out.abs() < threshold).float().mean().item()
    
    print("Sparsity Analysis (fraction of near-zero activations):")
    print(f"  ReLU:  {relu_zeros:.4f} ({relu_zeros*100:.1f}%)")
    print(f"  GELU:  {gelu_zeros:.4f} ({gelu_zeros*100:.1f}%)")
    print(f"  Swish: {swish_zeros:.4f} ({swish_zeros*100:.1f}%)")
    
    # More meaningful: fraction of "small" outputs
    small_threshold = 0.1
    relu_small = (relu_out.abs() < small_threshold).float().mean().item()
    gelu_small = (gelu_out.abs() < small_threshold).float().mean().item()
    swish_small = (swish_out.abs() < small_threshold).float().mean().item()
    
    print(f"
Fraction of small activations (|x| < 0.1):")
    print(f"  ReLU:  {relu_small:.4f} ({relu_small*100:.1f}%)")
    print(f"  GELU:  {gelu_small:.4f} ({gelu_small*100:.1f}%)")
    print(f"  Swish: {swish_small:.4f} ({swish_small*100:.1f}%)")
 
 
compare_activations()
sparsity_analysis()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class SwiGLU_FFN(nn.Module):
    """
    SwiGLU Feed-Forward Network as used in LLaMA and other modern LLMs.
    
    Architecture:
        gate = Swish(W_gate @ x)
        content = W_up @ x
        output = W_down @ (gate ⊙ content)
    
    The intermediate dimension is typically 8/3 * d_model to maintain
    similar parameter count to standard FFN with 4 * d_model.
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int = None,
        dropout: float = 0.0,
        bias: bool = False  # LLaMA omits biases
    ):
        super().__init__()
        
        # Default to 8/3 * d_model, rounded to multiple of 256
        if d_ff is None:
            d_ff = int(8 * d_model / 3)
            d_ff = (d_ff + 255) // 256 * 256  # Round up to multiple of 256
        
        # Three projections for GLU
        self.w_gate = nn.Linear(d_model, d_ff, bias=bias)  # Gate projection
        self.w_up = nn.Linear(d_model, d_ff, bias=bias)    # Up/content projection
        self.w_down = nn.Linear(d_ff, d_model, bias=bias)  # Down/output projection
        
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply SwiGLU FFN.
        
        Args:
            x: Input [batch_size, seq_len, d_model]
            
        Returns:
            Output [batch_size, seq_len, d_model]
        """
        # Compute gate and content in parallel
        gate = F.silu(self.w_gate(x))  # SiLU is Swish with β=1
        content = self.w_up(x)
        
        # Element-wise gating
        hidden = gate * content
        hidden = self.dropout(hidden)
        
        # Project back to model dimension
        output = self.w_down(hidden)
        
        return output
 
 
class GeGLU_FFN(nn.Module):
    """
    GeGLU Feed-Forward Network using GELU gating.
    
    Used in some T5 variants and other models.
    """
    def __init__(self, d_model: int, d_ff: int = None, bias: bool = True):
        super().__init__()
        
        if d_ff is None:
            d_ff = int(8 * d_model / 3)
        
        self.w_gate = nn.Linear(d_model, d_ff, bias=bias)
        self.w_up = nn.Linear(d_model, d_ff, bias=bias)
        self.w_down = nn.Linear(d_ff, d_model, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = F.gelu(self.w_gate(x))
        content = self.w_up(x)
        return self.w_down(gate * content)
 
 
class ReGLU_FFN(nn.Module):
    """
    ReGLU: GLU variant with ReLU gating.
    """
    def __init__(self, d_model: int, d_ff: int = None, bias: bool = True):
        super().__init__()
        
        if d_ff is None:
            d_ff = int(8 * d_model / 3)
        
        self.w_gate = nn.Linear(d_model, d_ff, bias=bias)
        self.w_up = nn.Linear(d_model, d_ff, bias=bias)
        self.w_down = nn.Linear(d_ff, d_model, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = F.relu(self.w_gate(x))
        content = self.w_up(x)
        return self.w_down(gate * content)
 
 
def compare_ffn_variants():
    """Compare standard FFN vs GLU variants."""
    d_model = 4096
    
    # Standard FFN
    d_ff_standard = 4 * d_model  # 16384
    standard_params = 2 * d_model * d_ff_standard  # Ignoring bias
    
    # GLU FFN (adjusted d_ff to match params)
    d_ff_glu = int(8 * d_model / 3)  # ~10922
    glu_params = 3 * d_model * d_ff_glu
    
    print("FFN Parameter Comparison (d_model = 4096):")
    print(f"  Standard FFN (d_ff = {d_ff_standard:,}):")
    print(f"    Parameters: {standard_params:,}")
    print(f"  SwiGLU FFN (d_ff = {d_ff_glu:,}):")
    print(f"    Parameters: {glu_params:,}")
    print(f"    Ratio: {glu_params / standard_params:.3f}")
    
    # Test forward pass
    batch_size, seq_len = 2, 128
    x = torch.randn(batch_size, seq_len, d_model)
    
    swiglu = SwiGLU_FFN(d_model, d_ff_glu)
    
    # Count actual parameters
    actual_params = sum(p.numel() for p in swiglu.parameters())
    print(f"
  Actual SwiGLU params: {actual_params:,}")
 
compare_ffn_variants()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
 
class MoELayer(nn.Module):
    """
    Simplified Mixture of Experts layer.
    
    Each token is routed to top-k experts. Only those experts
    perform computation for that token, enabling sparse scaling.
    
    Used in: Switch Transformer, Mixtral, GLaM, etc.
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int = 8,
        top_k: int = 2,
        dropout: float = 0.1
    ):
        super().__init__()
        
        self.num_experts = num_experts
        self.top_k = top_k
        self.d_model = d_model
        
        # Router: produces expert scores for each token
        self.router = nn.Linear(d_model, num_experts, bias=False)
        
        # Expert FFNs (each is a standard FFN)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.ReLU(),
                nn.Linear(d_ff, d_model),
                nn.Dropout(dropout)
            )
            for _ in range(num_experts)
        ])
        
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass with expert routing.
        
        Args:
            x: Input [batch_size, seq_len, d_model]
            
        Returns:
            output: Processed output [batch_size, seq_len, d_model]
            router_logits: For auxiliary loss [batch_size, seq_len, num_experts]
        """
        batch_size, seq_len, d_model = x.shape
        
        # Compute router scores
        router_logits = self.router(x)  # [batch, seq, num_experts]
        
        # Select top-k experts for each token
        routing_weights, selected_experts = torch.topk(
            router_logits, self.top_k, dim=-1
        )
        routing_weights = F.softmax(routing_weights, dim=-1)
        
        # Initialize output
        output = torch.zeros_like(x)
        
        # Flatten for easier processing
        flat_x = x.view(-1, d_model)  # [batch*seq, d_model]
        flat_output = output.view(-1, d_model)
        flat_selected = selected_experts.view(-1, self.top_k)
        flat_weights = routing_weights.view(-1, self.top_k)
        
        # Route tokens to experts (simplified, not optimized)
        for i in range(self.top_k):
            for expert_idx in range(self.num_experts):
                # Find tokens routed to this expert at position i
                mask = (flat_selected[:, i] == expert_idx)
                if mask.any():
                    expert_input = flat_x[mask]
                    expert_output = self.experts[expert_idx](expert_input)
                    # Weight by routing weight
                    flat_output[mask] += flat_weights[mask, i:i+1] * expert_output
        
        output = flat_output.view(batch_size, seq_len, d_model)
        
        return output, router_logits
    
    def load_balance_loss(self, router_logits: torch.Tensor) -> torch.Tensor:
        """
        Compute auxiliary load balancing loss.
        
        Encourages equal utilization of all experts.
        """
        # Average routing probability per expert
        routing_probs = F.softmax(router_logits, dim=-1)
        avg_per_expert = routing_probs.mean(dim=[0, 1])  # [num_experts]
        
        # Ideal is 1/num_experts for each
        target = 1.0 / self.num_experts
        
        # Variance from ideal
        balance_loss = torch.sum((avg_per_expert - target) ** 2)
        
        return balance_loss
 
 
def demonstrate_moe():
    """Show MoE behavior and scaling."""
    d_model = 512
    d_ff = 2048
    num_experts = 8
    top_k = 2
    
    moe = MoELayer(d_model, d_ff, num_experts, top_k)
    
    # Parameter count comparison
    single_ffn_params = 2 * d_model * d_ff
    moe_params = sum(p.numel() for p in moe.parameters())
    
    print("MoE Parameter Analysis:")
    print(f"  Single FFN: {single_ffn_params:,} parameters")
    print(f"  MoE ({num_experts} experts, top-{top_k}): {moe_params:,} parameters")
    print(f"  Capacity increase: {num_experts}x")
    print(f"  Compute increase: ~{top_k}x (only top-{top_k} active)")
    
    # Forward pass
    batch_size, seq_len = 2, 16
    x = torch.randn(batch_size, seq_len, d_model)
    
    output, router_logits = moe(x)
    
    print(f"
  Input shape: {x.shape}")
    print(f"  Output shape: {output.shape}")
    
    # Analyze routing
    selected = router_logits.argmax(dim=-1)  # Top-1 for simplicity
    unique, counts = torch.unique(selected, return_counts=True)
    print(f"
  Expert utilization (top-1):")
    for exp, cnt in zip(unique.tolist(), counts.tolist()):
        print(f"    Expert {exp}: {cnt} tokens ({cnt/(batch_size*seq_len)*100:.1f}%)")
 
demonstrate_moe()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
 
class EfficientFFN(nn.Module):
    """
    FFN with efficiency optimizations for large-scale training.
    
    Features:
    - Optional gradient checkpointing
    - Mixed precision support
    - Fused operations where available
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        dropout: float = 0.0,
        use_checkpoint: bool = False,
        bias: bool = False
    ):
        super().__init__()
        
        self.use_checkpoint = use_checkpoint
        
        # No bias improves efficiency slightly
        self.w1 = nn.Linear(d_model, d_ff, bias=bias)
        self.w2 = nn.Linear(d_ff, d_model, bias=bias)
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        
    def _forward_impl(self, x: torch.Tensor) -> torch.Tensor:
        """Core forward logic, separated for checkpointing."""
        # Fuse W1 + activation where possible
        hidden = self.w1(x)
        hidden = torch.nn.functional.gelu(hidden)
        hidden = self.dropout(hidden)
        return self.w2(hidden)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.use_checkpoint and self.training:
            # Gradient checkpointing: save memory, recompute in backward
            return checkpoint(
                self._forward_impl, x, 
                use_reentrant=False
            )
        return self._forward_impl(x)
 
 
class TensorParallelFFN(nn.Module):
    """
    FFN designed for tensor parallelism across GPUs.
    
    W1 is column-parallel: split along d_ff dimension
    W2 is row-parallel: split along input dimension
    
    This allows partial results to be computed independently,
    then all-reduced only after W2.
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_partitions: int,
        partition_id: int
    ):
        super().__init__()
        
        assert d_ff % num_partitions == 0
        local_d_ff = d_ff // num_partitions
        
        self.num_partitions = num_partitions
        self.partition_id = partition_id
        
        # Each GPU has a slice of the FFN
        self.w1 = nn.Linear(d_model, local_d_ff, bias=False)
        self.w2 = nn.Linear(local_d_ff, d_model, bias=False)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward with local computation.
        Caller must handle all-reduce after this.
        """
        hidden = torch.nn.functional.gelu(self.w1(x))
        output = self.w2(hidden)
        # In practice: output = all_reduce(output) across partitions
        return output
 
 
def benchmark_ffn_variants():
    """Benchmark different FFN implementations."""
    import time
    
    d_model = 4096
    d_ff = 16384
    batch_size = 8
    seq_len = 2048
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    
    # Standard FFN
    standard = nn.Sequential(
        nn.Linear(d_model, d_ff),
        nn.GELU(),
        nn.Linear(d_ff, d_model)
    ).to(device)
    
    # Checkpointed FFN
    checkpointed = EfficientFFN(d_model, d_ff, use_checkpoint=True).to(device)
    
    x = torch.randn(batch_size, seq_len, d_model, device=device)
    
    # Warmup
    for _ in range(3):
        _ = standard(x).sum().backward()
        standard.zero_grad()
    
    if device.type == 'cuda':
        torch.cuda.synchronize()
    
    # Benchmark forward + backward
    n_iters = 10
    
    start = time.perf_counter()
    for _ in range(n_iters):
        y = standard(x)
        y.sum().backward()
        standard.zero_grad()
    if device.type == 'cuda':
        torch.cuda.synchronize()
    standard_time = time.perf_counter() - start
    
    start = time.perf_counter()
    for _ in range(n_iters):
        y = checkpointed(x)
        y.sum().backward()
        checkpointed.zero_grad()
    if device.type == 'cuda':
        torch.cuda.synchronize()
    checkpoint_time = time.perf_counter() - start
    
    print(f"
Benchmark ({n_iters} iterations):")
    print(f"  Standard FFN: {standard_time:.3f}s")
    print(f"  Checkpointed: {checkpoint_time:.3f}s")
    print(f"  Overhead: {(checkpoint_time/standard_time - 1)*100:.1f}%")
    
    if device.type == 'cuda':
        # Memory comparison
        torch.cuda.reset_peak_memory_stats()
        _ = standard(x).sum().backward()
        standard_mem = torch.cuda.max_memory_allocated() / 1e9
        
        standard.zero_grad()
        torch.cuda.reset_peak_memory_stats()
        _ = checkpointed(x).sum().backward()
        checkpoint_mem = torch.cuda.max_memory_allocated() / 1e9
        
        print(f"
Peak memory:")
        print(f"  Standard: {standard_mem:.2f} GB")
        print(f"  Checkpointed: {checkpoint_mem:.2f} GB")
        print(f"  Savings: {(1 - checkpoint_mem/standard_mem)*100:.1f}%")
 
# Uncomment to run benchmark:
# benchmark_ffn_variants()