a2-task2-sophia

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: Sophia-G (bonus: Hessian-based optimizer), implemented from scratch on torch.optim.Optimizer. Hyperparameters: lr=0.000125, betas=[0.965, 0.99], rho=0.08, weight_decay=0.1, eps=1e-12, hessian_update_every=10, estimator=gnb. 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.1371	62.6	2.84	0.65	27.5 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'])