a2-task2-mars

ANLP Assignment 2, Task 2 (Optimizers): next-token-prediction pretraining from scratch on browndw/human-ai-parallel-corpus for 1x the dataset (37.8M tokens, 4,615 steps of 8,192 tokens).

Optimizer: MARS (category: variance-reduced AdamW variants), implemented from scratch on torch.optim.Optimizer. Hyperparameters: lr=0.002, betas=[0.9, 0.95], eps=1e-08, weight_decay=0.1, gamma=0.025. 300 warmup steps, cosine decay to 0.1x, gradient-norm clip 1.0.

Decoder-only transformer: d_model 512, 8 layers, 8 heads, context 256, learned absolute positions, pre-LayerNorm, bias-free GELU FFN, 32k BPE vocabulary with the embedding tied to the LM head. (Same architecture as the Task 1 dense baseline.)

final val loss	val ppl	test BLEU (7 refs)	test BLEU (human ref)	train time
4.0839	59.4	3.12	0.55	25.2 min

Files: model.pt (a torch.save dict: model state_dict, step, tokens_seen, config), tokenizer.json (HF tokenizers), config.yaml (training config), metrics.json (full eval output).

Loading (with the assignment code on the path):

import torch
from src.part1.config import TransformerConfig, resolve_ffn_dims
from src.part1.model import build_model

ckpt = torch.load('model.pt', map_location='cpu', weights_only=False)
model = build_model(resolve_ffn_dims(TransformerConfig(**ckpt['config'])))
model.load_state_dict(ckpt['model'])