跳转到主要内容
博客

PyTorch 2.4 发布博客

作者: 2024 年 7 月 24 日2024 年 11 月 12 日暂无评论

我们很高兴地宣布 PyTorch® 2.4 版本发布(发行说明)!PyTorch 2.4 为 torch.compile 增加了对最新版 Python(3.12)的支持。AOTInductor 冻结功能通过允许 MKLDNN 权重序列化,为运行 AOTInductor 的开发者提供了更多基于性能的优化。此外,还引入了利用 libuv 的新默认 TCPStore 服务器后端,这应能显著减少运行大规模作业的用户的初始化时间。最后,新的 Python 自定义运算符 API 使将自定义内核集成到 PyTorch 中比以前更容易,特别是对于 torch.compile

自 PyTorch 2.3 以来,此版本包含 3661 次提交和 475 位贡献者。我们衷心感谢我们敬业的社区所做的贡献。一如既往,我们鼓励您尝试这些功能并报告任何问题,以便我们改进 2.4。有关如何开始使用 PyTorch 2 系列的更多信息,请访问我们的入门页面。

测试版原型性能改进
torch.compile 支持 Python 3.12FSDP2:基于 DTensor 的逐参数分片 FSDPAWS Graviton (aarch64-linux) 处理器的 torch.compile 优化
AOTInductor CPU 冻结torch.distributed.pipelining,简化的流水线并行TorchInductor 中的 BF16 符号形状优化
新的高级 Python 自定义运算符 API可通过源代码构建获取 Intel GPU利用 CPU 设备的 GenAI 项目性能优化
将 TCPStore 的默认服务器后端切换到 libuv

*要查看完整的公开功能提交列表,请点击此处

Beta 功能

[Beta] torch.compile 支持 Python 3.12

torch.compile() 以前仅支持 Python 3.8-3.11。现在用户可以使用 Python 3.12 通过 torch.compile() 优化模型。

[Beta] AOTInductor CPU 冻结

此功能允许用户在 CPU 上使用 AOTInductor 时启用冻结标志。通过此功能,AOTInductor 可以覆盖相同的一组操作场景,并达到与 Inductor CPP 后端相媲美的性能。在此支持之前,当模型包含 MKLDNN 运算符(涉及计算密集型运算符,如卷积、线性、转置卷积等)且冻结功能开启时,这些模型将无法运行,因为 AOTInductor 不支持序列化不透明格式的 MKLDNN 权重。

工作流程如 AOTInductor 教程 中所述,此外,用户现在可以添加冻结标志以获得更好的性能。

export TORCHINDUCTOR_FREEZING=1

[Beta] 新的高级 Python 自定义运算符 API

我们添加了一个新的高级 Python 自定义运算符 API,使得通过自定义运算符扩展 PyTorch 比以前更容易,这些自定义运算符的行为类似于 PyTorch 的内置运算符。使用 新的高级 torch.library API 注册的运算符保证与 torch.compile 和其他 PyTorch 子系统兼容;使用以前的 低级 torch.library API 在 Python 中编写自定义运算符需要深入了解 PyTorch 内部结构,并且存在许多陷阱。

请参阅 教程 以获取更多信息。

[Beta] 将 TCPStore 的默认服务器后端切换到 libuv

引入了一个新的 TCPStore 默认服务器后端,该后端使用 libuv 构建,应能显著降低初始化时间和提高可伸缩性。这对于处理大规模作业的用户来说,理论上应能大大缩短启动时间。

有关动机 + 回退说明的更多信息,请参阅此 教程

原型功能

[原型] FSDP2:基于 DTensor 的逐参数分片 FSDP

FSDP2 是一种全新的完全分片数据并行实现,它使用 dim-0 逐参数分片来解决 FSDP1 的扁平参数分片带来的基本可组合性挑战。

有关 FSDP2 动机/设计的更多信息,请参阅 Github 上的 RFC

[原型] torch.distributed.pipelining,简化的流水线并行

流水线并行是深度学习中基本的并行技术之一。它允许对模型的执行进行分区,以便多个微批次可以同时执行模型的不同部分。

