欢迎使用 ⚡ 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
高级安装选项
从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 经过非常彻底的测试,所以预计这将需要一段时间。
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]])
编译后的函数 jfoo
接受并返回 PyTorch 张量,就像原始函数一样,因此由 Thunder 编译的模块和函数可以作为更大的 PyTorch 程序的一部分使用。
训练模型
Thunder 仍处于早期阶段,尚不应用于生产运行。
然而,它已经能够为 LitGPT 支持的 LLM(如 Mistral、Llama 2、Gemma、Falcon 等)的预训练和微调提供卓越的性能。
查看 LitGPT 集成 以了解如何一起运行 LitGPT 和 Thunder。
深入 Thunder:核心功能简介
给定一个 Python 可调用对象或 PyTorch 模块,Thunder 可以生成一个优化的程序,该程序:
- 计算其前向和反向传播
- 将操作合并为高效的融合区域
- 将计算分派到优化的内核
- 在机器之间优化分配计算
为此,Thunder 提供了:
- 一个 JIT,用于获取针对 PyTorch 和自定义操作的 Python 程序
- 一个多层中间表示(IR),用于将操作表示为简化操作集的跟踪
- 一组可扩展的计算图跟踪转换,如
grad
、融合、分布式(如ddp
、fsdp
)、函数式(如vmap
、vjp
、jvp
) - 一种将操作分派到可扩展的执行器集合的方法
Thunder 完全用 Python 编写。即使在所有转换阶段,其跟踪也表示为有效的 Python。这允许前所未有的内省和可扩展性水平。
Thunder 不直接为加速器(如 GPU)生成代码。它获取并转换用户程序,以便可以使用快速执行器最优地选择或生成设备代码,如:
- torch.compile
- nvFuser
- cuDNN
- Apex
- TransformerEngine
- PyTorch eager
- 通过 PyCUDA、Numba、CuPy 的自定义 CUDA 内核
- 用 OpenAI Triton 编写的自定义内核
用 Thunder 编译的模块和函数与原生 PyTorch 完全互操作,并支持 PyTorch 的自动微分。此外,Thunder 可以与 torch.compile 一起工作,以利用其最先进的优化。
文档
在线文档 可用。要在本地构建文档,您可以使用
make docs
并将浏览器指向 docs/build/index.html
生成的文档。
参与其中!
我们感谢您的反馈和贡献。如果您有功能请求、问题或想贡献代码或配置文件,请不要犹豫,使用 GitHub Issue 跟踪器。
我们欢迎所有个人贡献者,无论他们的经验水平或硬件如何。您的贡献都是宝贵的,我们很高兴看到您在这个协作和支持的环境中能够取得什么成就。
许可证
Lightning Thunder 根据 Apache 2.0 许可发布。 详情请参阅 LICENSE 文件。