Hugging Face's logo

Chiru20
/
ANLP_assignment2

Text Generation
English
Vietnamese
Japanese
moe_transformer
nlp
mixture-of-experts
Mixture of Experts
transformer
optimizers
decoding-strategies
adamw
adan
lion
muon
sophia
Model card Files
xet
 Community
ANLP_assignment2 / src /part1 /config.py
Chiru20's picture
Chiru20
Upload ANLP Assignment 2 Model Card, Configs, Report, and Evaluation Deliverables
f90f02b
Raw History Blame Contribute Delete
6.29 kB
"""
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