FAX：在 JAX 中可擴展且可微分的聯邦原語

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

作者: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

摘要

我們介紹了 FAX，這是一個基於 JAX 設計的庫，旨在支持數據中心和跨設備應用中的大規模分佈式和聯邦計算。FAX利用 JAX 的分片機制，實現對 TPU 和最先進的 JAX 運行時（包括 Pathways）的本地定位。FAX將聯邦計算的構建模塊作為 JAX 中的原始操作內嵌。這帶來三個關鍵好處。首先，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 .