FAX: Scalable and Differentiable Federated Primitives in JAX

Core Concepts

FAX is a JAX-based library designed for large-scale distributed and federated computations, providing efficient and scalable frameworks for data center federated computations.

Abstract

Abstract FAX leverages JAX's sharding mechanisms for native targeting of TPUs and state-of-the-art runtimes. Provides full implementation of federated automatic differentiation, simplifying expression of federated computations. Introduction Importance of scaling compute-intensive programs across distributed environments for modern ML success. Describes FAX as a library bringing benefits like sharding, JIT compilation, and AD to federated learning computations. Contributions Embeds a federated programming model into JAX via primitives mechanism. Enables forward- and reverse-mode differentiation for full implementation of Federated AD framework. System Design Represents federated arrays with extra dimension indicating placement. Implements federated computations on JAX arrays using Primitive mechanism. Scalability and Efficiency Demonstrates scalability through weak scaling experiments with FedAvg on transformer language models. Interpreting FAX to Production Platforms Discusses preserving placement information and integrating federated AD into production systems. Discussion Explains the importance of implementing federated AD for efficient algorithm development in FL.

Stats

FAXはJAXのシャーディングメカニズムを活用してTPUと最新のランタイムにネイティブにターゲティングします。 FAXはフル実装のフェデレーテッド自動微分を提供し、フェデレーテッド計算の表現を大幅に簡素化します。

Quotes

Key Insights Distilled From

FAX

by Keith Rush,Z... at arxiv.org 03-13-2024

https://arxiv.org/pdf/2403.07128.pdf

Deeper Inquiries

なぜFederated ADを実装する必要があるのですか？

Federated Automatic Differentiation（AD）を実装する主な理由は、効率的なアルゴリズム開発を容易にするためです。例えば、FAXプログラム内でクライアントごとの平均損失を計算する際に、jax.grad(federated_loss)を呼び出すだけで、クライアントごとの平均勾配を即座に取得できます。これにより、サーバーで最適化ステップ（例：Optaxなどのライブラリ使用）を行い、FedSGDアルゴリズム（McMahan et al., 2017）が直ちに得られます。 このようにしてFederated ADは、federated hypergradient descentやクライアント間の重み付け平均学習などのアルゴリズムも可能とし、差分プライバシーやセキュア集約（Bonawitz et al., 2017）などのプライバシー保護メカニズムと互換性が確保されます。FL全てのアルゴリズムがfederated AD経由で表現可能ではありませんが（Charles and Rush, 2022）、使いやすいfederated AD実装はFL分野におけるアルゴリズム開発や研究を加速させることが期待されています。

FAXの内部シャーディング注釈なしで弱いスケーリング挙動を得ることができますか

弱いスケーリング挙動は、「ワークロード量」と「コンピューティング資源」両方が同時に拡大した場合の計算時間変化を指します。一般的に現代MLシステムでは近似的な弱いスケーリングパフォーマンスを目指します。しかし、「GSPMD」等高度なコンパイラでも十分性能向上しなかったfor-loop実装から派生したダブルfor-loop処理方法も存在します。 具体的にはJITコンパイルされたfor-loop処理法と比較しても，そのランタイ ム時間はコホートサ イ ズ の増加割合通り 線形 ス カ レ を 示しました．つまり，追加された資源規模 を利用して演算 処 理 を 行うこと が不可 能だったこ tわけです．

Federated ADはFLアルゴリズム開発にどのような利点をもたらすと考えられますか

FAX-based FedAvg scales as the cohort size and number of TPU chips increase, across a range model sizes. FAX-exhibits near-constant runtime for a fixed model size, even up to cohort sizes of 128 or 512. This is highly non-trivial. Because FedAvg involves parallel model training across cohorts, and for multiple steps per client, the per-round workload size (in terms of total floating point operations) is at least as large as (4×(model size)×(cohort size)). To see this, note that in each round, each client updates its own local model four times. JIT compilation alone is not enough. ML research often involves writing custom training loops. For FL research, this is often done by writing a double for-loop that iterates over clients in a cohort, and over the batches held by each client. The outer loop here has no data dependency, meaning that the values returned by iterations of this loop are not processed as inputs to the next iteration. One might therefore imagine that a sufficiently advanced compiler could detect this fact, and parallelize client training when possible (e.g., within resource constraints in the data center environment). This can be a difficult task for a compiler. To illustrate this difficulty we implemented...