CS336 openwebtext 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: 4.140625 per token
- Training tokens represented by this checkpoint lineage: 647,168,000
- Parameters are 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("The United States", 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.