lightning thunder
Initial release
Machen Sie PyTorch-Modelle blitzschnell.
Lightning.ai • Leistung • Erste Schritte • Installation • Beispiele • Inside Thunder • Machen Sie mit! • Dokumentation
Thunder macht PyTorch-Modelle blitzschnell.
Thunder ist ein Source-to-Source-Compiler für PyTorch. Es macht PyTorch-Programme schneller, indem es verschiedene Hardware-Ausführer gleichzeitig kombiniert und verwendet (z. B. nvFuser, Torch.compile, cuDNN und TransformerEngine FP8).
Es unterstützt sowohl Einzel- als auch Multi-GPU-Konfigurationen. Thunder soll benutzerfreundlich, verständlich und erweiterbar sein.
Notiz
Lightning Thunder befindet sich in der Alpha-Phase. Machen Sie ruhig mit, aber rechnen Sie mit ein paar Überraschungen auf dem Weg.
Thunder kann im Vergleich zu nicht kompiliertem Standard-PyTorch-Code („PyTorch Eager“) erhebliche Geschwindigkeitssteigerungen erzielen, indem die Effekte von Optimierungen und die Verwendung erstklassiger Executoren kombiniert werden. Die folgende Abbildung zeigt den Vortrainingsdurchsatz für Llama 2 7B, wie er in LitGPT implementiert ist.
Wie im Diagramm oben gezeigt, erreicht Thunder mithilfe einer Kombination von Executoren, darunter nvFuser, Torch.compile, cuDNN und TransformerEngine FP8, eine Beschleunigung des Trainingsdurchsatzes um 40 % im Vergleich zu Eager-Code auf H100.
Thunder unterstützt auch verteilte Strategien wie DDP und FSDP für das Training von Modellen auf mehreren GPUs. Das folgende Diagramm zeigt den normalisierten Durchsatz, der für Llama 2 7B ohne gemischte FP8-Präzision gemessen wurde. Die Unterstützung für FSDP ist im Gange.
Der einfachste Weg, mit Thunder zu beginnen, ohne dass zusätzliche Installationen oder Setups erforderlich sind, ist die Verwendung unseres Zero to Thunder Tutorial Studio.
Thunder befindet sich in der Alpha-Phase und die neueste Entwicklung findet im main statt. Sie können die neueste Version von Thunder wie folgt vom main aus installieren:
pip install git+https://github.com/Lightning-AI/lightning-thunder.git@mainUm die beste Leistung zu erzielen, können Sie Thunder mit den folgenden zusätzlichen Abhängigkeiten installieren:
# install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25 # install cudnn
pip install nvidia-cudnn-frontendWenn Sie daran interessiert sind, an Thunder herumzubasteln und dazu beizutragen, empfehlen wir Ihnen, das Thunder-Repository zu klonen und es im bearbeitbaren Modus von pip zu installieren:
git clone https://github.com/Lightning-AI/lightning-thunder.git
cd lightning-thunder
pip install -e .Nachdem Sie das Lightning-Thunder-Repository geklont und wie oben erläutert als bearbeitbares Paket installiert haben, können Sie Ihre Umgebung für die Entwicklung von Thunder einrichten, indem Sie die Entwicklungsanforderungen installieren:
pip install -r requirements/devel.txtJetzt führen Sie Tests durch:
pytest thunder/testsThunder wurde sehr gründlich getestet, es kann also damit rechnen, dass dies eine Weile dauern wird.
Nachfolgend finden Sie ein einfaches Beispiel dafür, wie Sie mit Thunder PyTorch-Code kompilieren und ausführen können:
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]]) Die kompilierte Funktion jfoo nimmt PyTorch-Tensoren entgegen und gibt sie zurück, genau wie die Originalfunktion, sodass von Thunder kompilierte Module und Funktionen als Teil größerer PyTorch-Programme verwendet werden können.
Thunder befindet sich in einem frühen Stadium und sollte noch nicht für Produktionsläufe verwendet werden.
Es kann jedoch bereits eine hervorragende Leistung für das Vortraining und die Feinabstimmung von LLMs liefern, die von LitGPT unterstützt werden, wie Mistral, Llama 2, Gemma, Falcon und andere.
Schauen Sie sich die LitGPT-Integration an, um mehr über die gemeinsame Ausführung von LitGPT und Thunder zu erfahren.
Mit einem aufrufbaren Python- oder PyTorch-Modul kann Thunder ein optimiertes Programm generieren, das:
Zu diesem Zweck wird Thunder mit Folgendem geliefert:
grad , Fusionen, verteilt (wie ddp , fsdp ), funktional (wie vmap , vjp , jvp ).Thunder ist vollständig in Python geschrieben. Sogar seine Spur wird in allen Phasen der Transformation als gültiges Python dargestellt. Dies ermöglicht ein beispielloses Maß an Selbstbeobachtung und Erweiterbarkeit.
Thunder generiert keinen direkten Code für Beschleuniger wie GPUs. Es erfasst und transformiert Benutzerprogramme, sodass es möglich ist, Gerätecode mithilfe schneller Ausführer wie den folgenden optimal auszuwählen oder zu generieren:
Mit Thunder kompilierte Module und Funktionen funktionieren vollständig mit Vanilla PyTorch und unterstützen PyTorchs Autograd. Außerdem arbeitet Thunder mit Torch.compile zusammen, um seine hochmodernen Optimierungen zu nutzen.
Eine Online-Dokumentation ist verfügbar. Um die Dokumentation lokal zu erstellen, können Sie verwenden
make docs und verweisen Sie Ihren Browser auf die generierten Dokumente unter docs/build/index.html .
Wir freuen uns über Ihr Feedback und Ihre Beiträge. Wenn Sie Funktionswünsche oder Fragen haben oder Code- oder Konfigurationsdateien beisteuern möchten, zögern Sie bitte nicht, den GitHub Issue Tracker zu verwenden.
Wir begrüßen alle einzelnen Mitwirkenden, unabhängig von ihrem Erfahrungsstand oder ihrer Hardware. Ihre Beiträge sind wertvoll und wir sind gespannt, was Sie in diesem kollaborativen und unterstützenden Umfeld erreichen können.
Lightning Thunder wird unter der Apache 2.0-Lizenz veröffentlicht. Einzelheiten finden Sie in der LICENSE-Datei.
