与「快速路径（fastpath）」架构类似，自定义内核完全集成到 PyTorch Transformer API 中 —— 因此，使用 Transformer 和 MultiHeadAttention API 将使用户能够：

为了充分利用不同的硬件模型和 Transformer 用例，PyTorch 2.0 支持多个 SDPA 自定义内核，自定义内核选择逻辑是为给定模型和硬件类型选择最高性能的内核。除了现有的 Transformer API 之外，模型开发人员还可以通过调用新的 scaled_dot_product_attention 运算来直接使用缩放点积注意力内核。

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

Beta 功能

torch.compile

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

借助这些新技术，torch.compile 能够在 165 个开源模型上运行，并且在 float32 精度下平均运行速度提高 20%，在 AMP 精度下平均运行速度提高 36%。

PyTorch MPS 后端

MPS 后端在 Mac 平台上提供 GPU 加速的 PyTorch 训练。PyTorch 2.0 在正确性、稳定性和运算符覆盖率方面比之前的版本有所改进。

缩放点积注意力 2.0

PyTorch 2.0 引入了一个强大的缩放点积注意力函数。该函数包括多种实现，可以根据使用的输入和硬件无缝应用。

functorch → torch.func

functorch API 现在可以在 torch.func 模块中使用。其中，函数转换 API 与以前相同，但与 NN 模块交互的方式有所改变。

此外，PyTorch 2.0 还添加了对 torch.autograd.Function 的支持：现在可以在 torch.autograd.Function 上应用函数转换。

Dispatchable Collectives

Dispatchable Collectives 是对之前 init_process_group API 的改进，其中将后端更改为可选参数。对于用户来说，这个特性的主要优势在于，它将允许用户编写可以在 GPU 和 CPU 机器上运行的代码，而无需更改后端规范。

PyTorch 2.0 还将 torch.set_default_device 和 torch.device 作为语境管理器（context manager），将「X86」作为 x86 CPU 的新默认量化后端。

新的 X86 量化后端利用 FBGEMM 和 oneDNN 内核库，提供比原始 FBGEMM 后端更高的 INT8 推理性能。新后端在功能上与原始 FBGEMM 后端兼容。

此外，PyTorch 2.0 还包括多项关键优化，以提高 CPU 上 GNN 推理和训练的性能，并利用 oneDNN Graph 加速推理。

最后，PyTorch 2.0 还包含一些 Prototype 功能，包括：

参考链接：https://deploy-preview-1313--pytorch-dot-org-preview.netlify.app/blog/pytorch-2.0-release/

返回搜狐，查看更多