lightning thunder
Initial release
Créez des modèles PyTorch ultra-rapides.
Lightning.ai • Performances • Démarrez • Installez • Exemples • À l'intérieur de Thunder • Impliquez-vous ! •Documents
Thunder rend les modèles PyTorch ultra-rapides.
Thunder est un compilateur source à source pour PyTorch. Il rend les programmes PyTorch plus rapides en combinant et en utilisant différents exécuteurs matériels à la fois (par exemple, nvFuser, torch.compile, cuDNN et TransformerEngine FP8).
Il prend en charge les configurations mono et multi-GPU. Thunder vise à être utilisable, compréhensible et extensible.
Note
Lightning Thunder est en version alpha. N'hésitez pas à vous impliquer, mais attendez-vous à quelques obstacles en cours de route.
Thunder peut atteindre des accélérations significatives par rapport au code PyTorch standard non compilé (« PyTorch avide »), grâce aux effets cumulatifs des optimisations et à l'utilisation des meilleurs exécuteurs de leur catégorie. La figure ci-dessous montre le débit de pré-entraînement pour Llama 2 7B tel qu'implémenté dans LitGPT.
Comme le montre le graphique ci-dessus, Thunder atteint une accélération de 40 % du débit de formation par rapport au code impatient sur H100 en utilisant une combinaison d'exécuteurs, notamment nvFuser, torch.compile, cuDNN et TransformerEngine FP8.
Thunder prend également en charge les stratégies distribuées telles que DDP et FSDP pour la formation de modèles sur plusieurs GPU. Le graphique suivant affiche le débit normalisé mesuré pour Llama 2 7B sans précision mixte FP8 ; le soutien au FSDP est en cours.
Le moyen le plus simple de démarrer avec Thunder, sans nécessiter aucune installation ou configuration supplémentaire, consiste à utiliser notre studio de didacticiels Zero to Thunder.
Thunder est en alpha et le dernier développement est en cours sur la branche main . Vous pouvez installer la dernière version de Thunder à partir de la branche main comme suit :
pip install git+https://github.com/Lightning-AI/lightning-thunder.git@mainPour obtenir les meilleures performances, vous pouvez installer Thunder avec les dépendances supplémentaires suivantes :
# install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25 # install cudnn
pip install nvidia-cudnn-frontendSi vous souhaitez bricoler et contribuer à Thunder, nous vous recommandons de cloner le référentiel Thunder et de l'installer en mode modifiable de pip :
git clone https://github.com/Lightning-AI/lightning-thunder.git
cd lightning-thunder
pip install -e .Après avoir cloné le référentiel Lightning-Thunder et l'avoir installé en tant que package modifiable comme expliqué ci-dessus, vous pouvez configurer votre environnement pour développer Thunder en installant les exigences de développement :
pip install -r requirements/devel.txtMaintenant, vous exécutez des tests :
pytest thunder/testsThunder est testé de manière très approfondie, alors attendez-vous à ce que cela prenne un certain temps.
Vous trouverez ci-dessous un exemple simple de la façon dont Thunder vous permet de compiler et d'exécuter du code PyTorch :
import torch
import thunder
def foo ( a , b ):
return a + b
jfoo = thunder . jit ( foo )
a = torch . full (( 2 , 2 ), 1 )
b = torch . full (( 2 , 2 ), 3 )
result = jfoo ( a , b )
print ( result )
# prints
# tensor(
# [[4, 4]
# [4, 4]]) La fonction compilée jfoo prend et renvoie les tenseurs PyTorch, tout comme la fonction d'origine, de sorte que les modules et fonctions compilés par Thunder peuvent être utilisés dans le cadre de programmes PyTorch plus vastes.
Thunder en est à ses débuts et ne devrait pas encore être utilisé pour des séries de production.
Cependant, il peut déjà offrir des performances exceptionnelles pour le pré-entraînement et le réglage fin des LLM pris en charge par LitGPT, tels que Mistral, Llama 2, Gemma, Falcon et autres.
Découvrez l'intégration LitGPT pour en savoir plus sur l'exécution conjointe de LitGPT et Thunder.
Étant donné un module appelable Python ou PyTorch, Thunder peut générer un programme optimisé qui :
Pour ce faire, Thunder est livré avec :
grad , fusions, distribuées (comme ddp , fsdp ), fonctionnelles (comme vmap , vjp , jvp )Thunder est entièrement écrit en Python. Même sa trace est représentée comme Python valide à toutes les étapes de transformation. Cela permet des niveaux d’introspection et d’extensibilité sans précédent.
Thunder ne génère pas directement de code pour les accélérateurs, tels que les GPU. Il acquiert et transforme les programmes utilisateur afin qu'il soit possible de sélectionner ou de générer de manière optimale le code de l'appareil à l'aide d'exécuteurs rapides tels que :
Les modules et fonctions compilés avec Thunder interagissent entièrement avec Vanilla PyTorch et prennent en charge l'autograd de PyTorch. De plus, Thunder fonctionne aux côtés de torch.compile pour tirer parti de ses optimisations de pointe.
Une documentation en ligne est disponible. Pour créer de la documentation localement, vous pouvez utiliser
make docs et pointez votre navigateur vers les documents générés à l' docs/build/index.html .
Nous apprécions vos commentaires et vos contributions. Si vous avez des demandes de fonctionnalités, des questions ou si vous souhaitez contribuer au code ou aux fichiers de configuration, n'hésitez pas à utiliser le outil de suivi des problèmes GitHub.
Nous accueillons tous les contributeurs individuels, quel que soit leur niveau d'expérience ou de matériel. Vos contributions sont précieuses et nous sommes ravis de voir ce que vous pouvez accomplir dans cet environnement collaboratif et solidaire.
Lightning Thunder est publié sous la licence Apache 2.0. Voir le fichier LICENSE pour plus de détails.