torch.distributed.pipelining 提供了一个工具包,可以轻松地在通用模型上实现流水线并行,同时还提供了与其他常见的 PyTorch 分布式功能(如 DDP、FSDP 或张量并行)的可组合性。

有关此内容的更多信息,请参阅我们的文档教程

[原型] Intel GPU 可通过源代码构建获得

Linux 系统上 PyTorch 中的 Intel GPU 提供了 Intel® Data Center GPU Max 系列的基本功能:eager 模式和 torch.compile。

对于 eager 模式,常用的 Aten 运算符通过 SYCL 编程语言实现。大多数性能关键的图和运算符通过 oneAPI 深度神经网络 (oneDNN) 进行高度优化。对于 torch.compile 模式,Intel GPU 后端已集成到基于 Triton 的 Inductor 中。

有关 Intel GPU 源代码构建的更多信息,请参阅我们的 博客文章文档

性能改进

AWS Graviton (aarch64-linux) 处理器的 torch.compile 优化

AWS 优化了 PyTorch 的 torch.compile 功能,以适用于 AWS Graviton3 处理器。与在基于 AWS Graviton3 的 Amazon EC2 实例上对多种自然语言处理 (NLP)、计算机视觉 (CV) 和推荐模型进行默认 eager 模式推理相比,此优化使 Hugging Face 模型推理性能提升高达 2 倍(基于 33 个模型性能改进的几何平均值),TorchBench 模型推理性能提升高达 1.35 倍(45 个模型性能改进的几何平均值)。

有关具体技术细节的更多信息,请参阅博客文章

TorchInductor 中的 BF16 符号形状优化

PyTorch 用户现在可以通过 Beta BF16 符号形状支持体验改进的质量和性能提升。虽然静态形状可能比符号形状提供额外的优化机会,但对于批处理大小和序列长度不同的推理服务,或具有数据依赖输出形状的检测模型等场景,它是不够的。

使用 TorchBench、Huggingface 和 timms_model 的验证显示,通过率与 BF16 静态形状场景相似,加速效果也相当。结合了符号形状的优势与 Intel CPU 提供的 BF16 AMX 指令硬件加速,以及适用于 PyTorch 2.4 中静态和符号形状的通用 Inductor CPU 后端优化,BF16 符号形状的性能相比 PyTorch 2.3 有了显著提升。

使用此功能的 API

model = ….
model.eval()
with torch.autocast(device_type=”cpu”, dtype=torch.bfloat16), torch.no_grad():
   compiled_model = torch.compile(model, dynamic=True)

利用 CPU 设备的 GenAI 项目性能优化

通过对 “Segment Anything Fast”“Diffusion Fast” 项目的优化,突出了 PyTorch 在 CPU 上的增强性能。然而,模型中仅支持 CUDA 设备。我们已将 CPU 支持集成到项目中,使用户能够利用 CPU 增强的性能来运行项目实验。同时,我们还为 SDPA 采用了 块式注意力掩码,这可以显著降低峰值内存使用量并提高性能。我们还在 Inductor CPU 中优化了一系列 布局传播规则,以提高性能。

为了方便起见,我们更新了 README 文件。使用此功能的 API 如下所示,只需在命令行中提供 --device cpu 即可。

  • 对于 Segment Anything Fast
export SEGMENT_ANYTHING_FAST_USE_FLASH_4=0
python run_experiments.py 16 vit_b <pytorch_github> <segment-anything_github>
<path_to_experiments_data> --run-experiments --num-workers 32 --device cpu
  • 对于 Diffusion Fast
python run_benchmark.py --compile_unet --compile_vae --enable_fused_projections --device=cpu

用户可以按照指南运行实验,亲身体验性能提升,并探索 FP32 和 BF16 数据类型之间的性能改进趋势。

此外,用户可以通过使用 torch.compile 和 SDPA 获得良好的性能。通过观察这些不同因素的性能趋势,用户可以更深入地了解各种优化如何增强 PyTorch 在 CPU 上的性能。