anlp-a2-p2-muon

ANLP A2 checkpoint: muon

Muon, Newton-Schulz (Wen et al. 2025, Algorithm 8)

• Part 2 of the Advanced NLP Assignment 2 (Monsoon 2026).
• Task: next-token pretraining on HAP-E (chunk_1 + chunk_2 human text)
• Token budget: 9,715,712 tokens (1x dataset, HAP-E next-token prediction, batch 32)
• Model: 35.7M-parameter decoder-only transformer (D=512, 6 layers, 8 heads, vocab 32768, tied embedding), final checkpoint (end of token budget).
• Total params: 35,658,240 | Active params: 35,658,240

Files

• model.pt -- {"model": <state_dict>, "config": <TransformerConfig dict>}
• config.json -- the TransformerConfig
• tokenizer.json -- byte-level BPE (v32768) trained on the training split
• README.md -- this file

Loading (with the assignment repo checked out)

from src.part1.model import TransformerLM
from src.part1.data import load_tokenizer

model = TransformerLM.load("part2_b32/muon/final", device="cuda")
tokenizer = load_tokenizer("tokenizer")  # or: from this repo's tokenizer.json