LayanSh1/tinystories-llama-60m

TinyStories Llama 60M (trained from scratch)

A 60M-parameter Llama-style causal language model trained from random initialisation on English TinyStories, as a course assignment on using scaling laws to plan model size and training-token budget, then configuring and training a Llama-style model.

Architecture

Setting	Value
Total parameters	60,247,552
Layers (L)	19
Hidden size (H)	512
Attention heads (A)	8 (head dimension 64)
KV heads (K)	4 (grouped-query attention)
Intermediate size (I)	1,408 (SwiGLU)
Vocabulary size (V)	8,192
Context length	256 tokens (training block length)
Other	RMSNorm (eps 1e-5), RoPE (θ = 10,000), no biases, tied input/output embeddings

Parameter breakdown: SwiGLU feed-forward 41.1M (68.2%), attention 14.9M (24.8%), shared embedding/LM head 4.2M (7.0%), normalisation 20K.

Why this architecture:

• N ≈ 60M (the bottom of the allowed 60–100M range): my real compute budget is small, so the scaling law favours a much smaller model (see below). Choosing the smallest allowed size gives the most training tokens per parameter in a fixed time.
• 19 layers × 512 width: among ~60M options, this is closest to the target. It balances depth and width: a 36 × 384 model is slower on a GPU because the layers run one after another, and a 13 × 640 model is larger (64M), so it would see fewer tokens per parameter.
• GQA with 4 KV heads: fewer parameters and a smaller KV cache than full multi-head attention (8 KV heads), with little loss in quality.
• Intermediate size 1,408 ≈ 8/3 × 512, rounded to a multiple of 64. This keeps the 3-matrix SwiGLU block about the same size as a standard 4H MLP and is efficient on GPUs.
• Vocabulary 8,192: TinyStories uses simple language, so a small vocabulary is enough and keeps the embedding at only 7% of parameters.
• Tied embeddings save 4.2M parameters.

Dataset

• Source: roneneldan/TinyStories, training split
• Language: English
• Licence: CDLA-Sharing-1.0
• Content: short stories synthetically generated by GPT-3.5 and GPT-4 using a restricted vocabulary
• Why I chose it: TinyStories was designed for training small language models. Its simple vocabulary and grammar let a 60M model show measurable learning and produce readable text within a few GPU-hours, whereas a more varied corpus (e.g. Wikipedia) would need far more tokens before producing coherent text.

Preprocessing:

1. Streamed and shuffled the published training split (seed 42, buffer 10,000) and kept the first 300,000 unique stories.
2. Normalised whitespace and removed exact duplicates using SHA-256 hashes. This doesn't catch paraphrases or near-duplicates.
3. Split whole stories into train/validation/test (90/5/5) before tokenisation, so pieces of one story can't appear in two splits.
4. Trained a byte-level BPE tokenizer (vocabulary 8,192, min frequency 2) on the training split only. Special tokens: <pad>, <bos>, <eos>, <unk>.
5. Appended <eos> after every story, joined stories within each split, and cut them into non-overlapping 256-token blocks (the incomplete final block was dropped).

Split	Stories	Tokens	256-token blocks
Train	270,000	56,836,863	222,018
Validation	15,000	3,177,577	12,412
Test	15,000	3,146,434	12,290

The published TinyStories validation split was not used.

Scaling-law calculations

I used the simplified Chinchilla rule D ≈ 20N with training compute C ≈ 6ND. Substituting gives C = 120N², so N = √(C / 120).

Planned budget:

Quantity	Value
Planned parameters N	60,000,000 (actual model: 60,247,552)
Planned tokens D = 20N	1.2B (1,204,951,040 for the actual N)
Planned compute C = 6ND	4.32 × 10¹⁷ FLOPs

Actual run:

Quantity	Value
Hardware budget	1 × T4 GPU, ~90 minutes of training
Measured throughput	5,398 tokens/s
Input tokens processed	29,491,200 (1,800 updates × 16,384 tokens)
Scored next-token targets	29,376,000 (255 per 256-token block)
Actual compute (6ND)	1.07 × 10¹⁶ FLOPs
Coverage of planned D	2.45%
Tokens per parameter	0.49 (Chinchilla: ~20)
Unique training data used	51.9% of the 56.8M training tokens, each seen once (no repetition)

Interpretation: for my real compute (1.07 × 10¹⁶ FLOPs), the Chinchilla rule would recommend a model of only √(1.07 × 10¹⁶ / 120) ≈ 9.4M parameters. The assignment requires 60–100M, so I chose the smallest allowed size. My model processed about 40× fewer tokens per parameter than the Chinchilla ratio, so it is heavily under-trained relative to the scaling law. The loss was still slowly decreasing at the end, so more tokens would very likely improve it. The 6ND estimate ignores attention-score computation and implementation overhead.

Training setup

Setting	Value
Hardware	1 × NVIDIA Tesla T4 (16 GB), Google Colab
Precision	FP16 mixed precision with gradient scaling (FP32 weights)
Optimizer	AdamW, β = (0.9, 0.95), weight decay 0.1, gradient clipping 1.0
Learning rate	5 × 10⁻⁴ peak, 100 linear warmup steps, cosine decay to 0
Micro-batch	16 sequences × 256 tokens
Gradient accumulation	4 steps
Effective batch	64 sequences = 16,384 tokens per update
Training steps	1,800 optimizer updates
Training duration	~91 minutes of training compute (5,463 s), plus evaluation/saving time
Peak GPU memory	4.47 GiB
Seed	42
Software	PyTorch 2.14.1, transformers 4.46.3, datasets 3.1.0, tokenizers 0.20.3

