---
library_name: pytorch
tags:
- anlp-a2
- optimizers
- lion
language:
- en
datasets:
- browndw/human-ai-parallel-corpus
---

# a2-task2-lion

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:** Lion (category: memory-efficient optimizers), implemented from scratch on `torch.optim.Optimizer`. Hyperparameters: lr=0.00025, betas=[0.9, 0.95], weight_decay=0.7. 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.0437 | 57.0 | 2.73 | 0.77 | 23.6 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):

```python
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'])
```