File size: 6,294 Bytes
f90f02b | 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 | """
Configuration dataclass for Transformer and Mixture-of-Experts (MoE) architectures.
Matches Listing 1 from Appendix with full parameter-accounting formulas for the 5 ablations.
"""
from dataclasses import dataclass, field
from typing import Optional, Literal
@dataclass
class TransformerConfig:
# Model architecture dimensions
D_model: int = 256 # Hidden representation dimension
N_ctx: int = 256 # Maximum context window length
N_vocab: int = 16000 # Vocabulary size
N_layers: int = 6 # Number of Transformer blocks
N_heads: int = 8 # Multi-head attention heads
D_ff: int = 1024 # Feed-forward hidden dimension (4 * D_model)
dropout: float = 0.1 # Dropout probability
# MoE architecture configuration
use_moe: bool = False # False: Standard MLP; True: MoE
ablation_variant: Literal[
"dense",
"top1",
"top2",
"shared_routed",
"matched_active"
] = "dense"
# MoE details
num_experts: int = 4 # Total experts in MoE layer
num_active_experts: int = 1 # Active experts routed per token
num_shared_experts: int = 0 # Always-active shared experts (for variant 4)
aux_loss_coef: float = 0.01 # Switch/GShard load balancing auxiliary loss weight
capacity_factor: float = 1.25 # Capacity factor for expert routing if capped
# Training settings
learning_rate: float = 3e-4
weight_decay: float = 0.01
warmup_steps: int = 500
max_steps: int = 10000
batch_size: int = 32
def __post_init__(self):
"""Validates configuration parameters and assigns ablation-specific settings."""
assert self.D_model % self.N_heads == 0, f"D_model ({self.D_model}) must be divisible by N_heads ({self.N_heads})"
if self.ablation_variant == "dense":
self.use_moe = False
else:
self.use_moe = True
if self.ablation_variant == "top1":
# 4 experts, 1 active at a time. Total params == Dense baseline.
self.num_experts = 4
self.num_active_experts = 1
self.num_shared_experts = 0
self.expert_d_ff = self.D_ff // 4
elif self.ablation_variant == "top2":
# 4 experts, 2 active at a time. Total params == Dense baseline.
self.num_experts = 4
self.num_active_experts = 2
self.num_shared_experts = 0
self.expert_d_ff = self.D_ff // 4
elif self.ablation_variant == "shared_routed":
# 1 shared expert (always active) + 3 routed (1 active). Total 2 of 4 active.
# Total params == Dense baseline.
self.num_experts = 4
self.num_shared_experts = 1
self.num_active_experts = 1 # 1 routed active + 1 shared active = 2 active total
self.expert_d_ff = self.D_ff // 4
elif self.ablation_variant == "matched_active":
# 4 experts, 2 active at a time, ACTIVE params == Dense baseline.
# Since 2 are active, each expert has D_ff / 2, so active = 2 * (D_ff / 2) = D_ff.
# Total params = 4 * (D_ff / 2) = 2 * Dense.
self.num_experts = 4
self.num_active_experts = 2
self.num_shared_experts = 0
self.expert_d_ff = self.D_ff // 2
else:
self.expert_d_ff = self.D_ff
@property
def ffn_total_params_per_layer(self) -> int:
"""Computes the theoretical total parameter count of one FFN/MoE layer."""
if not self.use_moe or self.ablation_variant == "dense":
# Linear1: D_model * D_ff + D_ff; Linear2: D_ff * D_model + D_model
return 2 * self.D_model * self.D_ff + self.D_ff + self.D_model
elif self.ablation_variant in ["top1", "top2"]:
# Router: D_model * 4 + 4 experts of (2 * D_model * d_expert)
router = self.D_model * self.num_experts
experts = self.num_experts * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
return router + experts
elif self.ablation_variant == "shared_routed":
# Router: D_model * 3 + 1 shared expert + 3 routed experts
router = self.D_model * 3
shared = 1 * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
routed = 3 * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
return router + shared + routed
elif self.ablation_variant == "matched_active":
router = self.D_model * self.num_experts
experts = self.num_experts * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
return router + experts
return 0
@property
def ffn_active_params_per_layer(self) -> int:
"""Computes the theoretical active parameter count of one FFN/MoE layer per token."""
if not self.use_moe or self.ablation_variant == "dense":
return self.ffn_total_params_per_layer
elif self.ablation_variant == "top1":
router = self.D_model * self.num_experts
active_expert = 1 * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
return router + active_expert
elif self.ablation_variant == "top2":
router = self.D_model * self.num_experts
active_experts = 2 * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
return router + active_experts
elif self.ablation_variant == "shared_routed":
router = self.D_model * 3
shared_expert = 1 * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
active_routed = 1 * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
return router + shared_expert + active_routed
elif self.ablation_variant == "matched_active":
router = self.D_model * self.num_experts
active_experts = 2 * (2 * self.D_model * self.expert_d_ff + self.expert_d_ff + self.D_model)
return router + active_experts
return 0
|