FAX: JAX에서의 확장 가능하고 미분 가능한 연합 학습 프리미티브

FAX: Scalable and Differentiable Federated Primitives in JAX

March 11, 2024

저자: Keith Rush, Zachary Charles, Zachary Garrett

cs.AI

초록

우리는 데이터 센터 및 크로스 디바이스 애플리케이션에서 대규모 분산 및 연합 계산을 지원하기 위해 설계된 JAX 기반 라이브러리인 FAX를 소개합니다. FAX는 JAX의 샤딩 메커니즘을 활용하여 TPU 및 Pathways와 같은 최신 JAX 런타임을 네이티브로 타겟팅할 수 있도록 합니다. 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 .