""" 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()