ANLP Assignment 2 (part2): lion
Decoder-only transformer with Lion (memory-efficient sign momentum), implemented from scratch.
Trained from scratch for IIIT-H Advanced NLP, Monsoon 2026. Architecture, training code and evaluation live in the assignment repository; this repository holds only the checkpoint.
Architecture
| Setting | Value |
|---|---|
d_model |
512 |
n_layers |
8 |
n_heads |
8 |
n_kv_heads |
8 |
n_ctx |
256 |
vocab_size |
32000 |
ffn |
dense |
d_ff |
2048 |
n_routed_experts |
0 |
n_shared_experts |
0 |
top_k |
0 |
norm |
rmsnorm |
Training
- Scored target tokens: 43,021,354
- Best validation loss: 3.8859
- Best validation perplexity: nan
Loading
import torch
payload = torch.load("model.pt", map_location="cpu", weights_only=False)
state_dict, config = payload["model"], payload["config"]