AI开发大一统:谷歌OpenXLA开源,整合所有框架和AI芯片( 二 )


透明度:所有决定和理由都应该对公众清晰可见 。
OpenXLA 生态系统:性能、规模和可移植能力
OpenXLA 通过模块化工具链消除了机器学习开发人员的障碍,它通过通用编译器接口得到所有领先框架的支持,利用可移植的标准化模型表示,并提供具有强大的目标向和特定硬件优化的特定领域编译器 。该工具链包括 XLA、StableHLO 和 IREE,所有这些工具都利用 MLIR:一种编译器基础架构,使机器学习模型能够在硬件上一致地表示、优化和执行 。

AI开发大一统:谷歌OpenXLA开源,整合所有框架和AI芯片

文章插图
 
OpenXLA 主要亮点
机器学习用例的范围
OpenXLA 当前的使用涵盖了 ML 用例的范围,包括在阿里云上对 DeepMind 的 AlphaFold、GPT2 和 Swin Transformer 等模型进行全面训练,以及在 Amazon.com 上进行多模态 LLM 训练 。Waymo 等客户利用了 OpenXLA 进行车载实时推理 。此外,OpenXLA 还用于优化配备 AMD RDNA™ 3 的本地机器上的 Stable Diffusion 服务 。
最佳性能,开箱即用
OpenXLA 使开发人员无需编写特定于设备的代码,即可轻松加快模型性能 。它具有整体模型优化功能,包括简化代数表达式、优化内存数据布局以及改进调度以减少峰值内存使用和通信开销 。高级算子融合和内核生成有助于提高设备利用率并降低内存带宽要求 。
轻松扩展工作负载
开发高效的并行化算法非常耗时并且需要专业知识 。借助 GSPMD 等功能,开发人员只需注释关键张量的一个子集,然后编译器就可以使用这些子集自动生成并行计算 。这消除了跨多个硬件主机和加速器对模型进行分区和高效并行化所需的大量工作 。
便携性和可选性
OpenXLA 为多种硬件设备提供开箱即用的支持,包括 AMD 和 NVIDIA GPU、x86 CPU 和 Arm 架构以及 ML 加速器,如 Google TPU、AWS Trainium 和 Inferentia、Graphcore IPU、Cerebras Wafer-Scale Engine 等等 。OpenXLA 还通过 StableHLO 支持 TensorFlow、PyTorch 和 JAX,StableHLO 是一个用作 OpenXLA 输入格式的可移植层 。
灵活性
OpenXLA 为用户提供了手动调整模型热点的灵活性 。自定义调用等扩展机制使用户能够用 CUDA、HIP、SYCL、Triton 和其他内核语言编写深度学习原语,从而能够充分利用硬件特性 。
StableHLO
StableHLO 是 ML 框架和 ML 编译器之间的一个可移植层,是一个支持动态、量化和稀疏性的高级运算(HLO)的运算集 。此外,它可以被序列化为 MLIR 字节码以提供兼容性保证 。所有主要的 ML 框架(JAX、PyTorch、TensorFlow)都可以产生 StableHLO 。2023 年,谷歌计划与 PyTorch 团队紧密合作,实现与 PyTorch 2.0 版本的整合 。
参考内容:
https://opensource.googleblog.com/2023/03/openxla-is-ready-to-accelerate-and-simplify-ml-development.html?m=1
https://venturebeat.com/ai/google-reveals-whats-next-for-cloud-ai/




推荐阅读