Choices compared with the course lab: the effective batch is 4× larger (16,384 vs 4,096 tokens), so I raised the learning rate from 3e-4 to 5e-4. I turned off gradient checkpointing because memory was plentiful, which increased throughput from ~3,050 to ~5,400 tokens/s. I used 10× more stories (300K vs 30K) so the run never repeats data.

Evaluation

Next-token cross-entropy (natural log) and perplexity = exp(loss), on held-out stories:

Evaluation	Split	Loss	Perplexity	Targets scored
Before training (random weights)	validation subset (20 batches)	≈ 9.1	≈ 9,000	81,600
End of training	validation subset (20 batches)	2.043	7.71	81,600
Final	full validation	2.041	7.69	3,165,060
Final	full test	2.037	7.67	3,133,950

• The starting loss is close to ln(8,192) ≈ 9.01, as expected when a random model guesses evenly across the vocabulary.
• The validation split guided training. The test split was evaluated only once, after training finished.
• Validation and test scores are almost identical, so the model generalises to unseen stories.
• Perplexity is only comparable between models using the same tokenizer, data and evaluation procedure.

Loss curve

The blue line is the loss on a fixed 20-batch validation subset, measured every 100 updates. The orange line is the loss of the last effective training batch at each of those points. The loss drops quickly from 9.1 to about 3 in the first 300 updates, then decreases slowly to 2.04. The two curves overlap throughout, so there is no sign of overfitting, which is expected because every training story was seen at most once. The curve is still slightly decreasing at step 1,800.

Example outputs

Before vs after training (same prompt, greedy decoding, 60 new tokens):

Before: Once upon a time, a little robot found a book. dessert dessert dessert dessert dessert dessert … curt curt curt curt curt …

After: Once upon a time, a little robot found a book. The robot was very happy and wanted to read it. But the little robot was scared. The robot did not know what to do. The robot asked the little robot, "Why are you sad?" The little robot said, "I am sad because I can't read my book. I can't

Sampled outputs (temperature 0.7, top-k 50, up to 80 new tokens; prompt in bold):

Once upon a time, there was a little girl named Lily. She loved to play with her toys and sing songs. One day, Lily's mom asked her to help with a special massage. Lily was so excited to help and started to sing a song. When the party was over, Lily's mom was happy and said, "Thank you for making me feel better, Lily." Lily felt happy and proud of herself. From that day on, Lily

The dog was very sad because he had no friends to play with and no one could play with him anymore. The end.

Tom and his mom went to the park. They saw a big slide and decided to climb it. They climbed up the ladder and went down. They laughed and laughed as they climbed up. At the top, they saw a big dog. The dog was barking and running around. Tom and his mom were scared. They wanted to run away. They all ran to the slide and threw down. It was a bad ending for Tom and his mom.

One day, a big storm came to the village. The people were scared, but they all wanted to help them. So, they went on a long way to get to the storm. As they reached the storm, they saw that the lightning was coming down. The people were scared and didn't know what to do. They knew they had to be brave, so they ran up to the storm. They tried to stay calm, but the storm was

Limitations

Observed in the samples:

• Weak logic and cause-and-effect: events don't always follow ("help with a special massage" then "when the party was over"; "ran up to the storm" to be brave).
• Confused characters and references: "The robot asked the little robot" when there is only one robot; "the people… wanted to help them" with no clear "them".
• Odd word choices and actions: "they all ran to the slide and threw down".
• Sudden or unsatisfying endings: some stories stop after one sentence ("The end.") or end abruptly ("It was a bad ending").
• Mild repetition of words and phrases ("The robot… the robot… the little robot").
• Longer outputs are often cut off mid-sentence by the token limit.

General limitations:

• Heavily under-trained relative to the scaling law (0.49 tokens per parameter vs ~20).
• Trained only on simple, synthetic English children's stories. It doesn't handle other topics, languages, factual questions or instructions, and is not a chat or instruction-following model.
• Trained on 256-token blocks; quality on longer inputs is untested and likely worse.
• Inherits any biases or errors in the GPT-3.5/GPT-4-generated TinyStories data.
• Documents are packed together, so the model could attend across story boundaries during training.
• Intended for education and experimentation only.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "LayanSh1/tinystories-llama-60m"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForCausalLM.from_pretrained(repo_id)
model.eval()

prompt = "Once upon a time, there was a little girl named"
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=80,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
print(tokenizer.decode(output[0], skip_special_tokens=True))

This code was tested in a fresh Google Colab notebook on CPU, without logging in to Hugging Face, using Colab's default transformers version. Use add_special_tokens=False, which matches how the model was trained.

Files in this repository

• model.safetensors, config.json, generation_config.json: model weights and configuration
• tokenizer.json, tokenizer_config.json, special_tokens_map.json: tokenizer
• loss_curve.png: training and validation loss curves
• history.json: validation/training loss at each evaluation step
• experiment.json: full record of settings, split sizes, token counts and results

References

• Hoffmann et al. (2022), Training Compute-Optimal Large Language Models (Chinchilla): https://arxiv.org/abs/2203.15556
• Eldan & Li (2023), TinyStories: How Small Can Language Models Be and Still Speak Coherent English?: https://arxiv.org/abs/2305.07759