FAX: JAXにおけるスケーラブルで微分可能なフェデレーテッドプリミティブ

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

著者: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

要旨

本論文では、データセンターおよびクロスデバイスアプリケーションにおける大規模分散型および連合型計算をサポートするために設計されたJAXベースのライブラリ、FAXを紹介します。FAXは、JAXのシャーディングメカニズムを活用し、TPUやPathwaysを含む最先端のJAXランタイムをネイティブにターゲットとします。FAXは、連合型計算のためのビルディングブロックをJAXのプリミティブとして組み込みます。これにより、3つの主要な利点が得られます。第一に、FAXの計算はXLA HLOに変換可能です。第二に、FAXは連合型自動微分の完全な実装を提供し、連合型計算の表現を大幅に簡素化します。最後に、FAXの計算は既存の生産環境のクロスデバイス連合型計算システムに解釈可能です。FAXが、データセンターにおける連合型計算のためのプログラミングが容易で、高性能かつスケーラブルなフレームワークを提供することを示します。FAXはhttps://github.com/google-research/google-research/tree/master/faxで利用可能です。

English

We present FAX, a JAX-based library designed to support large-scale distributed and federated computations in both data center and cross-device applications. FAX leverages JAX's sharding mechanisms to enable native targeting of TPUs and state-of-the-art JAX runtimes, including Pathways. FAX embeds building blocks for federated computations as primitives in JAX. This enables three key benefits. First, FAX computations can be translated to XLA HLO. Second, FAX provides a full implementation of federated automatic differentiation, greatly simplifying the expression of federated computations. Last, FAX computations can be interpreted out to existing production cross-device federated compute systems. We show that FAX provides an easily programmable, performant, and scalable framework for federated computations in the data center. FAX is available at https://github.com/google-research/google-research/tree/master/fax .