NVIDIA|NVIDIA独家绝技 CUDA正在被赶下神坛？( 五 ) 显卡|CUDA|NVIDIA独家绝技|CUDA正在

他们终于用 TorchDynamo 破解了这个难题。TorchDynamo 将摄取任何 PyTorch 用户脚本，包括调用外部 3rd 方库的脚本，并生成FX 图形。
Dynamo 将所有复杂算子减少到 PrimTorch 中的约 250 个原始算子。
一旦图形成，未使用的算子将被丢弃，图确定哪些中间算子需要存储或写入内存，哪些可能被融合。这极大地减少了模型内的开销，同时对用户来说也是无缝的。
TorchDynamo 已经适用于7,000 个经过测试的 PyTorch 模型中的 99% 以上，包括来自 OpenAI、HuggingFace、Meta、Nvidia、Stability.AI 等的模型，而无需对原始代码进行任何更改。
测试的 7,000 个模型是从 GitHub 上使用 PyTorch 的最流行项目中不分青红皂白地挑选出来的。

文章图片

谷歌的 TensorFlow/Jax 和其他图形模式执行管道通常要求用户确保他们的模型适合编译器架构，以便可以捕获图形。Dynamo 通过启用部分图捕获、受保护的图捕获和即时重新捕获来改变这一点。
部分图形捕获允许模型包含不受支持的/非 python 构造。当无法为模型的那部分生成图时，将插入图中断，并且将在部分图之间以急切模式执行不支持的构造。
受保护的图捕获检查捕获的图是否对执行有效。守卫是需要重新编译的更改。这很重要，因为多次运行相同的代码不会多次重新编译。
【NVIDIA|NVIDIA独家绝技 CUDA正在被赶下神坛？】如果捕获的图对于执行无效，则即时重新捕获允许重新捕获图。

文章图片

PyTorch 的目标是创建一个具有流畅 UX 的统一前端，该前端利用 Dynamo 生成图形。该解决方案的用户体验不会发生变化，但性能可以得到显着提升。捕获图形意味着可以在大量计算资源上更有效地并行执行。
Dynamo 和AOT Autograd然后将优化的 FX 图传递给 PyTorch 本机编译器级别 TorchInductor 。硬件公司也可以将此图输入到他们自己的后端编译器中。
TorchInductor
TorchInductor 是一个 python 原生深度学习编译器，可以为多个加速器和后端生成快速代码。Inductor 将采用具有约 250 个算子的 FX 图，并将它们降低到约 50 个算子。
Inductor 然后进入调度阶段，在该阶段融合运算符，并确定内存规划。
Inductor 然后进入“Wrapper Codegen”，它生成在 CPU、GPU 或其他 AI 加速器上运行的代码。包装器 codegen 取代了编译器堆栈的解释器部分，可以调用内核和分配内存。
后端代码生成部分利用适用于 GPU 的 OpenAI Triton 并输出 PTX 代码。对于 CPU，英特尔编译器生成 C++（也适用于非英特尔 CPU）。
未来他们将支持更多硬件，但关键是 Inductor 大大减少了编译器团队在为其 AI 硬件加速器制作编译器时必须做的工作量。此外，代码针对性能进行了更优化。显着降低了内存带宽和容量要求。
分析人士表示，我们不想构建只支持 GPU 的编译器。我们想要一些可以扩展以支持各种硬件后端的东西，并且拥有 C++ 和 [OpenAI] Triton 会强制实现这种通用性。
OpenAI 的Triton
OpenAI 的 Triton 对 Nvidia 的机器学习闭源软件护城河具有颠覆性的角度。Triton 直接采用 Python 或通过PyTorch Inductor 堆栈提供数据。
后者将是最常见的用例。Triton 然后将输入转换为 LLVM 中间表示，然后生成代码。对于 Nvidia GPU，它直接生成 PTX 代码，跳过 Nvidia 的闭源 CUDA 库（例如 cuBLAS），转而使用开源库（例如 cutlass）。