FAX : Primitives fédérées scalables et différentiables dans JAX

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

Auteurs: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

Résumé

Nous présentons FAX, une bibliothèque basée sur JAX conçue pour supporter des calculs distribués et fédérés à grande échelle, aussi bien dans les centres de données que dans les applications inter-appareils. FAX exploite les mécanismes de partitionnement de JAX pour cibler nativement les TPUs et les environnements d'exécution de pointe de JAX, y compris Pathways. FAX intègre des blocs de construction pour les calculs fédérés en tant que primitives dans JAX. Cela offre trois avantages clés. Premièrement, les calculs FAX peuvent être traduits en XLA HLO. Deuxièmement, FAX fournit une implémentation complète de la différenciation automatique fédérée, simplifiant grandement l'expression des calculs fédérés. Enfin, les calculs FAX peuvent être interprétés pour fonctionner avec les systèmes de calcul fédéré inter-appareils existants en production. Nous démontrons que FAX offre un cadre facilement programmable, performant et évolutif pour les calculs fédérés dans les centres de données. FAX est disponible à l'adresse suivante : 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 .