PyTorch 2.0正式版发布！一行代码提速2倍，100%向后兼容

新智元报道

编辑：Britta 桃子

【新智元导读】PyTorch 2.0正式发布。

PyTorch 2.0正式版终于来了！

去年12月，PyTorch基金会在PyTorch Conference 2022上发布了PyTorch 2.0的第一个预览版本。

跟先前1.0版本相比，2.0有了颠覆式的变化。在PyTorch 2.0中，最大的改进是torch.compile。

新的编译器比以前PyTorch 1.0中默认的「eager mode」所提供的即时生成代码的速度快得多，让PyTorch性能进一步提升。

除了2.0之外，还发布了一系列PyTorch域库的beta更新，包括那些在树中的库，以及包括 TorchAudio、TorchVision和TorchText在内的独立库。TorchX的更新也同时发布，可以提供社区支持模式。

亮点总结

-torch.compile是PyTorch 2.0的主要API，它包装并返回编译后的模型，torch.compile是一个完全附加（和可选）的特性，因此2.0版本是100%向后兼容的。

-作为torch.compile的基础技术，带有Nvidia和AMD GPU的TorchInductor将依赖OpenAI Triton深度学习编译器来生成高性能代码，并隐藏低级硬件细节。OpenAI Triton生成的内核实现的性能，与手写内核和cublas等专门的cuda库相当。

-Accelerated Transformers引入了对训练和推理的高性能支持，使用自定义内核架构实现缩放点积注意力 (SPDA)。API与torch.compile () 集成，模型开发人员也可以通过调用新的scaled_dot_product_attention () 运算符，直接使用缩放的点积注意力内核。

-Metal Performance Shaders (MPS) 后端在Mac平台上提供GPU加速的PyTorch训练，并增加了对前60个最常用操作的支持，覆盖了300多个操作符。

-Amazon AWS优化了基于AWS Graviton3的C7g实例上的PyTorch CPU推理。与之前的版本相比，PyTorch 2.0提高了Graviton的推理性能，包括对Resnet50和Bert的改进。

-跨TensorParallel、DTensor、2D parallel、TorchDynamo、AOTAutograd、PrimTorch和TorchInductor的新原型功能和技术。

编译，还是编译！

PyTorch 2.0的最新编译器技术包括：TorchDynamo、AOTAutograd、PrimTorch和TorchInductor。所有这些都是用Python开发的，而不是C++（Python与之兼容）。

并且还支持dynamic shape，无需重新编译就能发送不同大小的向量，灵活且易学。

1. 
   

   TorchDynamo

   
   

它可以借助Python Frame Evaluation Hooks，安全地获取PyTorch程序，这项重大创新是PyTorch过去 5 年来在安全图结构捕获 (safe graph capture) 方面的研发成果汇总。

1. 
   

   AOTAutograd

   
   

重载PyTorch autograd engine，作为一个 tracing autodiff，用于生成超前的backward trace。

1. 
   

   PrimTorch

   
   

将 2000+ PyTorch 算子归纳为约 250 个 primitive operator 闭集 (closed set)，开发者可以针对这些算子构建一个完整的 PyTorch 后端。PrimTorch 大大简化了编写 PyTorch 功能或后端的流程。

4. TorchInductor

TorchInductor一个深度学习编译器，可以为多个加速器和后端生成 fast code。 对于 NVIDIA GPU，它使用 OpenAI Triton 作为关键构建模块。

PyTorch基金会称，2.0的推出会推动「从C++回到Python」，并补充说这是PyTorch的一个实质性的新方向。

「从第一天起，我们就知道「eager execution」的性能限制。2017年7月，我们开始了第一个研究项目，为PyTorch开发一个编译器。编译器需要使PyTorch程序快速运行，但不能以PyTorch的体验为代价，还要保留灵活易用性，这样的话可以支持研究人员在不同探索阶段使用动态的模型和程序。」

当然了，非编译的「eager mode」使用动态即时代码生成器，在2.0中仍然可用。开发者可以使用porch.compile命令迅速升级到编译模式，只需要增加一行代码。

用户可以看到2.0的编译时间比1.0提高43%。

这个数据来自PyTorch基金会在Nvidia A100 GPU上使用PyTorch 2.0对163个开源模型进行的基准测试，其中包括包括图像分类、目标检测、图像生成等任务，以及各种 NLP 任务。

