TinyLM
A 3.4M parameter causal language model trained from scratch, for experimentation.
Architecture
| Hyperparameter | Value |
|---|---|
| Parameters | 3.403.968 |
| Layers | 4 |
| Hidden size | 64 |
| Attention heads | 4 |
| FFN dim | 192 |
| Embedding rank | 32 |
| Context length | 256 |
| Tokenizer | GPT-2 (50257 vocab) |
Uses a factored (low-rank) embedding to keep the vocab projection from eating the entire parameter budget, with weight tying on the output head.
Training
| Datasets | Skylion007/openwebtext (10k samples), roneneldan/TinyStories (10k samples) |
| Optimizer | AdamW (lr=3e-3, weight_decay=0.01) |
| Scheduler | Cosine annealing with warm restarts |
| Mixed precision | fp16 (torch.cuda.amp) |
| Hardware | Nvidia P100 |
Usage
from huggingface_hub import snapshot_download
import importlib.util
import torch
# Download files
snapshot_download(repo_id="Fu01978/TinyLM", local_dir="./tinylm")
# Load via script
spec = importlib.util.spec_from_file_location("modeling_tinylm", "./tinylm/modeling_tinylm.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
model, tokenizer, config = module.load_tinylm("./tinylm")
model.eval()
# Generate
output = module.generate(model, tokenizer, "Once upon a time, ")
print(output)