• 使用 torch.compile 作为 PyTorch 2.0 的主要 API；

• scaled_dot_product_attention 函数作为 torch.nn.functional 的一部分；

• MPS 后端；

• torch.func 模块中的 functorch API。

• torch.compile 是 PyTorch 2.0 的主要 API，它能包装并返回编译后的模型。这个是一个完全附加（和可选）的功能，PyTorch 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 上的 PyTorch CPU 推理。与之前的版本相比，PyTorch 2.0 提高了 Graviton 的推理性能，包括针对 ResNet-50 和 BERT 的改进。

• 其他一些跨 TensorParallel、DTensor、2D parallel、TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 的新 prototype 功能和方法。

稳定功能

PyTorch 2.0 版本包括 PyTorch Transformer API 新的高性能实现，以前称为「Better Transformer API」，现在更名为 「Accelerated PyTorch 2 Transformers」。研发团队表示他们希望整个行业都能负担得起训练和部署 SOTA Transformer 模型的成本。新版本引入了对训练和推理的高性能支持，使用自定义内核架构实现缩放点积注意力 (SPDA)。

• 显著提升模型速度；

• 支持更多用例，包括使用交叉注意力模型、Transformer 解码器，并且可以用于训练模型；

• 继续对固定和可变的序列长度 Transformer 编码器和自注意力用例使用 fastpath 推理。

将缩放点积注意力与自定义内核和 torch.compile 结合使用可为训练大型语言模型（上图以 nanoGPT 为例）提供显著加速。

torch.compile

torch.compile 是 PyTorch 2.0 的主要 API，它包装并返回编译后的模型。torch.compile 的背后是 PyTorch 团队研发的新技术 ——TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

• [Prototype] DTensor

• [Prototype] TensorParallel

• [Prototype] 2D Parallel

• [Prototype] torch.compile (dynamic=True)