PyTorch 的源到源编译器。 易于理解。可检查。可扩展。

✅ 让 PyTorch 运行速度提升 40%   ✅ 量化                ✅ 核融合        
✅ 训练配方         ✅ FP4/FP6/FP8 精度       ✅ 分布式 TP/PP/DP 
✅ 推理配方        ✅ 适用于 NVIDIA Blackwell  ✅ CUDA 图          
✅ LLM、非 LLM 及更多  ✅ 自定义 Triton 内核       ✅ 组合以上所有功能

Thunder 是一个面向 PyTorch 的源到源深度学习编译器，专注于简化模型在训练和推理中的优化流程。

它提供了：

• 一种简单、符合 Python 风格的中间表示（IR），用于捕捉整个计算过程
• 一套丰富的转换系统，能够同时作用于计算 IR、模型和权重
• 一个可扩展的调度机制，用于连接融合器和优化的内核库

借助 Thunder，你可以：

• 轻松剖析深度学习程序，将单个操作映射到内核，并以交互方式检查程序
• 通过编程方式将一系列操作替换为优化后的版本，并观察对性能的影响
• 通过灵活扩展解释器，在不中断图的情况下获取完整的计算图
• 修改程序以充分利用特定硬件上的前沿内核库
• 编写适用于单 GPU 的模型，并将其转换为分布式运行模式
• 快速迭代混合精度和量化策略，寻找对模型质量影响最小的组合
• 将所有优化打包成可组合的配方，以便在不同模型系列之间移植

归根结底，你可以把 Thunder 看作是从“未优化”到“已优化”的高效工具。

如果你对此感兴趣，请继续阅读以安装 Thunder 并快速上手。

快速入门

通过 pip 安装 Thunder（更多选项）：

pip install lightning-thunder

pip install -U torch torchvision
pip install nvfuser-cu128-torch28 nvidia-cudnn-frontend  # 如果有 NVIDIA GPU

pip install lightning-thunder

pip install torch==2.7.0 torchvision==0.22
pip install nvfuser-cu128-torch27 nvidia-cudnn-frontend  # 如果有 NVIDIA GPU

pip install lightning-thunder

pip install torch==2.6.0 torchvision==0.21
pip install nvfuser-cu126-torch26 nvidia-cudnn-frontend  # 如果有 NVIDIA GPU

pip install lightning-thunder

pip install torch==2.5.0 torchvision==0.20
pip install nvfuser-cu124-torch25 nvidia-cudnn-frontend  # 如果有 NVIDIA GPU

# Float8 支持（这将从源代码编译，需耐心等待）
pip install "transformer_engine[pytorch]"

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

git clone https://github.com/Lightning-AI/lightning-thunder.git
cd lightning-thunder
pip install -e .

Hello world

定义一个函数或一个 PyTorch 模块：

import torch.nn as nn

model = nn.Sequential(nn.Linear(2048, 4096), nn.ReLU(), nn.Linear(4096, 64))

用 Thunder 对其进行优化：

import thunder
import torch

thunder_model = thunder.compile(model)

x = torch.randn(64, 2048)

y = thunder_model(x)

torch.testing.assert_close(y, model(x))

示例

LLM 训练

安装 LitGPT（不更新其他依赖项）

pip install --no-deps 'litgpt[all]'

然后运行：

import thunder
import torch
import litgpt

with torch.device("cuda"):
    model = litgpt.GPT.from_name("Llama-3.2-1B").to(torch.bfloat16)

thunder_model = thunder.compile(model)

inp = torch.ones((1, 2048), device="cuda", dtype=torch.int64)

out = thunder_model(inp)
out.sum().backward()

HuggingFace BERT 推理

安装 Hugging Face Transformers（推荐版本为 4.50.2 及以上）

pip install -U transformers

然后运行以下代码：

import thunder
import torch
import transformers

model_name = "bert-large-uncased"

tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)

with torch.device("cuda"):
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16
    )
    model.requires_grad_(False)
    model.eval()

    inp = tokenizer(["Hello world!"], return_tensors="pt")

thunder_model = thunder.compile(model)

out = thunder_model(**inp)
print(out)

