FAX: Primitivas Federadas Escalables y Diferenciables en JAX

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

Autores: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

Resumen

Presentamos FAX, una biblioteca basada en JAX diseñada para soportar cálculos distribuidos y federados a gran escala, tanto en aplicaciones de centros de datos como en dispositivos cruzados. FAX aprovecha los mecanismos de fragmentación de JAX para permitir la orientación nativa de TPUs y los entornos de ejecución más avanzados de JAX, incluyendo Pathways. FAX incorpora bloques de construcción para cálculos federados como primitivas en JAX. Esto ofrece tres beneficios clave. Primero, los cálculos de FAX pueden traducirse a XLA HLO. Segundo, FAX proporciona una implementación completa de la diferenciación automática federada, simplificando enormemente la expresión de cálculos federados. Por último, los cálculos de FAX pueden interpretarse en sistemas de cómputo federado de dispositivos cruzados ya existentes en producción. Demostramos que FAX ofrece un marco de trabajo fácilmente programable, eficiente y escalable para cálculos federados en el centro de datos. FAX está disponible en 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 .