作者：Douglas Yarrington (Google TPgM)、James Rubin (Google PM)、Neal Vaidya (NVIDIA TME)、Jay Rodge (NVIDIA PMM)

NVIDIA 和 Google 共同宣布了新的里程碑和计划，通过利用 XLA 的强大功能，优化 TensorFlow 和 JAX 在 Ampere 和最近发布的 Hopper GPU 架构上的性能：XLA 是由 Google 构建的高效、灵活且可扩展的 ML 编译器。我们将加深与专注于在当前可用的 A100 GPU 上提供更高性能的专门工程团队之间的合作。NVIDIA 和 Google 还将共同支持最近发布的 H100 GPU 中的独特功能，包括 Transformer Engine，它支持硬件加速的 8 位浮点 (FP8) 数据类型和 transformer 库。

我们宣布了 TensorFlow 性能的提升，XLA 中新的 NVIDIA GPU 特定功能，以及 JAX 首次发布用于多节点、多 GPU 训练的功能，这将显著提高大型语言模型 (LLM) 的训练速度。我们预计 Hopper 架构将特别适合 LLM。

NVIDIA H100 Tensor Core GPU

用于 GPU 的 XLA

Google 在 NVIDIA GPU 上使用 LLM 能够实现高性能，这得益于一项名为 XLA 的重要技术，它支持所有领先的 ML 框架，例如 TensorFlow、JAX 和 PyTorch。在 Google，超过 90% 的 ML 编译（涵盖研究和生产）都依赖于 XLA。这些应用涵盖了各种 ML 使用场景，从 DeepMind 和 Google Research 中的超大规模模型训练，到我们产品中的优化部署，再到 Waymo 的边缘推理。

XLA 的强大功能集加速了大型语言模型的性能，并 解决了当今行业中大多数大型模型挑战。例如，XLA 独有的 SPMD 功能可以自动化大多数将模型分配到多个核心和设备所需的作业，使大型模型训练的扩展性和性能大幅提高。XLA 还可以自动识别并选择最适合目标后端的最佳手写库实现，例如 cuDNN 用于 CUDA 芯片组。否则，XLA 可以本地生成优化代码以实现高效执行。

我们一直在与 NVIDIA 合作开发一些令人振奋的功能和集成，这些功能和集成将进一步优化 GPU 上的 LLM。我们最近使全约简等集合操作能够与计算并行运行。这显著减少了客户的端到端延迟。此外，我们启用了对 bfloat16 的支持，这使得计算性能比 32 位浮点提高了 4.5 倍，同时保持了相同的值动态范围。

我们共同的努力意味着 XLA 与 NVIDIA 的 AI 工具深度集成，并且可以更好地利用 NVIDIA 的 AI 硬件优化库。在 2023 年第一季度，我们将发布 XLA-cuDNN 图 API 集成，该集成为客户提供卷积/矩阵乘法运算和 transformer 中的多头注意力机制的优化融合，从而改进内存使用并加快 GPU 内核执行速度。结果，开销显著降低，性能显著提升。