FAX: Schaalbare en Differentieerbare Federatieve Primitieven in JAX

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

Auteurs: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

Samenvatting

We presenteren FAX, een op JAX gebaseerde bibliotheek die is ontworpen om grootschalige gedistribueerde en federatieve berekeningen te ondersteunen, zowel in datacenter- als cross-device-toepassingen. FAX maakt gebruik van JAX's sharding-mechanismen om native ondersteuning te bieden voor TPU's en state-of-the-art JAX-runtimes, waaronder Pathways. FAX integreert bouwstenen voor federatieve berekeningen als primitieven in JAX. Dit biedt drie belangrijke voordelen. Ten eerste kunnen FAX-berekeningen worden vertaald naar XLA HLO. Ten tweede biedt FAX een volledige implementatie van federatieve automatische differentiatie, wat het uitdrukken van federatieve berekeningen aanzienlijk vereenvoudigt. Tot slot kunnen FAX-berekeningen worden geïnterpreteerd naar bestaande productie-omgevingen voor cross-device federatieve berekeningen. We tonen aan dat FAX een eenvoudig programmeerbaar, performant en schaalbaar raamwerk biedt voor federatieve berekeningen in het datacenter. FAX is beschikbaar op 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 .