Make PyTorch models Lightning fast.

Lightning.ai • Performance • Get started • Install • Examples • Features • Documentation •

Welcome to ⚡ Lightning Thunder

Thunder makes PyTorch models Lightning fast.

Thunder is a source-to-source compiler for PyTorch. It makes PyTorch programs faster by combining and using different hardware executors at once (ie: nvFuser, torch.compile, cuDNN, and TransformerEngine FP8).

Works on single accelerators and in multi-GPU settings. Thunder aims to be usable, understandable, and extensible.

Performance

Thunder can achieve significant speedups over standard PyTorch eager code, through the compounding effects of optimizations and the use of best-in-class executors. Here is an example of the pretraining throughput for Llama 2 7B as implemented in LitGPT.

Thunder achieves a 40% speedup in training throughput compared to eager code on H100 using a combination of executors including nvFuser, torch.compile, cuDNN, and TransformerEngine FP8.

Thunder supports distributed strategies like DDP and FSDP (ZeRO2 and ZeRO3). Here is the normalized throughput measured for Llama 2 7B (this time without FP8 mixed precision, support for FSDP is underway).

NOTE: Lightning Thunder is alpha. Feel free to get involved, expect a few bumps along the way.

Get started

Try Thunder without installing by using our Zero to Thunder Tutorial Studio.

Install Thunder

Install nvFuser nightly, and Thunder together

# install nvFuser which installs the matching nightly PyTorch
pip install --pre 'nvfuser-cu121[torch]' --extra-index-url https://pypi.nvidia.com

# install thunder
pip install lightning-thunder

Advanced install options

Install from main

pip install git+https://github.com/Lightning-AI/lightning-thunder.git

Install to tinker and contribute

Install this way to tinker with the internals and contribute:

pip install -e .

Hello World

Here is a simple example of how Thunder lets you compile and run PyTorch code:

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]])

The compiled function jfoo takes and returns PyTorch tensors, just like the original function, so modules and functions compiled by Thunder can be used as part of larger PyTorch programs.

Train models

Thunder is in its early stages and should not be used for production runs yet.

However, it can already deliver outstanding performance on LLM model supported by LitGPT, such as Mistral, Llama 2, Gemma, Falcon, and others.

Check out the LitGPT integration to learn about running LitGPT and Thunder together.

Features

Given a Python callable or PyTorch module, Thunder can generate an optimized program that:

Computes its forward and backward passes
Coalesces operations into efficient fusion regions
Dispatches computations to optimized kernels
Distributes computations optimally across machines

To do so, Thunder ships with:

A JIT for acquiring Python programs targeting PyTorch and custom operations
A multi-level IR to represent operations as a trace of a reduced op-set
An extensible set of transformations on the trace, such as grad, fusions, distributed (like ddp, fsdp), functional (like vmap, vjp, jvp)
A way to dispatch operations to an extensible collection of executors

Thunder is written entirely in Python. Even its trace is represented as valid Python at all stages of transformation. This allows unprecedented levels of introspection and extensibility.

Thunder doesn't generate code for accelerators directly. It acquires and transforms user programs so that it's possible to optimally select or generate device code using fast executors like:

torch.compile
nvFuser
cuDNN
Apex
TransformerEngine
PyTorch eager
custom kernels, including those written with OpenAI Triton

Modules and functions compiled with Thunder fully interoperate with vanilla PyTorch and support PyTorch's autograd. Also, Thunder works alongside torch.compile to leverage its state-of-the-art optimizations.

Documentation

Docs are currently not hosted publicly. However you can build them locally really quickly:

make docs

and point your browser to the generated docs at docs/build/index.html.

Develop and run tests

You can set up your environment for developing Thunder by installing the development requirements:

pip install -r requirements/devel.txt

Install Thunder as an editable package (optional):

pip install -e .

Now you run tests:

pytest thunder/tests

Thunder is very thoroughly tested, so expect this to take a while.

License

Lightning Thunder is released under the Apache 2.0 license. See the LICENSE file for details.

Hartmon8/lightning-thunder