Triton:重新定义深度学习原语的编程语言和编译器 🌊💻

267次阅读
一条评论

大家好,欢迎来到我们的 AI 播客。我是你的主持人,一位资深的 AI 专家。今天,我们将聚焦于一个名为 Triton 的开源项目,它正在重新定义我们编写高效深度学习原语的方式。🚀

Triton 的魅力 🌟

Triton 是 OpenAI 的开发项目,是一种用于编写极其高效的自定义深度学习原语的语言和编译器。📚 这个项目的宗旨在于提供一个开源环境,让开发者可以以比 CUDA 更高的效率编写快速代码,同时比其他现有的 DSL(领域特定语言)具有更高的灵活性。🔗

基于这个项目的基础理论已经在 MAPL2019 的出版物 ”Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations” 中描述,如果你使用 Triton,一定要引用这篇文章。📝

快速安装 Triton 🚀

Triton 的安装非常直接。你只需要使用 PIP(Python 的包管理器),通过下面的命令就可以安装最新的稳定版本:

pip install triton

如果你想要安装最新的每夜版本,可以使用以下命令:

pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly

从源码安装 Triton 🧩

当然,你也可以选择从源代码安装 Triton。首先,你需要复制代码库,然后在其根目录下安装相关的依赖项,最后安装 Triton 的 Python 包。以下是详细的步骤:

git clone https://github.com/openai/triton.git;
cd triton;

pip install ninja cmake wheel; # build-time dependencies
pip install -e python

Triton 和自定义 LLVM 🔧

值得一提的是,Triton 使用 LLVM(一种用于开发编译器的库)来为 GPU 和 CPU 生成代码。通常情况下,Triton 会下载一个预构建的 LLVM,但你也可以自己从源代码构建 LLVM。

然而,需要注意的是,LLVM 并没有稳定的 API,因此在任意的 LLVM 版本上,Triton 的构建都可能无法正常工作。

关于 Triton 更深入的内容,我想邀请你们访问其 Github 页面 自行探索。在这个存储库中,你可以找到详细的文档,以及一些关于如何使用 Triton 的教程。📖

结语 🌟

Triton 是一个令人兴奋的项目,它正在改变我们编写深度学习原语的方式。它提供了一种新的,更高效和灵活的方法,让我们能够更好地探索神经网络的潜力。🔥

希望你们喜欢今天的播客,咱们下次再见!👋

正文完
 
评论(一条评论)
2023-12-06 11:21:11 回复

Triton是一个用于编写高效自定义深度学习原语的开源语言和编译器项目。它由OpenAI发起,目的是提供一个开发环境,可以以比CUDA更高的生产力来编写性能接近CUDA的代码,同时又比其他现有的领域特定语言更具灵活性。

Triton的编程接口最初是基于C++的,后来引入了Python接口,大大提高了其易用性。它采用了“Block-wise编程”的方式,将Block内部的优化逻辑交给Triton编译器自动完成,使开发者可以更关注业务逻辑。

Triton编译器的优化核心是Layout抽象和各种Passes。Layout抽象描述了计算资源与输入输出元素的映射关系。各种Passes则实现了一些NV GPU计算的常见优化技巧,如pipeline、swizzling等。这些都降低了开发者对底层优化的关注。

Triton后端经历了从手写IR到基于MLIR重构的过程。MLIR重构后架构更清晰,也便于支持多种硬件后端。目前Triton在MLIR使用上还比较务实,没有过多追求语法抽象优雅。这在项目早期是可以接受的。

Triton目前应用最多的场景是被集成到PyTorch的Inductor中,为其提供高效的自定义算子实现。JAX和XLA也有集成Triton的工作。这些集成有不同的实现策略。

Triton找到了一个适合自己的“产品-技术”契合点,解决了高性能定制算子开发的问题,是这个领域一个值得关注的开源实现。其接口易用性和底层优化自动化降低了开发门槛。基于MLIR的实现也为其带来了更好的前景。希望Triton能成为AI系统中一个易用且高效的组件,也期待该领域的更多创新。

 Macintosh  Edge  澳大利亚新南威尔士悉尼