GPT-2-Modell mit `llm.c` günstig nachbauen

• Erklärung, wie sich das GPT-2-(124M)-Modell mit llm.c in 90 Minuten für 20 $ reproduzieren lässt
• GPT-2 (124M) ist das kleinste Modell, das OpenAI 2019 veröffentlicht hat
• Mit einem 8X-A100-80GB-SXM-Node bei Lambda liegen die Kosten bei etwa 14 $ pro Stunde, insgesamt also bei rund 20 $
  • Das Training ist auch auf einer einzelnen GPU möglich, dauert dann aber länger (4–24 Stunden)

Vergleich der Ergebnisse

• Auf dem FineWeb-Validierungsdatensatz zeigt es bessere Leistung als der von OpenAI veröffentlichte Checkpoint
  • Allerdings wurde GPT-2 mit WebText trainiert, daher ist der Vergleich nicht vollständig fair
• Auch die HellaSwag-Genauigkeit wurde gemessen: Mit 29,9 liegt sie nahe an GPT-3 Small (124M) mit 33,7
  • GPT-2 (124M) mit 29,4 wird bereits übertroffen
  • Hier wurde jedoch mit 10B Tokens trainiert, während GPT-3 mit 300B Tokens trainiert wurde

Minimale Umgebungseinrichtung

• GPU erforderlich (Lambda Labs empfohlen)
• Anleitung basiert auf Linux x86 64bit Ubuntu 22.04 mit CUDA 12
• Installation von miniconda und anschließend der PyTorch-Nightly-Version (optional)
• Installation der für den Tokenizer benötigten Pakete
• Installation von cuDNN für mehr Geschwindigkeit (optional)
• Installation von MPI bei Nutzung mehrerer GPUs (optional)
• Vorverarbeitung des FineWeb-10B-Token-Datensatzes (~1 Stunde)
• Kompilieren von llm.c (gemischte Präzision, cuDNN FlashAttention)

Training ausführen

• Beispielbefehl für die Nutzung einer einzelnen GPU
• Ausführung mit mpirun bei Multi-GPU-Betrieb (8 GPUs)
• Erklärung der wichtigsten Argumente
  • -i, -j: Pfade zu Trainings-/Validierungsdaten
  • -o: Speicherpfad für Logs und Checkpoints
  • -e: Modellinitialisierung (GPT-2 mit Tiefe 12)
  • -b: Micro-Batch-Größe (bei Speichermangel reduzieren)
  • -t: maximale Sequenzlänge
  • -d: gesamte Batch-Größe (siehe GPT-3-Paper)
  • -r: Recompute-Einstellung (spart Speicher)
  • -z: ZeRO-1 (Sharding des Optimizer-Status)
  • Weitere Einstellungen wie Weight Decay, Lernrate, Checkpoint-Intervall usw.

Trainingsverlauf

• Bei 10B Trainingstokens und einer Batch-Größe von 0,5M werden etwa 20K Schritte erwartet
• Für eine A100-40GB-PCIe-GPU werden Zeit pro Schritt, MFU und Token-Durchsatz ausgegeben
• Zu Beginn des Trainings tritt Gradient Exploding auf, wird aber durch Clipping behoben

Visualisierung

• Bereitgestellt wird ein Jupyter-Notebook, das Log-Dateien parst und die Trainingskurven visualisiert

Tokenizer

• Wird benötigt, um Integer-Tokens in Zeichenketten umzuwandeln
• Kann mit einem PyTorch-Skript erzeugt werden

Sampling

• Aktuell nicht für Inferenz optimiert
• Mit kleineren Code-Anpassungen sind unconditional/conditional Sampling möglich

Codestruktur

• Der Großteil der Implementierung befindet sich in der Datei train_gpt2.cu
• Die ersten 500 Zeilen enthalten Setups für MPI, NCCL, cuDNN, cuBLAS usw.
• Die folgenden 1500 Zeilen enthalten Forward/Backward des Transformers
• Die nächsten 1000 Zeilen implementieren das GPT-2-Modell
• Die letzten 1000 Zeilen enthalten Trainingsschleife, Argument-Parsing usw.

350M-Modell

• 10B Tokens reichen nicht aus, stattdessen werden 30B Tokens verwendet
• Mit 8X A100 80GB dauert das 14 Stunden und kostet etwa 200 $

FAQ

• Sampling möglich?: Ja, aber ineffizient.
• Chat möglich?: Derzeit nur Pretraining, kein Chat-Fine-Tuning.
• Verteiltes Training über mehrere Nodes?: Möglich, aber noch nicht getestet.
• Bitgenau deterministisch?: Nahezu deterministisch, aber einige Kernel-Patches sind nötig.
• FP8-Training möglich?: Aktuell wird mit BF16 trainiert, FP8-Unterstützung folgt bald.
• Unterstützung für Nicht-NVIDIA-GPUs?: Derzeit nur C/CUDA.

Meinung von GN⁺

• GPT-2 ist als Ursprung moderner LLMs ein äußerst wichtiges Modell. Auch GPT-3 und andere LLMs unterscheiden sich im Kern nicht grundlegend von GPT-2.
• Dieses Projekt ermöglicht es praktisch jedem, ein Modell auf GPT-2-Niveau selbst zu vertretbaren Kosten zu trainieren. Das dürfte sehr dabei helfen, das Verständnis für LLMs zu vertiefen.
• Allerdings ist es derzeit noch nicht für Inferenz optimiert, weshalb der praktische Einsatz in realen Services eingeschränkt ist. Auch Fine-Tuning zu einem dialogorientierten Modell wird nicht unterstützt.
• Aktuell werden nur NVIDIA-GPUs unterstützt; künftig wäre Unterstützung für verschiedene Plattformen wie AMD oder Apple Silicon wünschenswert.
• Ähnliche Open-Source-Projekte mit vergleichbarem Ziel sind Megatron-LM, DeepSpeed und FairSeq. Sie haben jeweils eigene Vor- und Nachteile und sollten je nach Einsatzzweck gewählt werden.
• Im Sinne eines lebendigen LLM-Entwicklungsökosystems ist dies ein äußerst ermutigendes Projekt. Die weitere Entwicklung bleibt spannend.