跳转到主要内容
博客

PyTorch 2.3 发布博客

作者: 2024 年 4 月 24 日2024 年 11 月 13 日无评论

我们很高兴地宣布 PyTorch® 2.3 已发布(发布说明)!PyTorch 2.3 支持 torch.compile 中的用户定义 Triton 内核,允许用户将其自己的 Triton 内核从 Eager 模式迁移,而不会遇到性能下降或图中断。张量并行性改进了使用原生 PyTorch 函数训练大型语言模型的体验,这已通过 1000 亿参数模型的训练运行得到验证。此外,半结构化稀疏性将半结构化稀疏性实现为张量子类,与密集矩阵乘法相比,观察到的速度提升高达 1.6 倍。

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

Beta原型性能改进
torch.compile 中的用户定义 Triton 内核torch.export 添加了新的 API 来指定 dynamic_shapes仅权重量化引入到 Inductor CPU 后端
PyTorch 分布式中的张量并行性异步检查点生成
支持半结构化稀疏性

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

Beta 功能

[Beta] torch.compile 中用户定义 Triton 内核的支持

允许包含 Triton 内核的 PyTorch 代码使用 torch.compile 本地执行。这使用户能够将包含 Triton 内核的代码从 Eager PyTorch 迁移到 torch.compile,而不会遇到性能下降或图中断。原生支持还为 Torch Inductor 提供了预编译用户定义 Triton 内核的机会,并更好地组织 Triton 内核周围的代码,从而实现进一步优化。

您可以在此教程中找到有关如何在 torch.compile 中利用用户定义 Triton 内核的更多信息。

[Beta] 张量并行性引入了更高效的 LLM 训练方式

张量并行 API 促进了跨 GPU/主机进行各种张量操作,并与 FSDP 集成以实现 2D 并行性(跨设备的张量并行性 + 跨主机的数据并行性)。它还提供了一个低级 API,用于构建更高级别的张量并行 API。此 API 已验证支持训练参数超过 1000 亿的 Transformer 模型。

您可以在此教程中找到有关如何在工作流中利用此功能的更多信息。

[Beta] 半结构化稀疏性为用户提供了一种利用加速稀疏推理和内存节省的方法

torch.sparse.SparseSemiStructuredTensor 将半结构化稀疏性实现为张量子类,与密集矩阵乘法相比,观察到的速度提升高达 1.6 倍。

特别是,它增加了

  • 对量化可组合性的额外支持(混合数据类型、去量化融合)
  • 更新了 cuSPARSELt 和 CUTLASS 内核
  • torch.compile 支持

您可以在此处找到有关如何利用半结构化稀疏性的更多信息。

原型功能

[原型] torch.export 添加了新的 API 来指定 dynamic_shapes

您现在可以使用 torch.export.Dim 来更好地表示动态形状,使开发人员能够指定范围(最小值和最大值),这些范围可以在不同的输入维度之间重复使用,这些维度被约束为相等。

要了解有关 torch.export.Dim 以及如何使用它来表达更有趣的关系(例如线性算术表达式)的更多信息,请查看此处的教程。

[原型] 异步检查点生成

异步检查点生成允许用户在生成检查点时继续其训练循环,从而基本卸载了大部分检查点成本。

您可以通过此示例了解如何在自己的工作流中利用此功能。

性能改进

[原型] 仅权重量化引入到 Inductor CPU 后端

PyTorch 2.3 增强了 torch Inductor CPU 后端的 LLM 推理性能。gpt-fast 项目为使用 torch.compile 的 Transformer 文本生成提供了简单高效的 PyTorch 原生加速。在 2.3 之前,仅支持 CUDA 设备,此功能通过为 int4 和 int8 仅权重量化线性提供高度优化的内核来支持 CPU 对应功能。

有关此功能或如何利用此功能的更多信息,请参阅 gpt-fast README