| """ |
| 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_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,})" |
| ) |
|
|
| |
| 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") |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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") |
|
|
| |
| 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() |
|
|
