FAX: Skalierbare und differenzierbare föderierte Primitive in JAX

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

Autoren: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

Zusammenfassung

Wir präsentieren FAX, eine auf JAX basierende Bibliothek, die für die Unterstützung von verteilten und föderierten Berechnungen im großen Maßstab in Rechenzentren und plattformübergreifenden Anwendungen entwickelt wurde. FAX nutzt JAX's Sharding-Mechanismen, um eine native Ausrichtung auf TPUs und modernste JAX-Laufzeiten, einschließlich Pathways, zu ermöglichen. FAX integriert Bausteine für föderierte Berechnungen als Primitiven in JAX. Dies bietet drei wesentliche Vorteile. Erstens können FAX-Berechnungen in XLA HLO übersetzt werden. Zweitens bietet FAX eine vollständige Implementierung der föderierten automatischen Differentiation, was die Ausdrucksweise von föderierten Berechnungen erheblich vereinfacht. Schließlich können FAX-Berechnungen in bestehende produktionsreife plattformübergreifende föderierte Rechensysteme interpretiert werden. Wir zeigen, dass FAX ein leicht programmierbares, leistungsstarkes und skalierbares Framework für föderierte Berechnungen im Rechenzentrum bietet. FAX ist verfügbar unter 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 .