HuggingFace DeepSeek R1 distill 推理

安装 Hugging Face Transformers（推荐版本为 4.50.2 及以上）

pip install -U transformers

然后运行以下代码：

import torch
import transformers
import thunder

model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"

tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)

with torch.device("cuda"):
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16
    )
    model.requires_grad_(False)
    model.eval()

    inp = tokenizer(["Hello world! Here's a long story"], return_tensors="pt")

thunder_model = thunder.compile(model)

out = thunder_model.generate(
    **inp, do_sample=False, cache_implementation="static", max_new_tokens=100
)
print(out)

Vision Transformer 推理

import thunder
import torch
import torchvision as tv

with torch.device("cuda"):
    model = tv.models.vit_b_16()
    model.requires_grad_(False)
    model.eval()

    inp = torch.randn(128, 3, 224, 224)

out = model(inp)

thunder_model = thunder.compile(model)

out = thunder_model(inp)

基准测试

尽管 Thunder 是一个用于优化模型的工具，而不是一个开箱即用就能带来加速的黑盒编译器，但这里仍提供一组基准测试结果。

从性能上看，Thunder 在开箱即用的情况下与 PyTorch 的 torch.compile 处于同一水平，尤其是在使用 CUDAGraphs 时。不过需要注意的是，Thunder 并不是 PyTorch torch.compile 的竞争对手！它实际上可以将 torch.compile 作为其融合执行器之一。

脚本 examples/quickstart/hf_llm.py 展示了如何对文本生成模型进行基准测试，包括前向传播、带损失的前向传播以及完整的前向 + 后向计算。

在配备 H100 显卡、PyTorch 版本为 2.8.0、nvFuser-cu128-torch28 和 Transformers 4.55.4 的环境下，运行 Llama 3.2 1B 模型时，我们观察到以下时间数据：

使用 torch.compile 和 CUDAGraphs（减少开销模式）：  521ms
使用 torch.compile 但不使用 CUDAGraphs（默认模式）：       814ms
不使用 torch.compile：                                    1493ms
使用 Thunder 和 CUDAGraphs：                                                542ms

插件

插件是一种对模型应用优化的方式，例如并行化和量化。

Thunder 自带一些插件，但也很容易编写新的插件。

• 使用 DDP、FSDP、TP 等分布式策略进行扩展
• 使用 FP8、MXFP8 优化数值精度
• 通过量化节省显存
• 使用 CUDAGraphs 减少延迟
• 调试和性能分析

例如，为了通过 CUDAGraphs 减少 CPU 开销，可以在 thunder.compile 的 plugins= 参数中添加 "reduce-overhead"：

thunder_model = thunder.compile(model, plugins="reduce-overhead")

这可能会或可能不会带来显著的效果。Thunder 的优势在于，你可以轻松地切换不同的优化方法，探索最适合你环境的组合。

工作原理

Thunder 的工作分为三个阶段：

1. ⚡️ 它通过解释 Python 字节码来获取你的模型，并生成一段线性 Python 程序。
2. ⚡️ 它转换模型和计算轨迹，以实现分布式计算或改变精度。
3. ⚡️ 它将轨迹的不同部分路由到相应的执行方式：
   • 融合（如 NVFuser、torch.compile）
   • 专用库（如 cuDNN SDPA、TransformerEngine）
   • 自定义 Triton 和 CUDA 内核
   • PyTorch eager 操作

以下是针对一个简单 MLP 的轨迹示例：

import thunder
import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(1024, 2048), nn.ReLU(), nn.Linear(2048, 256))

thunder_model = thunder.compile(model)
y = thunder_model(torch.randn(4, 1024))

print(thunder.last_traces(thunder_model)[-1])

这是获取到的轨迹，已准备好进行转换和执行：