这些Benchmark分为三类：HuggingFace Tranformers、TIMM和TorchBench。

NVIDIA A100 GPU eager mode torch.compile 针对不同模型的提速表现

据PyTorch基金会称，新编译器在使用Float32精度模式时运行速度提高了21%，在使用自动混合精度（AMP）模式时运行速度提高了51%。

在这163个模型中，torch.compile可以在93%模型上正常运行。

「在PyTorch 2.x的路线图中，我们希望在性能和可扩展性方面让编译模式越走越远。有一些工作还没有开始。有些工作因为带宽不够而办法落地。」

训练LLM提速2倍

此外，性能是PyTorch 2.0的另一个主要重点，也是开发人员一直不吝于宣传的一个重点。

事实上，新功能的亮点之一是Accelerated Transformers，之前被称为Better Transformers。

另外，PyTorch 2.0正式版包含了一个新的高性能PyTorch TransformAPI实现。

PyTorch项目的一个目标，是让最先进的transformer模型的训练和部署更加容易、快速。

Transformers是帮助实现现代生成式人工智能时代的基础技术，包括GPT-3以及GPT-4这样的OpenAI模型。

在PyTorch 2.0 Accelerated Transformers中，使用了自定义内核架构的方法（也被称为缩放点积注意力SDPA），为训练和推理提供高性能的支持。

由于有多种类型的硬件可以支持Transformers，PyTorch 2.0可以支持多个SDPA定制内核。更进一步，PyTorch集成了自定义内核选择逻辑，将为给定的模型和硬件类型挑选最高性能的内核。

加速的影响非同小可，因为它有助于使开发人员比以前的PyTorch迭代更快地训练模型。

新版本能够实现对训练和推理的高性能支持，使用定制的内核架构来处理缩放点积注意力(SPDA) ，扩展了推理的快速路径架构。

与fastpath架构类似，定制内核完全集成到PyTorch TransformerAPI中--因此，使用本地Transformer和MultiHeadAttention API将使用户能够：

-看到速度明显提升；

-支持更多的用例，包括使用交叉注意的模型、Transformer解码器和训练模型；

-继续将快速路径推理用于固定和可变序列长度的变形器编码器和自注意力机制的用例。

为了充分利用不同的硬件模型和Transformer用例，支持多个SDPA自定义内核，自定义内核选择逻辑将为特定模型和硬件类型挑选最高性能的内核。

除了现有的Transformer API，开发者还可以通过调用新的scaled_dot_product_attention()操作符直接使用缩放点积注意力关注内核，加速PyTorch 2 Transformers与torch.compile()集成。

为了在使用模型的同时，还能获得PT2编译的额外加速（用于推理或训练），可以使用model = torch.compile(model)对模型进行预处理。

目前，已经使用自定义内核和torch.compile()的组合，在训练Transformer模型，特别是使用加速的PyTorch 2 Transformer的大型语言模型方面取得实质性加速提升。

使用自定义内核和 torch.compile来提供对大型语言模型训练显著加速

HuggingFace Transformers的主要维护者Sylvain Gugger在PyTorch项目发表的一份声明中写道「只需添加一行代码，PyTorch 2.0就能在训练Transformers模型时提供1.5倍至2.0倍的速度。这是自混合精度训练推出以来最令人兴奋的事情！」

PyTorch和谷歌的TensorFlow是两个最流行的深度学习框架。世界上有数千家机构正在使用PyTorch开发深度学习应用程序，而且它的使用量正在不断增加。

PyTorch 2.0的推出将有助于加速深度学习和人工智能应用的发展，Lightning AI的首席技术官和PyTorch Lightning的主要维护者之一Luca Antiga表示:

「PyTorch 2.0 体现了深度学习框架的未来。不需要用户干预即可捕获PyTorch 程序，开箱即用的程序生成，以及巨大的设备加速，这种可能性为人工智能开发人员打开了一个全新的维度。」

参考资料：

https://pytorch.org/blog/pytorch-2.0-release/

https://venturebeat.com/ai/pytorch-2-0-brings-new-fire-to-open-source-machine-learning/

https://www.datanami.com/2023/03/15/new-pytorch-2-0-compiler-promises-big-speedup-for-ai-developers/