adamroberts/tinystories-5090

llm.kittens TinyStories 124M BF16

This is a 124M-parameter GPT-2-style causal language model trained from scratch on TinyStories with the llm.kittens C++/CUDA trainer, which is a fork of Kaparthy's llm.c with some optimisations for SM120, and multi-stack kernel optimisations.

The model is published as a standard Hugging Face Transformers checkpoint with BF16 safetensors weights. It was trained on a single RTX 5090 in 14 hours.

Result

• Model weights: model.safetensors
• Training step: 20000 / 20000
• Final train loss: 0.785740
• Final validation loss: 0.875080
• Final throughput: 207135 tokens/s
• Final step time: 2531.04 ms
• Final reported BF16 MFU: 39.7%
• Average iteration time: 2605.014347 ms
• Safetensors size: 248,894,656 bytes
• Parameter count: 124,475,904

The TinyStories paper reports eval losses of 1.33 to 1.58 for the 768-hidden-size 1- and 2-layer attention-head ablations in Figure 24. This run's 0.875080 validation loss is lower, but the comparison is not apples-to-apples: this model is a 12-layer GPT-2-style model using GPT-2 tokenization, a 1024-token context, and a different implementation/training setup.

Architecture

• Family: GPT-2-style decoder-only Transformer
• Descriptor: d12
• Layers: 12
• Attention heads: 12
• Hidden size: 768
• Context length: 1024
• Vocabulary size: 50,257
• Precision: BF16 weights

Training

The run used the TinyStories GPT-2 dataset files generated by dev/data/tinystories.py in llm.kittens.

./train_gpt2cu \
    -i "dev/data/tinystories/TinyStories_train.bin" \
    -j "dev/data/tinystories/TinyStories_val.bin" \
    -o "log124M/5090_S" \
    -v 250 -s 20000 -g 144 \
    -h 0 \
    -b 64 -t 1024 -d 524288 \
    -r 0 \
    -z 1 \
    -c 0.1 \
    -l 0.0006 -q 0.0 -u 700 -n 5000 \
    -y 0 \
    -e "d12" \
    -x 20000

Key settings:

• Hardware target: RTX 5090 / SM120
• Micro batch: 64
• Sequence length: 1024
• Total desired batch size: 524,288 tokens
• Max steps: 20,000
• Optimizer: AdamW as implemented in llm.kittens
• Peak learning rate: 6e-4
• Scheduler: cosine
• Warmup: 700 steps
• Final LR fraction: 0.0
• Weight decay: 0.1
• Recompute: off
• ZeRO stage: 1
• Checkpoint interval: 5000 steps

Sample

Prompt/sample emitted at the final checkpoint:

Once upon a time, there was a little girl named Lily. She loved to play in the park with her friends. One day, while they were playing, they saw a big dog. The dog was very scary and Lily's friends ran away. But Lily stayed behind to protect herself.
The dog came closer and closer, but then it suddenly stopped and wagged its tail. "Don't be afraid," said the dog. "I just want to play too." Lily was surprised and happy that the dog was friendly. She realized that sometimes things that seem scary can turn out to be wonderful.
Later that day, Lily's mom asked her if she had any questions about the park. Lily smiled and said, "Yes, I saw a big dog and it was nice. I didn't run away because my friends were scared, but I stayed calm." Her mom was proud of her for being brave and not running away. From that day on, Lily knew that it's important to face your fears and stay calm, because you never know what might happen.

Files

• model.safetensors: BF16 Transformers weights.
• config.json: GPT-2 model configuration.
• generation_config.json: default generation settings.
• tokenizer.json: GPT-2 tokenizer.

Loading

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "adamroberts/tinystories-5090"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16)

inputs = tokenizer("Once upon a time", return_tensors="pt")
with torch.inference_mode():
    outputs = model.generate(**inputs, max_new_tokens=800, do_sample=True, temperature=0.8)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Source implementation: https://github.com/adamdroberts/llm.kittens

TinyStories reference paper: https://arxiv.org/abs/2305.07759