anlp-a2-part2-optimizers

Part 2 of ANLP Assignment 2: the dense decoder-only transformer pretrained for next-token prediction on browndw/human-ai-parallel-corpus for 1x the dataset (43.6M tokens), one checkpoint per optimizer, all from identical initial weights on identical data. Optimizers implemented from scratch on torch.optim.Optimizer (Wen et al. 2025, Appendix A); hyperparameters chosen by a validation-loss sweep.

Numbers below are copied from the evaluation files uploaded next to each checkpoint.

optimizer	tokens seen	final val loss	final val ppl	final test BLEU
adamw/	43,626,496	3.6559	38.70	3.25
mars/	43,626,496	3.5245	33.94	3.71
lion/	43,626,496	3.6727	39.36	3.29
muon/	43,626,496	3.6598	38.85	3.55
sophia/	43,626,496	3.6595	38.84	3.22

Each folder holds model.safetensors (the tied lm_head.weight is stored once as embedding.weight), config.json (model + training configuration and the W&B run id) and the run's metrics. The tokenizer for every checkpoint is tokenizer.json at the repository root. Training logs: https://wandb.ai/skv-balasubramanian-iiit-hyderabad/anlp-a2-part2

import json
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from tokenizers import Tokenizer
from src.part1.model import Transformer, TransformerConfig   # model code from the assignment repository

cfg = json.load(open(hf_hub_download("bala-skv/anlp-a2-part2-optimizers", "adamw/config.json")))
model = Transformer(TransformerConfig(**cfg["model"]))
model.load_state_dict(load_file(hf_hub_download("bala-skv/anlp-a2-part2-optimizers", "adamw/model.safetensors")), strict=False)  # lm_head is tied
tok = Tokenizer.from_file(hf_hub_download("bala-skv/anlp-a2-part2-optimizers", "tokenizer.json"))