12月2日，PyTorch2.0正式发布！本次更新不仅将PyTorch的性能推向了新的高度，还增加了对动态形状和分布的支持。此外，2.0系列还将PyTorch的部分代码从C++移回Python。目前，PyTorch2.0仍处于测试阶段，第一个稳定版本预计将于2023年3月上旬面世。PyTorch2.x：更快，更Pythonic！在过去的几年里，PyTorch从1.0到最近的1.13进行了创新和迭代，并作为Linux基金会的一部分搬到了新成立的PyTorch基金会。当前版本的PyTorch面临的挑战是急切模式难以跟上不断增加的GPU带宽和更疯狂的模型架构。PyTorch2.0的诞生将从根本上改变和改进PyTorch在编译器层面的运行方式。众所周知，PyTorch中的(Py)来自数据科学中广泛使用的开源Python编程语言。不过PyTorch的代码并没有完全使用Python，而是有一部分交给了C++。不过在未来的2.x系列中，PyTorch项目组计划将与torch.nn相关的代码移回Python。除此之外，由于PyTorch2.0是一个完全附加（和可选）的功能，因此2.0是100%向后兼容的。也就是说，代码库相同，API相同，您编写模型的方式也相同。更多技术支持TorchDynamo使用Python框架评估钩子安全捕获PyTorch程序，这是团队过去5年在图捕获方面的重大创新。AOTAutograd重载PyTorch的autograd引擎作为跟踪的autodiff，用于生成高级向后跟踪。PrimTorch将大约2000多个PyTorch运算符概括为一个包含大约250个原始运算符的封闭集合，开发人员可以根据这些运算符构建完整的PyTorch后端。大大降低了编写PyTorch函数或后端的门槛。TorchInductor一种深度学习编译器，可为多个加速器和后端生成快速代码。对于Nvidia的GPU，它使用OpenAITriton作为关键构建块。值得注意的是，TorchDynamo、AOTAutograd、PrimTorch和TorchInductor都是用Python编写的，并且支持动态形状。更快的训练速度通过引入一种新的编译模式“torch.compile”，PyTorch2.0可以用一行代码加速模型训练。这里没有技巧，运行torch.compile()即可，仅此而已：opt_module=torch.compile(module)为了验证这些技术，团队精心搭建了测试基准，包括图像分类、目标检测、图像生成等。任务，以及各种NLP任务，如语言建模、问答、序列分类、推荐系统和强化学习。其中，这些基准测试可以分为三类：HuggingFaceTransformers的46个模型TIMM的61个模型：RossWightman收集的最先进的PyTorch图像模型TorchBench的56个模型：github上的一组流行代码库测试结果表明，在这163个跨越视觉、NLP等领域的开源模型上，训练速度提升了38%-76%。在NVIDIAA100GPU上的比较此外，该团队还对一些流行的开源PyTorch模型进行了基准测试，并获得了30%到2倍的大幅加速。开发者SylvainGugger表示：“只需添加一行代码，PyTorch2.0就可以在训练Transformers模型时实现1.5倍到2.0倍的加速。这是自混合精度训练出现以来最令人兴奋的事情！”技术概述PyTorch编译器可以分为三个部分：图获取图缩减图编译其中，图获取是构建PyTorch编译器时更困难的挑战。TorchDynamo今年早些时候，该团队开始研究TorchDynamo，这是一种使用PEP-0523中引入的称为框架评估API的CPython功能的方法。为此，该团队采用数据驱动的方法，通过使用PyTorch编写的7,000多个Github项目作为验证集来验证TorchDynamo在图形捕获方面的有效性。结果表明，TorchDynamo可以在99%的时间内正确安全地执行图形捕获，而开销可以忽略不计。TorchInductor对于PyTorch2.0的新编译器后端，该团队从用户编写高性能自定义内核的方式中获得灵感：越来越多地使用Triton语言。TorchInductor使用Pythonic定义的逐周期级IR自动将PyTorch模型映射到GPU上生成的Triton代码和CPU上的C++/OpenMP。TorchInductor的核心循环级IR仅包含约50个运算符，并且使用Python实现，易于扩展。如果AOTAutograd想要加快训练速度，它不仅需要捕获用户级代码，还需要捕获反向传播。AOTAutograd可以使用PyTorch的torch_dispatch扩展机制来跟踪Autograd引擎，“提前”捕获反向传播，然后能够使用TorchInductor加速前向和后向通道。PrimTorchPyTorch有1200多个运算符，如果考虑每个运算符的各种重载，则超过2000个。因此，编写后端或横切功能成为一项劳动密集型工作。在PrimTorch项目中，团队定义了两个更小且更稳定的运算符集：Primops有大约250个运算符，适用于编译器。由于它足够低级，因此只需将它们融合在一起以获得良好的性能。ATenops有大约750个典型的操作符，适合按原样输出。这些适用于已在ATen级别集成的后端，或尚未编译的后端，可恢复低级运算符集（如Primops）的性能。动态形状在研究支持PyTorch代码的通用性的必要条件时，一个关键要求是支持动态形状并允许模型接受不同大小的张量，而不会在每次形状更改时都导致重新编译。在不支持动态形状的情况下，一种常见的解决方法是将它们填充到最接近的2的幂。但是，正如我们从下图中可以看到的那样，它会产生显着的性能开销，并显着延长编译时间。现在，由于支持动态形状，PyTorch2.0的性能比Eager高出40%。最后，在PyTorch2.x路线图中，团队希望在性能和可扩展性方面进一步推动编译模型。
