让PyTorch模型快如闪电。

<p align="center"> <a href="https://lightning.ai/">Lightning.ai</a> • <a href="#performance">性能</a> • <a href="#get-started">开始使用</a> • <a href="#install-thunder">安装</a> • <a href="#hello-world">示例</a> • <a href="#inside-thunder-a-brief-look-at-the-core-features">Thunder内部</a> • <a href="#get-involved">参与其中</a> • <a href="https://lightning-thunder.readthedocs.io/en/latest/">文档</a> </p>

</div>

欢迎使用 ⚡ Lightning Thunder

Thunder让PyTorch模型快如闪电。

Thunder是PyTorch的源到源编译器。它通过结合并同时使用不同的硬件执行器（例如，nvFuser、torch.compile、cuDNN和TransformerEngine FP8）来加速PyTorch程序。

它支持单GPU和多GPU配置。 Thunder旨在易用、易懂且可扩展。

[!Note] Lightning Thunder处于alpha阶段。欢迎参与其中，但请做好遇到一些问题的准备。

单GPU性能

相比标准的非编译PyTorch代码（"PyTorch eager"），Thunder可以通过优化的叠加效应和使用最佳执行器实现显著的加速。下图展示了LitGPT中实现的Llama 2 7B模型的预训练吞吐量。

如上图所示，在H100上使用包括nvFuser、torch.compile、cuDNN和TransformerEngine FP8在内的执行器组合，Thunder相比eager代码在训练吞吐量上实现了40%的加速。

多GPU性能

Thunder还支持分布式策略，如DDP和FSDP，用于在多个GPU上训练模型。以下图表显示了Llama 2 7B模型在不使用FP8混合精度的情况下测得的标准化吞吐量；对FSDP的支持正在进行中。

开始使用

开始使用Thunder最简单的方法是使用我们的Zero to Thunder教程工作室，无需额外安装或设置。

安装Thunder

要在本地机器上使用Thunder：

按以下方式一起安装nvFuser每日版和PyTorch每日版：

# 安装nvFuser，它会安装匹配的每日版PyTorch
pip install --pre 'nvfuser-cu121[torch]' --extra-index-url https://pypi.nvidia.com

按以下方式安装cudnn：

# 安装cudnn
pip install nvidia-cudnn-frontend

最后，按以下方式安装Thunder：

# 安装thunder
pip install lightning-thunder

<details> <summary>高级安装选项</summary>

从main分支安装

或者，您可以直接从此GitHub仓库安装Thunder的最新版本，方法如下：

# 1) 安装nvFuser和PyTorch每日版依赖：
pip install --pre 'nvfuser-cu121[torch]' --extra-index-url https://pypi.nvidia.com

# 2) 安装 Thunder 本身
pip install git+https://github.com/Lightning-AI/lightning-thunder.git

安装以进行修改和贡献

如果您有兴趣修改和为 Thunder 做出贡献，我们建议克隆 Thunder 仓库并以 pip 的可编辑模式安装：

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

开发和运行测试

在按上述方法克隆 lightning-thunder 仓库并将其安装为可编辑包后，您可以通过安装开发要求来设置用于开发 Thunder 的环境：

pip install -r requirements/devel.txt

现在您可以运行测试：

pytest thunder/tests

Thunder 经过非常彻底的测试，所以预计这将需要一段时间。

</details>

Hello World

以下是一个简单的示例，展示了 Thunder 如何允许您编译和运行 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)

# 输出
# tensor(
#  [[4, 4]
#   [4, 4]])