我们很高兴地宣布 PyTorch® 2.3 版本发布了 (发布说明)! PyTorch 2.3 支持在 torch.compile 中使用用户自定义的 Triton 内核,用户可以将自己的 Triton 内核从 Eager 模式迁移过来,而不会遇到性能退化或图断裂。Tensor Parallelism 改进了使用 PyTorch 原生函数训练大型语言模型的体验,在训练 100B 参数模型时已得到验证。此外,半结构化稀疏性 (semi-structured sparsity) 作为 Tensor 的一个子类实现,相比密集矩阵乘法,观察到高达 1.6 倍的加速。
自 PyTorch 2.2 发布以来,此版本包含 3393 个提交和 426 位贡献者。我们衷心感谢我们专注的社区所做的贡献。一如既往,我们鼓励大家尝试这些新特性,并在我们改进 2.3 版本的过程中报告遇到的任何问题。有关如何开始使用 PyTorch 2 系列的更多信息,请访问我们的入门页面。
Beta 功能 | 原型功能 | 性能改进 |
torch.compile 中的用户自定义 Triton 内核 | torch.export 添加新 API 用于指定 dynamic_shapes | 在 Inductor CPU 后端引入了 Weight-Only-Quantization |
PyTorch Distributed 中的 Tensor parallelism | 异步检查点生成 | |
支持半结构化稀疏性 |
*查看完整公共特性提交列表,请点击这里。
Beta 功能
[Beta] 在 torch.compile 中支持用户自定义 Triton 内核
允许包含 Triton 内核的 PyTorch 代码使用 torch.compile 本地执行。这使得用户可以将包含 Triton 内核的代码从 Eager PyTorch 迁移到 torch.compile,而不会遇到性能退化或图断裂。原生支持也为 Torch Inductor 预编译用户自定义的 Triton 内核创造了机会,并能更好地组织 Triton 内核周围的代码,从而实现进一步优化。
有关如何在 torch.compile 中使用用户自定义 Triton 内核的更多信息,请参阅此教程。
[Beta] Tensor Parallelism 引入了更高效的 LLM 训练方法
Tensor Parallel API 方便了跨 GPU/主机进行各种张量操作,并与 FSDP 集成以实现二维并行(跨设备进行张量并行 + 跨主机进行数据并行)。它还提供了构建更高级别 Tensor 并行 API 的低级 API。该 API 已经过验证,支持训练参数超过 1000 亿的 Transformer 模型。
有关如何在您的工作流程中使用此功能的更多信息,请参阅此教程。
[Beta] 半结构化稀疏性为用户提供了一种利用加速稀疏推理和节省内存的方法
torch.sparse.SparseSemiStructuredTensor 将半结构化稀疏性实现为 Tensor 的一个子类,相比密集矩阵乘法,观察到高达 1.6 倍的加速。
特别是,它增加了
- 对量化组合性的额外支持(混合数据类型,反量化融合)
- 更新了 cuSPARSELt 和 CUTLASS 内核
- 支持 torch.compile
有关如何利用半结构化稀疏性的更多信息,请参阅这里。
原型功能
[原型] torch.export 添加新 API 用于指定 dynamic_shapes
现在您可以使用 torch.export.Dim 更好地表示动态形状,开发者可以指定可在不同输入维度之间重用(并约束其相等)的范围(最小值和最大值)。
要了解更多关于 torch.export.Dim 以及如何用它表达更有趣的关系(例如线性算术表达式),请查阅此教程。
[原型] 异步检查点生成
异步检查点生成允许用户在生成检查点的同时继续训练循环,这本质上分担了大部分检查点生成的开销。
您可以通过此示例了解如何在您自己的工作流程中使用此功能。
性能改进
[原型] 在 Inductor CPU 后端引入 Weight-Only-Quantization
PyTorch 2.3 增强了 torch inductor CPU 后端上的 LLM 推理性能。项目 gpt-fast 提供了一种简单高效的 PyTorch 原生加速方法,用于通过 torch.compile 进行 Transformer 文本生成。在 2.3 版本之前,仅支持 CUDA 设备,此功能通过为 int4 和 int8 仅权重量化 Linear 提供高度优化的内核,从而支持 CPU。
有关此功能的更多信息/如何使用此功能,请参阅 gpt-fast README。