FAX: Primitivas Federadas Escaláveis e Diferenciáveis em JAX

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

Autores: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

Resumo

Apresentamos o FAX, uma biblioteca baseada em JAX projetada para suportar computações distribuídas e federadas em grande escala, tanto em aplicações de data center quanto em dispositivos cruzados. O FAX aproveita os mecanismos de fragmentação do JAX para permitir o direcionamento nativo de TPUs e os runtimes de última geração do JAX, incluindo o Pathways. O FAX incorpora blocos de construção para computações federadas como primitivas no JAX. Isso possibilita três benefícios principais. Primeiro, as computações do FAX podem ser traduzidas para XLA HLO. Segundo, o FAX fornece uma implementação completa de diferenciação automática federada, simplificando significativamente a expressão de computações federadas. Por último, as computações do FAX podem ser interpretadas em sistemas de computação federada em dispositivos cruzados já existentes em produção. Demonstramos que o FAX oferece uma estrutura facilmente programável, de alto desempenho e escalável para computações federadas no data center. O FAX está disponível em 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 .