FAX: Scalable and Differentiable Federated Primitives in JAX
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 .
65f11795c8c8fba63192afc4
Keith Rush
claimed_verified
2024-03-13T16:15:00.532Z
65f1cf3aea7a16c6295f4737
/avatars/905106897916e87e114ccc2d33e8e1c2.svg
Keith Rush
user
jkr26
65f11795c8c8fba63192afc5
Zachary Charles
admin_assigned
2024-03-14T08:09:56.031Z
65f1c0bf22ea2a37a28b6466
/avatars/d9a11158bad338ada0d044476d8681c1.svg
Zachary Charles
user
zcharles8
65f11795c8c8fba63192afc6
Zachary Garrett
admin_assigned
2024-03-14T08:10:16.343Z
65f1be9109a3517bd1c8e17b
/avatars/bc4112a3f749fe969ebb633b2e0cbc14.svg
Zachary Garrett
user
ZacharyGarrett
FAX: Scalable and Differentiable Federated Primitives in JAX
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 .
FAX is a JAX-based library enabling distributed and federated computations using TPUs, XLA HLO, and native federated automatic differentiation.
TPUs XLA HLO federated automatic differentiation