PyTorch 2.0 现已发布!
新版本包括稳定版的 Accelerated Transformers(以前称为 Better Transformers);Beta 功能包括:
- 使用 torch.compile 作为 PyTorch 2.0 的主要 API(有关 torch.compile 的全面介绍和技术概述,请访问 2.0 入门页面。)
- scaled_dot_product_attention 函数作为 torch.nn.functional 的一部分
- MPS 后端
- torch.func 模块中的 functorch API
- 以及 GPU 和 CPU 上各种推理、性能和训练优化功能的其他 Beta/Prototype 改进。
新功能的概括介绍:
- 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 ,PyTorch 官方还发布了 PyTorch 域库的一系列 beta 更新,包括那些在树中的库,以及包括 TorchAudio、TorchVision 和 TorchText 在内的独立库。TorchX 转向社区支持模式,正在发布更新。可以在该博客中找到更多详细信息。
PyTorch 2.0 官方发布公告中对 2.0 的每个新特性作了详细介绍。