File size: 7,927 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 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 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | """
Main CLI entrypoint for ANLP Assignment 2.
Supports running:
- Part 1: MoE Transformer Ablations (5 variants)
- Part 2: Custom Optimizers Pretraining Dynamics (AdamW, Adan, Lion, Muon, Sophia)
- Part 3: Decoding Strategies Benchmark on ROCStories (Greedy, Top-k, Top-p, Beam Search 1/2/4)
- Smoke tests & parameter verification
"""
import sys
import os
import argparse
import logging
import torch
from src.utils import set_seed, count_parameters
from src.part1.config import TransformerConfig
from src.part1.model import Transformer
from src.part1.moe import MoE, MLP
from src.part2.optimizers import (
CustomAdamW,
CustomAdan,
CustomLion,
CustomMuon,
CustomSophia,
)
from src.part3.decoding import (
greedy_decode,
top_k_decode,
top_p_decode,
beam_search_decode,
)
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger("ANLP_Main")
def run_part1(args):
logger.info(f"Running Part 1: MoE Ablation variant '{args.ablation}'")
config = TransformerConfig(ablation_variant=args.ablation)
model = Transformer(config)
params = count_parameters(model)
logger.info(f"Total Parameters: {params['total_parameters']:,}")
logger.info(f"Active Parameters: {params['active_parameters']:,}")
if getattr(args, "train", False):
from src.part1.train import run_part1_training
logger.info(f"Starting actual training for {args.ablation} (samples={args.samples}, epochs={args.epochs})...")
run_part1_training(
ablation_variant=args.ablation,
max_samples=args.samples,
epochs=args.epochs
)
else:
# Dummy tensor verification
dummy_input = torch.randint(0, config.N_vocab, (2, config.N_ctx))
logits, loss, aux_loss = model(dummy_input, targets=dummy_input)
logger.info(f"Forward output logits shape: {logits.shape}, Loss: {loss.item():.4f}, Aux Loss: {aux_loss.item():.4f}")
def run_part2(args):
logger.info(f"Running Part 2: Pretraining Optimizer benchmark '{args.optimizer}'")
config = TransformerConfig(ablation_variant="dense")
model = Transformer(config)
from src.part2.train import get_optimizer
if args.optimizer == "all":
opts = ["adamw", "adan", "lion", "muon", "sophia"]
else:
opts = [args.optimizer]
if getattr(args, "train", False):
from src.part2.train import run_part2_pipeline
logger.info(f"Starting actual pretraining benchmark for {opts} (tokens={args.tokens})...")
run_part2_pipeline(optimizer_names=opts, total_budget_tokens=args.tokens)
else:
for opt_name in opts:
opt = get_optimizer(opt_name, model, lr=config.learning_rate)
logger.info(f"Successfully initialized custom optimizer: {type(opt).__name__}")
def run_part3(args):
logger.info(f"Running Part 3: Decoding Strategies evaluation on {args.samples} samples")
from src.part3.evaluate import run_decoding_benchmark
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
run_decoding_benchmark(
model_name=args.model_name,
num_samples=args.samples,
device=device,
output_dir=args.output_dir
)
def verify_all_architectures():
"""Verifies all mathematical parameter counts, optimizers, and decoding strategies."""
logger.info("=== Verifying Part 1: All 5 FFN Architecture Variants ===")
variants = ["dense", "top1", "top2", "shared_routed", "matched_active"]
dense_total, dense_active = 0, 0
for v in variants:
cfg = TransformerConfig(ablation_variant=v)
model = Transformer(cfg)
params = count_parameters(model)
tot = params["total_parameters"]
act = params["active_parameters"]
if v == "dense":
dense_total, dense_active = tot, act
diff_tot = tot - dense_total
diff_act = act - dense_active
logger.info(
f"Variant: {v:<15} | Total Params: {tot:>10,} (diff vs dense: {diff_tot:>+8,}) | "
f"Active Params: {act:>10,} (diff vs dense: {diff_act:>+8,})"
)
# Quick forward pass test
dummy_idx = torch.randint(0, cfg.N_vocab, (2, 16))
logits, loss, aux = model(dummy_idx, targets=dummy_idx)
assert logits.shape == (2, 16, cfg.N_vocab), f"Unexpected shape for {v}: {logits.shape}"
assert not torch.isnan(loss), f"Loss is NaN for {v}"
logger.info("=== Verifying Part 2: All 5 Custom Optimizers ===")
toy_p = torch.nn.Parameter(torch.randn(10, 10))
toy_grad = torch.randn(10, 10)
for OptCls in [CustomAdamW, CustomAdan, CustomLion, CustomMuon, CustomSophia]:
opt = OptCls([toy_p], lr=1e-3)
toy_p.grad = toy_grad.clone()
opt.step()
logger.info(f"Verified single step for {OptCls.__name__}")
logger.info("=== Verifying Part 3: All 4 Decoding Strategies from Scratch ===")
class ToyLM(torch.nn.Module):
def __init__(self):
super().__init__()
self.emb = torch.nn.Embedding(50, 16)
self.head = torch.nn.Linear(16, 50)
def forward(self, x):
return self.head(self.emb(x))
toy_lm = ToyLM()
test_input = torch.tensor([[1, 2, 3]], dtype=torch.long)
g_out, _, _ = greedy_decode(toy_lm, test_input, max_new_tokens=3)
k_out, _, _ = top_k_decode(toy_lm, test_input, k=5, max_new_tokens=3)
p_out, _, _ = top_p_decode(toy_lm, test_input, p=0.8, max_new_tokens=3)
b_out, _, _ = beam_search_decode(toy_lm, test_input, beam_width=2, max_new_tokens=3)
assert g_out.shape == (1, 6)
assert k_out.shape == (1, 6)
assert p_out.shape == (1, 6)
assert b_out.shape == (1, 6)
logger.info("Verified greedy, top-k, top-p, and beam search decoding.")
logger.info("=== All verification checks passed successfully! ===")
def main():
parser = argparse.ArgumentParser(description="ANLP Assignment 2: MoE, Optimizers, Decoding")
subparsers = parser.add_subparsers(dest="command", help="Sub-commands")
# Part 1 parser
p1 = subparsers.add_parser("part1", help="Run Part 1 MoE Ablations")
p1.add_argument("--ablation", choices=["dense", "top1", "top2", "shared_routed", "matched_active"], default="dense")
p1.add_argument("--train", action="store_true", default=True, help="Run actual training and generate heatmap")
p1.add_argument("--no-train", action="store_false", dest="train", help="Skip training, run sanity forward only")
p1.add_argument("--samples", type=int, default=10000)
p1.add_argument("--epochs", type=int, default=1)
# Part 2 parser
p2 = subparsers.add_parser("part2", help="Run Part 2 Optimizers")
p2.add_argument("--optimizer", choices=["adamw", "adan", "lion", "muon", "sophia", "all"], default="all")
p2.add_argument("--train", action="store_true", default=True, help="Run actual pretraining benchmark and generate dynamics plot")
p2.add_argument("--no-train", action="store_false", dest="train", help="Skip training, run sanity checks only")
p2.add_argument("--tokens", type=int, default=200000)
# Part 3 parser
p3 = subparsers.add_parser("part3", help="Run Part 3 Decoding Strategies")
p3.add_argument("--model_name", default="EleutherAI/Pythia-160M")
p3.add_argument("--samples", type=int, default=1000)
p3.add_argument("--output_dir", default="assets/part3")
# Verify parser
subparsers.add_parser("verify", help="Run architecture, optimizer, and decoding sanity verification")
args = parser.parse_args()
if args.command == "part1":
run_part1(args)
elif args.command == "part2":
run_part2(args)
elif args.command == "part3":
run_part3(args)
elif args.command == "verify" or args.command is None:
verify_all_architectures()
if __name__ == "__main__":
main()
|