CS336 tinystories causal language model
Four-layer decoder-only Transformer trained for Stanford CS336 Assignment 1. It uses a byte-level BPE tokenizer, context length 256, width 512, 16 attention heads, SwiGLU width 1344, RMSNorm, and RoPE.
- Validation loss: 1.389844 per token
- Training tokens represented by this checkpoint lineage: 311,296,000
- Parameters are not tied between the token embedding and LM head.
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("PATH_OR_REPO", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("PATH_OR_REPO", trust_remote_code=True)
inputs = tokenizer("Once upon a time", return_tensors="pt")
output = model.generate(**inputs, max_new_tokens=128, temperature=0.8, top_p=0.9, do_sample=True)
print(tokenizer.decode(output[0], skip_special_tokens=False))
This is a small educational model and may hallucinate, repeat, or produce incorrect information.