tiny-llama – ein kleines Llama-artiges Sprachmodell von Grund auf

Komplett selbst implementiertes Decoder-only-Modell im Llama-Stil (PyTorch), trainiert auf TinyStories.

Struktur

Datei	Inhalt
model.py	Architektur: RMSNorm, RoPE (precompute_rope/apply_rope), SwiGLU-MLP, kausale Attention mit GQA (n_kv_heads < n_heads), Block (Pre-Norm + Residual), Llama (Embedding, N Blöcke, Final-Norm, LM-Head mit Weight Tying), generate() mit Temperatur + Top-k
prepare_data.py	Trainiert einen Byte-Level-BPE-Tokenizer (HF tokenizers, Vokabular 8192) und schreibt data/train.bin / data/val.bin (uint16-Token-IDs, 95/5-Split)
train.py	Training mit AdamW (Weight Decay nur auf Matrizen), linearem Warmup + Cosine-LR, Gradient Clipping, regelmäßiger Evaluation und Checkpoint (checkpoints/ckpt.pt). Konfiguration über die Dataclass TrainConfig
sample.py	Lädt den Checkpoint und generiert Text
data/	Rohtext, Tokenizer, tokenisierte Daten
checkpoints/	ckpt.pt + eval_log.json (Loss-Verlauf)

Schnellstart

cd /workspace/tiny-llama && source .venv/bin/activate
python prepare_data.py                      # Tokenizer + Daten
python train.py                             # Training (Standardwerte aus TrainConfig)
python train.py max_iters=5000 lr=6e-4      # Werte per key=value überschreiben
python train.py resume=true max_iters=4000  # vom Checkpoint weitertrainieren
python sample.py --prompt "Once upon a time" --n 3 --temperature 0.8

Hochskalieren

• Modellgröße: dim, n_layers, n_heads, n_kv_heads, hidden_dim (≈ 8/3·dim, auf Vielfaches von 64 runden). Beispiele:
  • ~25M: dim=512 n_layers=8 n_heads=8 n_kv_heads=4 hidden_dim=1376
  • ~110M: dim=768 n_layers=12 n_heads=12 n_kv_heads=4 hidden_dim=2048
• Kontext: seq_len=512 oder 1024 (RoPE wird automatisch angepasst).
• Daten: Das aktuelle Training nutzt nur die TinyStories-Validierungsdatei (~4,6M Tokens). Für mehr Qualität TinyStoriesV2-GPT4-train.txt (~2 GB) von Hugging Face laden und prepare_data.py --input ... ausführen. Faustregel (Chinchilla): ~20 Tokens pro Parameter.
• GPU: train.py nutzt automatisch CUDA. Für Geschwindigkeit zusätzlich sinnvoll: torch.autocast(dtype=torch.bfloat16), torch.compile(model), größere batch_size bzw. Gradient Accumulation, bei mehreren GPUs DDP.
• Lernrate: bei größeren Modellen kleiner wählen (z. B. 3e-4 bis 6e-4).
• Vokabular: bei vielfältigeren Daten (Code, mehrere Sprachen) --vocab 32000; dann uint16 reicht bis 65535.

Mit eigenen Daten trainieren

1. Text als UTF-8-Datei ablegen, z. B. data/meine_daten.txt. Dokumente mit <|endoftext|> trennen (oder einen eigenen Trenner per --sep angeben).
2. python prepare_data.py --input data/meine_daten.txt --out data_eigen --vocab 8192 (Ordner vorher anlegen).
3. python train.py data_dir=data_eigen out_dir=ckpt_eigen
4. python sample.py --ckpt ckpt_eigen/ckpt.pt --tokenizer data_eigen/tokenizer.json --prompt "..."

Hinweis: Deutscher Text funktioniert genauso – der Byte-Level-BPE lernt Umlaute automatisch. Bei sehr wenig Daten (< 1M Tokens) Modell kleiner halten und dropout=0.1 in ModelConfig setzen, sonst Overfitting (val-Loss steigt, während train-Loss fällt).

Ergebnis des ersten Trainingslaufs

• Hardware: CPU (8 Kerne, keine GPU), PyTorch 2.14.1+cpu
• Modell: 6,45M Parameter (dim 256, 6 Layer, 8 Heads / 4 KV-Heads, SwiGLU 688, Kontext 256)
• Training: 1500 Iterationen × 32 × 256 Tokens (~12M Tokens, ~2,6 Epochen), ~17,5 Minuten
• Loss: 9,08 → train 2,23 / val 2,44

Ergebnis des zweiten Trainingslaufs (v0.2)

• Daten: neu tokenisiertes, größeres TinyStories-Subset mit ~49,5M Trainings-Tokens und ~1,0M Validierungs-Tokens
• Training: fortgesetzt bis Iteration 9000 (Batch 12 × 256, LR 6e-4 → 6e-5 Cosine), ~63 Minuten auf der CPU
• Bester Checkpoint (checkpoints/ckpt.pt, Iteration 7750): train 1,89 / val 1,82
• Letzter Stand (Iteration 9000): train 1,87 / val 1,87

Version	Train-Loss	Val-Loss
v0.1	2,23	2,44
v0.2	1,89	1,82
v0.3	1,76	1,84

Beispiel (v0.2):

Once upon a time, there was a little fish named Fin. Fin lived in the sea and liked to swim all day. One day, he saw a big shell near the pond. He was very curious. Fin swam up to the shell and said, "Hello, shell! Can you help me?"

Der Checkpoint ist als Asset am Release v0.2 angehängt.

Ergebnis des dritten Trainingslaufs (v0.3)

• Fortsetzung ab v0.2 von Iteration 9000 bis 11750 (Batch 12 × 256, LR 6e-4 → 3e-5), ~44 Minuten auf der CPU
• Zusätzlich ~8,4M Tokens (insgesamt ~36M Token-Positionen in diesem Datensatz-Lauf)
• Bester Checkpoint (Iteration 11750): train 1,76 / val 1,84
• Hinweis: v0.3 nutzt eval_iters=30 statt 10 bei v0.2; der Val-Wert ist dadurch stabiler und nicht direkt vergleichbar (v0.2 lag mit 30 Eval-Batches bei ~1,92 zum Start dieses Laufs).

Beispiel (v0.3):

Once upon a time, there was a little girl named Lily. Lily liked to design things. She made pretty pictures with her hands. [...] Her mom said, "Wow, Lily! You did a great job helper too!" Lily felt proud of herself.

Der Checkpoint ist als Asset am Release v0.3 angehängt.