在我们庆祝 OpenXLA、PyTorch 2.0 和 PyTorch/XLA 2.0 发布之际，值得回顾一下并分享我们对中短期未来发展方向的看法。随着 PyTorch 在 AI 领域的领先地位和 XLA 对一流编译器功能的支持，PyTorch/XLA 已准备好为模型训练和推理提供尖端开发堆栈。为了实现这一目标，我们将在三个主要领域进行投资
- 训练大型模型 – 大型语言模型 (LLM) 和扩散模型已迅速普及，当今许多尖端应用都建立在它们之上。此外，训练这些模型需要规模，更具体地说，需要在数千个加速器上进行训练的能力。为了实现这一目标，我们正在投资于以下功能：用于混合精度训练的 AMP、用于提高运行时性能的 PjRt、用于高效模型分片的 SPMD / FSDP、用于实现新研究方法的动态形状、通过 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 完全在开源中开发，因此请在 GitHub 上提交问题、提交拉取请求并发送 RFC，以便我们公开协作。您还可以 亲自尝试 在各种 XLA 设备（包括 TPU 和 GPU）上使用 PyTorch/XLA。
祝好，
Google 的 PyTorch/XLA 团队