随着我们庆祝 OpenXLA、PyTorch 2.0 和 PyTorch/XLA 2.0 的发布，值得回顾一下在短期内到中期内我们认为它会走向何方。随着 PyTorch 在 AI 领域的广泛应用和 XLA 支持一流的编译器特性，PyTorch/XLA 正处于为模型训练和推理提供尖端开发栈的有利位置。为了实现这一点，我们认为需要在以下三个主要领域进行投资：
- 训练大型模型 - 大型语言模型和扩散模型迅速崛起，并在今天许多前沿应用中得到了广泛应用。更进一步，训练这些模型需要规模，特别是能够在数千个加速器上训练的能力。为了实现这一点，我们正在投资于以下功能：混合精度训练的 AMP、提高运行时性能的 PjRt、高效模型分片技术的 SPMD/FSDP、Dynamic Shapes 以实现新的研究方法、通过 Ray 和 tf.data 加快数据加载，以及将这些功能打包成一个无缝工作流程的工具链。其中一些功能已经在实验或测试阶段可用，而其他功能将在今年推出，许多都将大量利用底层的 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 完全开源开发，请将问题、提交拉取请求和发送 RFC 发送到 GitHub，以便我们公开协作。您还可以在包括 TPUs 和 GPU 在内的各种 XLA 设备上尝试 PyTorch/XLA。
干杯，
谷歌 PyTorch/XLA 团队
