Project Icon

lightning-thunder

PyTorch模型优化编译器 显著提升训练速度

Lightning Thunder是一款专为PyTorch设计的源到源编译器。它通过整合nvFuser、torch.compile、cuDNN等多种硬件执行器,大幅提升PyTorch程序的执行效率。支持单GPU和多GPU环境,在Llama 2 7B模型训练中实现40%的吞吐量提升。Thunder具有易用性、可理解性和可扩展性,是PyTorch开发者提升模型性能的有力工具。

Thunder Thunder

让PyTorch模型快如闪电。


Lightning.ai性能开始使用安装示例Thunder内部参与其中文档

license CI testing General checks Documentation Status pre-commit.ci status

欢迎使用 ⚡ Lightning Thunder

Thunder让PyTorch模型快如闪电。

Thunder是PyTorch的源到源编译器。它通过结合并同时使用不同的硬件执行器(例如,nvFusertorch.compilecuDNNTransformerEngine FP8)来加速PyTorch程序。

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

 

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

 

单GPU性能

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

Thunder

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

 

多GPU性能

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

Thunder

 

开始使用

开始使用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、融合、分布式(如 ddpfsdp)、函数式(如 vmapvjpjvp
  • 一种将操作分派到可扩展的执行器集合的方法

Thunder 完全用 Python 编写。即使在所有转换阶段,其跟踪也表示为有效的 Python。这允许前所未有的内省和可扩展性水平。

Thunder 不直接为加速器(如 GPU)生成代码。它获取并转换用户程序,以便可以使用快速执行器最优地选择或生成设备代码,如:

用 Thunder 编译的模块和函数与原生 PyTorch 完全互操作,并支持 PyTorch 的自动微分。此外,Thunder 可以与 torch.compile 一起工作,以利用其最先进的优化。

 

文档

在线文档 可用。要在本地构建文档,您可以使用

make docs

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

 

参与其中!

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

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

 

许可证

Lightning Thunder 根据 Apache 2.0 许可发布。 详情请参阅 LICENSE 文件。

项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手,通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能,支持100+编程语言,与主流编辑器无缝集成,显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台,能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐,Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

白日梦AI

白日梦AI提供专注于AI视频生成的多样化功能,包括文生视频、动态画面和形象生成等,帮助用户快速上手,创造专业级内容。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案,通过智能技术简化视频制作流程。无论是企业宣传还是个人分享,有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持,能够阅读和理解用户上传的文件内容,解析网页信息,并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题,Kimi都能以友好、专业的方式提供帮助。

Project Cover

讯飞绘镜

讯飞绘镜是一个支持从创意到完整视频创作的智能平台,用户可以快速生成视频素材并创作独特的音乐视频和故事。平台提供多样化的主题和精选作品,帮助用户探索创意灵感。

Project Cover

讯飞文书

讯飞文书依托讯飞星火大模型,为文书写作者提供从素材筹备到稿件撰写及审稿的全程支持。通过录音智记和以稿写稿等功能,满足事务性工作的高频需求,帮助撰稿人节省精力,提高效率,优化工作与生活。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术,为商家提供一键生成商品图和营销文案的服务,显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台,让商品第一时间被种草。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具,简化了选题、文献检索至论文撰写的整个过程。通过简单设定,平台可快速生成高质量论文大纲和全文,配合图表、参考文献等一应俱全,同时提供开题报告和答辩PPT等增值服务,保障数据安全,有效提升写作效率和论文质量。

投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号