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.

Downloads last month
-
Safetensors
Model size
28.8M params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train balalida/gpt-22M-owt