据人工智能研究公司 OpenAI 称,Nvidia 的图形处理单元太难编程,包括使用 Nvidia 自己的编程工具 CUDA。这家总部位于旧金山的人工智能初创公司在微软和风险投资公司 Khosla Ventures 的支持下,于周三推出了 1.0 版本,这是一种专门为减轻这一负担而设计的新编程语言 Triton,在一篇链接到 GitHub 源代码的博客文章中进行了详细介绍。 OpenAI 声称与 CUDA 中的编码相比,Triton 可以为一些处于 AI 机器学习形式核心的神经网络任务(例如矩阵乘法)提供显着的易用性优势。 “我们的目标是使其成为深度学习 CUDA 的可行替代方案,”该项目的负责人、OpenAI 科学家 Philippe Tillet 通过电子邮件告诉 ZDNet。 Tillet 表示,Triton “适用于尽管拥有良好的软件工程技能但不熟悉 GPU 编程的机器学习研究人员和工程师。”该语言来自 OpenAI,它开发了风靡全球的 GPT-3 自然语言处理程序,这一事实可能会使该代码在 AI 领域中获得一些额外的优势。该软件是作为开源软件提供的,但要求版权声明和许可包含在代码的大量副本的任何分发中。
最初的 Triton 揭幕发生在 Tillet 于 2019 年发表的一篇论文中,当时他还是哈佛大学的一名研究生,以及他的顾问 HT Kung 和 David Cox。 Tillet 着手解决的问题是如何制作一种比供应商特定的 AI 库更具表现力的语言,例如 Nvidia 的 cuDNN,意思是能够处理神经网络中涉及的矩阵的各种操作;同时具有可移植性和可与 cuDNN 和类似供应商库相媲美的性能。根据 Tillet 和团队的说法,直接在 CUDA 中编程 GPU 太难了。例如,为 GPU 编写原生内核或函数“由于 GPU 编程的许多复杂性而出奇地困难,”Tillet 和团队在博文中写道。特别是,正如 Triton 文档所解释的那样,“GPU 在优化局部性和并行性方面仍然极具挑战性”。但 Tillet 还希望该语言比迄今为止的定制工作更容易编程,基于所谓的“微内核”,“涉及大量手动工作”。特别是,Triton 被用作替代供应商库的两种主要方法的替代方案,这两种方法被称为多面体编译和调度语言。 Tillet 确定的是一种称为瓷砖的方法。在 CUDA 编程中广泛使用的 Tiles 将机器学习程序中使用的矩阵分解成片段,这些片段可以有效地分布在共享 SRAM 内存和快速寄存器内存中,并通过多个并行指令线程有效地进行操作。但是,在 CUDA 中进行并行化很困难,因为需要在程序的指令线程之间进行显式同步语句。
Triton 的语义将切片指定为内置类型,以便 Triton 编译器可以确定如何在 GPU 的多个内核及其随附寄存器之间有效分配这些片段。实际上,并行化和优化代码的工作从语言下推到编译器中。正如 Tillet 所说,编译器“自动执行各种重要的程序优化”。 “例如,通过查看计算密集型块级操作的操作数,数据可以自动存储到共享内存中。” Triton 程序员的高级代码首先被转换为中间表示,其灵感来自开源 LLVM 编译器基础结构中的中间表示。正如 Tillet 在原始论文中所描述的那样,“只需对 LLVM-IR 进行一些数据和控制流扩展,就可以实现各种平铺级优化传递,从而使性能与供应商库相提并论。”然后将中间表示馈送到即时编译器,该编译器执行将各种矩阵塑造成片段的工作,这种方式将最适合共享内存和 GPU 内核的寄存器。 JIT 在 GPU 内核中组织指令线程以从主内存中提取相同的值,称为“内存合并”。同样,JIT 将此类线程共同感兴趣的数据放入共享内存中以进行高效操作,称为“共享内存分配”。
正如 Tillet 所描述的那样,结果是“单线程和自动并行化”的程序。 JIT 正在执行自动调整图块、数据片段的工作,以最有效地将它们分布在核心之间。在最初的 Triton 论文中,Tillet 提出了一种基于 CUDA 语法的类似 C 形式的 Triton。然而,在这个新的 1.0 版本中,Triton 与 Python 集成。详细信息在博客文章中详细说明。使用 Triton 的好处应该是立即加快开发神经网络的一些基本操作。正如 Tillet 在博客文章中阐述的那样,“它可以用来编写与 cuBLAS 性能相匹配的 FP16 矩阵乘法内核”,这是一个实现开源基本线性代数子程序的 Nvidia 库,“这是许多 GPU 程序员无法做到的”不做——不到 25 行代码。”他说,Tillet 在 OpenAI 的超级计算负责人 Chris Berner 的监督下,在 OpenAI 全职从事该项目。但也有几位 OpenAI 员工对 Triton 项目的帮助。 “最近,几位 OpenAI 员工和研究人员——都没有 GPU 编程经验——为该项目贡献了代码和想法,”Tillet 告诉 ZDNet。 “我们已经使用它来加速和重写我们的大部分 GPU 内核,并且我们致力于通过后续版本使其更广泛地适用。” Tillet 指出,该项目得到了 OpenAI 外部的“有意义的贡献”,包括香港科技大学的大燕、微软 DeepSpeed 优化库的团队以及商业 AI 初创公司 Anthropic。周三的博客文章并没有强调性能指标,只是说 Triton 可以匹配 CuBLAS。然而,在 Tillet 的原始论文中,该语言的 Triton-C 版本在运行所谓的深度卷积时能够获得比 Nvidia 的 CuDNN 库更好的性能,这些操作将输入视为本地相关数据的组,例如图像像素。
请注意,目前该软件仅适用于 Nvidia GPU;它还不能用于 AMD 的 GPU,也不能编译到 CPU。作者邀请对这些芯片感兴趣的合作者加入这项工作。 Tillet 的语言努力正值 AI 硬件加速领域的一个有趣时刻。英伟达面临着来自人工智能芯片和系统初创公司的激烈竞争,例如 Cerebras Systems、Graphcore 和 SambaNova。这些公司都有各种芯片架构,可以将并行计算分配给多个片上内核。事实上,SambaNova 的芯片具有所谓的数据流架构,该架构共享 Triton 的一些原理。然而,所有这些供应商都不得不开发自己的软件工具来优化 PyTorch 和 TensorFlow 程序在他们的计算机上的移动。相比之下,英伟达拥有超过十年的 CUDA 开发经验和广泛的软件开发人员基础。可以想象,Triton 可能成为竞争者为其芯片获得广泛的开源平台所需的新软件工具之一。