Réimplémenter à bas coût le modèle GPT-2 avec llm.c

• Explication de la méthode pour reproduire le modèle GPT-2 (124M) avec llm.c en 90 minutes pour 20 $
• GPT-2 (124M) est le plus petit modèle publié par OpenAI en 2019
• En utilisant un nœud Lambda équipé de 8X A100 80GB SXM, le coût est d’environ 14 $/heure, soit environ 20 $ au total
  • L’entraînement est aussi possible sur un seul GPU, mais cela prend plus de temps (4 à 24 heures)

Comparaison des résultats

• Sur le jeu de validation FineWeb, les performances sont meilleures que celles du checkpoint publié par OpenAI
  • Cela dit, GPT-2 a été entraîné sur WebText, la comparaison n’est donc pas totalement équitable
• La précision sur HellaSwag a également été mesurée, avec un score de 29,9, proche des 33,7 de GPT-3 Small (124M)
  • Le score de GPT-2 (124M), à 29,4, est déjà dépassé
  • Ici toutefois, l’entraînement a été fait sur 10B tokens, contre 300B tokens pour GPT-3

Configuration minimale

• GPU requis (Lambda Labs recommandé)
• Guide basé sur Linux x86 64 bits Ubuntu 22.04 avec CUDA 12
• Installation de miniconda puis de la version nightly de PyTorch (optionnel)
• Installation des paquets nécessaires pour le tokenizer
• Installation de cuDNN pour de meilleures performances (optionnel)
• Installation de MPI en cas d’utilisation de plusieurs GPU (optionnel)
• Prétraitement du jeu de données FineWeb 10B tokens (~1 heure)
• Compilation de llm.c (précision mixte, utilisation de cuDNN FlashAttention)

Lancer l’entraînement

• Exemple de commande pour l’utilisation sur un seul GPU
• Exécution via mpirun pour une configuration multi-GPU (8 GPU)
• Explication des principaux arguments
  • -i, -j : chemins des données d’entraînement / validation
  • -o : chemin de sauvegarde des logs et checkpoints
  • -e : initialisation du modèle (GPT-2 depth 12)
  • -b : taille du micro-batch (à réduire en cas de manque de mémoire)
  • -t : longueur maximale de séquence
  • -d : taille totale du batch (voir l’article GPT-3)
  • -r : paramètre Recompute (économie de mémoire)
  • -z : ZeRO-1 (sharding de l’état de l’optimiseur)
  • Autres paramètres comme le weight decay, le learning rate, la fréquence des checkpoints, etc.

Processus d’entraînement

• Avec 10B tokens d’entraînement et une taille de batch de 0.5M, environ 20K étapes sont prévues
• Sur un GPU A100 40GB PCIe, le temps par étape, le MFU et le débit de traitement des tokens sont affichés
• Au début de l’entraînement, un phénomène de gradient exploding apparaît, mais il est résolu par clipping

Visualisation

• Un notebook Jupyter est fourni pour parser les logs et visualiser les courbes d’entraînement

Tokenizer

• Nécessaire pour convertir des tokens entiers en chaînes de caractères
• Peut être généré via un script PyTorch

Sampling

• Pour l’instant, ce n’est pas optimisé pour l’inférence
• Avec quelques modifications de code, le sampling unconditional/conditional est possible

Structure du code

• L’essentiel de l’implémentation se trouve dans le fichier train_gpt2.cu
• Les 500 premières lignes concernent la configuration de MPI, NCCL, cuDNN, cuBLAS, etc.
• Les 1500 lignes suivantes couvrent le forward/backward du Transformer
• Les 1000 lignes suivantes implémentent le modèle GPT-2
• Les 1000 dernières lignes couvrent la boucle d’entraînement, le parsing des arguments, etc.

Modèle 350M

• 10B tokens ne suffisent pas, 30B tokens sont utilisés
• Avec 8X A100 80GB, l’entraînement dure 14 heures pour un coût d’environ 200 $

FAQ

• Sampling possible ? : Oui, mais c’est inefficace.
• Possibilité de chat ? : Pour l’instant, seul le préentraînement est possible ; le fine-tuning pour le chat n’est pas pris en charge.
• Entraînement distribué multi-nœuds : Possible, mais pas encore testé.
• Déterminisme au bit près : Presque déterministe, mais certains patchs de kernel sont nécessaires.
• Entraînement en FP8 possible ? : Actuellement l’entraînement se fait en BF16, le support FP8 est prévu bientôt.
• Prise en charge des GPU non NVIDIA ? : Pour l’instant, seul C/CUDA est pris en charge.

L’avis de GN⁺

• GPT-2 est un modèle extrêmement important comme point de départ des LLM modernes. Par la suite, GPT-3 et les autres LLM ne diffèrent pas fondamentalement de GPT-2.
• Ce projet permet à tout le monde d’entraîner directement un modèle de niveau GPT-2 à un coût raisonnable. Cela devrait beaucoup aider à mieux comprendre les LLM.
• En revanche, comme ce n’est pas encore optimisé pour l’inférence, son utilisation dans un service réel reste limitée. Le fine-tuning en modèle conversationnel n’est pas non plus pris en charge.
• À l’heure actuelle, seuls les GPU NVIDIA sont pris en charge, mais on peut espérer à l’avenir un support de plateformes variées comme AMD ou Apple Silicon.
• Parmi les projets open source poursuivant des objectifs similaires, on peut citer Megatron-LM, DeepSpeed et FairSeq. Chacun a ses avantages et ses inconvénients, et le choix dépendra de l’usage visé.
• C’est un projet très encourageant pour dynamiser l’écosystème de développement des LLM. La suite est très prometteuse.