PyTorch 2.5 发布概览

PyTorch 基金会近日发布了 PyTorch 2.5，该版本引入了对 Intel GPU 的支持，并包含多项性能优化。此次更新共包含自 PyTorch 2.4 以来的 4095 次提交。

Intel GPU 支持

Intel GPU 的支​​持在最近的 PyTorch 会议 上首次预览。Intel 工程师 Eikan Wang 和 Min Jean Cho 详细介绍了 PyTorch 的改动，包括通用化 PyTorch 运行时和设备层，以简化新硬件后端的集成。此外，还为 torch.compile 和 torch.distributed 实现了 Intel 特定的后端。据 Intel 的 AI 框架工程副总裁 Kismat Singh 表示：

我们在 PyTorch 2.5 中添加了对 Intel 客户端 GPU 的支​​持，这意味着用户可以在基于最新 Intel 处理器的笔记本电脑和台式机上运行 PyTorch。预计今年将为 PyTorch 用户解锁 4000 万台设备，明年年底这一数字将达到 1 亿台。

FlexAttention API

新发布的 FlexAttention API 使 PyTorch 用户能够更轻松地在模型中尝试不同的注意力机制。以往，研究人员需要手动编写 PyTorch 运算符来实现新的注意力变体，这可能导致运行缓慢和 CUDA 内存不足问题。新的 API 允许用户用“几行简洁的 PyTorch 代码”实现这些机制，编译器会将其转换为优化的内核，性能接近手动编写的代码。

性能优化

多项性能优化功能以测试版发布：

• Fused Flash Attention：为 NVIDIA H100 GPU 提供比 FlashAttentionV2 高达 75% 的速度提升。
• 区域性编译：减少对整个模型的编译需求，仅编译重复的 nn.Modules（如 Transformer 层），从而减少编译延迟，性能损失仅几个百分点。
• TorchInductor CPU 后端优化：对 CPU 后端进行了多项优化。

Flight Recorder 调试工具

新引入的 Flight Recorder 工具用于调试分布式训练中的“卡住”任务。该工具使用内存循环缓冲区捕获诊断信息，当检测到任务卡住时，将诊断数据转储到文件中，用户可通过脚本分析数据以确定根本原因。

用户反馈

在 Reddit 上，许多用户对 Intel GPU 的支持表示欢迎，称其为“改变游戏规则的更新”。一位用户写道：

对 torch.compile 的改进感到兴奋，尤其是重复模块的复用功能，这对于具有大量相似组件的大型模型来说可能是革命性的。FlexAttention API 也极具潜力，仅用几行代码即可实现各种注意力机制并获得接近手动编写的性能，这非常了不起。感谢 PyTorch 团队和贡献者的又一次出色发布！

代码与发布说明

PyTorch 2.5 的 代码和发布说明 可在 GitHub 上获取。