在我们庆祝 OpenXLA、PyTorch 2.0 和 PyTorch/XLA 2.0 发布之际，有必要回顾一下，分享我们对短期到中期发展的看法。随着 PyTorch 在 AI 领域的主导地位以及 XLA 对一流编译器功能的支持，PyTorch/XLA 已蓄势待发，可为模型训练和推理提供尖端开发堆栈。为实现这一目标，我们计划在以下三个主要领域进行投资：
- 大型模型训练 – 大型语言模型 (LLM) 和扩散模型已迅速普及，如今许多尖端应用都建立在它们之上。此外，训练这些模型需要规模，更具体地说，是跨数千个加速器进行训练的能力。为实现这一目标，我们正在投资的功能包括：用于混合精度训练的 AMP、用于提高运行时性能的 PjRt、用于高效模型分片的 SPMD/FSDP、用于实现新研究方法的动态形状、通过 Ray 和 tf.data 加速数据加载，以及将所有这些功能打包成无缝工作流程的工具链。其中一些功能已处于实验或 Beta 阶段，另一些功能将于今年推出，其中许多功能都大量利用了底层的 OpenXLA 编译器堆栈。
- 模型推理 – 随着大型模型在规模和计算成本上持续增长，当这些模型不断进入应用程序时，部署成为下一个挑战。随着 PyTorch 2.0 版本中 Dynamo 的引入，PyTorch/XLA 提供了具有竞争力的推理性能。然而，我们正在整合其他面向推理的功能，包括模型服务支持、用于分片大型模型的 Dynamo，以及通过 Torch.Export 和 StableHLO 进行量化。
- 生态系统集成 – 我们正在扩展与 Hugging Face 和 PyTorch Lightning 的集成，以便用户可以通过熟悉的 API 利用即将推出的 PyTorch/XLA 尖端功能（例如 Hugging Face 中的 FSDP 支持）和下游 OpenXLA 功能（例如量化）。
此外，PyTorch/XLA 将迁移到开源 OpenXLA 作为其默认的下游编译器；这使得 PyTorch 社区能够访问领先的、与框架无关的编译器堆栈，该堆栈享有行业范围内的贡献和创新。为实现这一目标，我们将开始支持 StableHLO。因此，OpenXLA 将取代现有的 TF:XLA 依赖项，从而全面简化依赖项并从更广泛的编译器生态系统中获得优势。PyTorch/XLA 还将在迁移后停止支持 XRT 运行时。您可以在下面看到由此产生的高级堆栈，其中 TensorFlow 依赖项已被划掉
图：此处展示了即将推出的 PyTorch/XLA 功能和集成
我们对 PyTorch/XLA 的未来充满期待，并邀请社区加入我们。PyTorch/XLA 完全在开源中开发，因此请在 GitHub 上提交问题、发送拉取请求和 RFC，以便我们公开协作。您也可以在各种 XLA 设备（包括 TPU 和 GPU）上亲自试用 PyTorch/XLA。
祝好，
Google 的 PyTorch/XLA 团队