FAX: Масштабируемые и дифференцируемые федеративные примитивы в JAX

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

Авторы: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

Аннотация

Мы представляем FAX, библиотеку на основе JAX, разработанную для поддержки крупномасштабных распределенных и федеративных вычислений как в центрах обработки данных, так и в приложениях, работающих на нескольких устройствах. FAX использует механизмы фрагментации JAX для нативной поддержки TPU и передовых времен выполнения JAX, включая Pathways. FAX встраивает строительные блоки для федеративных вычислений как примитивы в JAX. Это обеспечивает три ключевых преимущества. Во-первых, вычисления FAX могут быть преобразованы в XLA HLO. Во-вторых, FAX предоставляет полную реализацию федеративного автоматического дифференцирования, что значительно упрощает выражение федеративных вычислений. Наконец, вычисления FAX могут быть интерпретированы в существующие производственные системы федеративных вычислений на нескольких устройствах. Мы показываем, что FAX предоставляет легко программируемую, производительную и масштабируемую среду для федеративных вычислений в центрах обработки данных. FAX доступен по адресу 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 .