rtarun1/anlp-assignment2-checkpoints

ANLP Assignment 2: checkpoints

Final weights (Lightning final.ckpt, saved at the last training step) and the tokenizer of every run.

• part1/: decoder-only translation models (vi/ja to en), five FFN variants: dense MLP, 4 experts top-1, 4 experts top-2, 1 shared + 3 routed top-1, 4 wide experts top-2.
• part2/: the dense model pretrained on browndw/human-ai-parallel-corpus with different optimizers (AdamW, MARS, Lion, SOAP, Sophia, and the PyTorch AdamW reference), including the learning-rate sweep. Run names ending in -lr<value> are sweep runs; names without it use the default learning rate.

Load with the code repository (sys.path must include src/):

from part1.train import LitTranslator          # Part 1
model = LitTranslator.load_from_checkpoint("part1/part1-dense/final.ckpt", map_location="cpu")
from part2.train import LitLM                  # Part 2
model = LitLM.load_from_checkpoint("part2/part2-soap-lr6e-3/final.ckpt", map_location="cpu")

Training logs: https://wandb.ai/rtarun1-github/anlp-assignment-2