FAX: Scalable and Differentiable Federated Primitives in JAX
CoRR(2024)
摘要
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 .
更多查看译文
AI 理解论文
溯源树
样例
生成溯源树,研究论文发展脉络
Chat Paper
正在生成论文摘要