def computation(input, t_0_bias, t_0_weight, t_2_bias, t_2_weight):
# input: "cuda:0 f32[4, 1024]"
# t_0_bias: "cuda:0 f32[2048]"
# t_0_weight: "cuda:0 f32[2048, 1024]"
# t_2_bias: "cuda:0 f32[256]"
# t_2_weight: "cuda:0 f32[256, 2048]"
t3 = ltorch.linear(input, t_0_weight, t_0_bias) # t3: "cuda:0 f32[4, 2048]"
t6 = ltorch.relu(t3, False) # t6: "cuda:0 f32[4, 2048]"
t10 = ltorch.linear(t6, t_2_weight, t_2_bias) # t10: "cuda:0 f32[4, 256]"
return (t10,)

请注意，Thunder 的中间表示形式仅仅是 Python 的子集！

性能

Thunder 非常快速。以下是使用 LitGPT 在 H100 和 B200 硬件上进行预训练任务时，相对于 PyTorch eager 所获得的加速效果。

社区

Thunder 是一个开源项目，由社区协作开发，并得到了 NVIDIA 的重要贡献。 ⚡️ 关于贡献 Thunder 的文档 💬 在 Discord 上寻求帮助 📋 许可证：Apache 2.0

Lightning Thunder 快速上手指南

Lightning Thunder 是一个专为 PyTorch 设计的源码到源码（source-to-source）深度学习编译器。它旨在通过简单的 Python 接口，让开发者轻松优化模型的训练和推理性能，支持内核融合、量化、分布式策略以及最新硬件特性（如 NVIDIA Blackwell）。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: 3.10 或更高版本
• GPU: 推荐使用 NVIDIA GPU (需安装对应的 CUDA 驱动)
• PyTorch: 建议安装最新版本 (2.8.0+) 以获得最佳兼容性，同时也支持 2.5 - 2.7 版本。

安装步骤

1. 基础安装

使用 pip 安装 Lightning Thunder 及最新的 PyTorch 版本：

pip install lightning-thunder

pip install -U torch torchvision

2. 可选依赖 (针对 NVIDIA GPU)

如果您使用 NVIDIA GPU，建议安装 nvfuser 和 cudnn-frontend 以启用内核融合等高级优化功能。请根据您的 PyTorch 版本选择对应的命令：

对于 PyTorch 2.8 + CUDA 12.8:

pip install nvfuser-cu128-torch28 nvidia-cudnn-frontend

对于旧版本 PyTorch:

• PyTorch 2.7 + CUDA 12.6:

  pip install torch==2.7.0 torchvision==0.22
  pip install nvfuser-cu126-torch26 nvidia-cudnn-frontend
  

• PyTorch 2.6 + CUDA 12.4:

  pip install torch==2.6.0 torchvision==0.21
  pip install nvfuser-cu124-torch25 nvidia-cudnn-frontend
  

提示: 国内用户若遇到下载缓慢，可添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像源。

3. 高级功能 (可选)

如需启用 FP8 精度支持或其他前沿功能：

# 安装 Float8 支持 (可能需要从源码编译，耗时较长)
pip install "transformer_engine[pytorch]"

# 或者安装开发版
pip install git+https://github.com/Lightning-AI/lightning-thunder.git@main

基本使用

Thunder 的使用非常直观，只需导入库并使用 thunder.compile 包装您的 PyTorch 模型即可。

Hello World 示例

以下示例展示如何定义一个简单的神经网络并通过 Thunder 进行编译和优化：

import torch.nn as nn
import thunder
import torch

# 1. 定义标准的 PyTorch 模型
model = nn.Sequential(
    nn.Linear(2048, 4096), 
    nn.ReLU(), 
    nn.Linear(4096, 64)
)

# 2. 使用 Thunder 编译模型
thunder_model = thunder.compile(model)

# 3. 准备输入数据
x = torch.randn(64, 2048)

# 4. 运行优化后的模型
y = thunder_model(x)

# 5. 验证输出结果是否与原始模型一致
torch.testing.assert_close(y, model(x))

进阶用法：应用插件

Thunder 支持通过插件轻松应用各种优化策略，例如使用 CUDA Graphs 减少 CPU 开销：

# 启用 "reduce-overhead" 插件以减少启动开销
thunder_model = thunder.compile(model, plugins="reduce-overhead")

您现在可以像使用普通 PyTorch 模型一样使用 thunder_model，但它将自动利用底层的高效算子和融合技术。