lightning thunder

让 PyTorch 模型快如闪电。

---

Lightning.ai • 性能 • 入门 • 安装 • 示例 • Thunder 内部 • 参与其中！ • 文档

欢迎来到⚡闪电雷霆

Thunder 使 PyTorch 模型快如闪电。

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

它支持单 GPU 和多 GPU 配置。 Thunder 的目标是可用、可理解和可扩展。

笔记

闪电雷霆处于阿尔法阶段。请随意参与，但预计一路上会遇到一些坎坷。

单GPU性能

通过优化的复合效应和使用一流的执行器，Thunder 可以比标准的非编译 PyTorch 代码（“PyTorch eager”）实现显着的加速。下图显示了 LitGPT 中实现的 Llama 2 7B 的预训练吞吐量。

如上图所示，与使用 nvFuser、torch.compile、cuDNN 和 TransformerEngine FP8 等执行器组合的 H100 上的 eager 代码相比，Thunder 的训练吞吐量提高了 40%。

多GPU性能

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

开始使用

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

安装迅雷

Thunder 处于 alpha 阶段，最新的开发正在main分支上进行。您可以从main分支安装最新版本的 Thunder，如下所示：

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

为了获得最佳性能，您可以安装 Thunder 并附加以下依赖项：

• 安装为 PyTorch 2.5.1 构建的预发布 nvFuser，如下所示：

 # install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25

• 安装cudnn如下：

 # install cudnn
pip install nvidia-cudnn-frontend

高级安装选项

安装以修补和贡献

如果您有兴趣修改 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 经过了非常彻底的测试，因此预计这需要一段时间。

你好世界

下面是 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 )

# prints
# tensor(
#  [[4, 4]
#   [4, 4]])

编译后的函数jfoo接受并返回 PyTorch 张量，就像原始函数一样，因此 Thunder 编译的模块和函数可以用作更大的 PyTorch 程序的一部分。

火车模型

Thunder 还处于早期阶段，尚不应该用于生产运行。

然而，它已经可以为 LitGPT 支持的 LLM 预训练和微调提供出色的性能，例如 Mistral、Llama 2、Gemma、Falcon 等。

查看 LitGPT 集成，了解如何一起运行 LitGPT 和 Thunder。

Thunder 内部：核心功能简介

给定一个 Python 可调用或 PyTorch 模块，Thunder 可以生成一个优化的程序：

• 计算其前向和后向传递
• 将操作合并到高效的融合区域
• 将计算分派给优化的内核
• 在机器之间最佳地分配计算

为此，Thunder 附带了：

• 用于获取针对 PyTorch 和自定义操作的 Python 程序的 JIT
• 多级中间表示 (IR)，将操作表示为缩减操作集的踪迹
• 计算图轨迹上的一组可扩展转换，例如grad 、 fusions 、分布式（如ddp 、 fsdp ）、函数（如vmap 、 vjp 、 jvp ）
• 一种将操作分派给可扩展执行器集合的方法

Thunder 完全用 Python 编写。甚至它的踪迹在转换的所有阶段都被表示为有效的 Python。这使得内省和可扩展性达到前所未有的水平。

Thunder 不会直接为 GPU 等加速器生成代码。它获取并转换用户程序，以便可以使用快速执行器以最佳方式选择或生成设备代码，例如：

• 火炬编译
• 热熔器
• CUDNN
• 顶点
• 变压器引擎
• PyTorch 渴望
• 通过 PyCUDA、Numba、CuPy 自定义 CUDA 内核
• 用 OpenAI Triton 编写的自定义内核

使用 Thunder 编译的模块和函数与 vanilla PyTorch 完全互操作，并支持 PyTorch 的 autograd。此外，Thunder 与 torch.compile 一起工作，以利用其最先进的优化。

文档

提供在线文档。要在本地构建文档，您可以使用

make docs

并将浏览器指向docs/build/index.html处生成的文档。

介入！

我们感谢您的反馈和贡献。如果您有功能请求、问题，或者想要贡献代码或配置文件，请随时使用 GitHub 问题跟踪器。

我们欢迎所有个人贡献者，无论他们的经验或硬件水平如何。您的贡献很有价值，我们很高兴看到您在这个协作和支持的环境中能够取得什么成就。

执照

Lightning Thunder 在 Apache 2.0 许可证下发布。有关详细信息，请参阅许可证文件。