跳转到主要内容
博客

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 系列的更多信息,请访问我们的入门页面。

测试版原型性能改进
torch.compile 中用户定义的 Triton 内核torch.export 添加了新 API 以指定 dynamic_shapesInductor CPU 后端引入了仅权重量化
PyTorch Distributed 中的张量并行性异步检查点生成
支持半结构化稀疏性

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

测试版功能

[测试版] 在 torch.compile 中支持用户定义的 Triton 内核

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

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

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

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

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

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

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