FAX: Primitivi Federati Scalabili e Differenziabili in JAX

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

Autori: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

Abstract

Presentiamo FAX, una libreria basata su JAX progettata per supportare calcoli distribuiti e federati su larga scala, sia in applicazioni di data center che cross-device. FAX sfrutta i meccanismi di partizionamento di JAX per consentire il targeting nativo di TPU e runtime JAX all'avanguardia, inclusi Pathways. FAX integra blocchi costitutivi per calcoli federati come primitive in JAX. Ciò offre tre vantaggi principali. Innanzitutto, i calcoli FAX possono essere tradotti in XLA HLO. In secondo luogo, FAX fornisce un'implementazione completa della differenziazione automatica federata, semplificando notevolmente l'espressione di calcoli federati. Infine, i calcoli FAX possono essere interpretati nei sistemi di calcolo federato cross-device esistenti in produzione. Dimostriamo che FAX offre un framework facilmente programmabile, performante e scalabile per calcoli federati nel data center. FAX è disponibile all'indirizzo 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 .