Scalable Cross-Facility Federated Learning for Scientific Foundation Models on Multiple Supercomputers

TL;DR

This paper develops a scalable federated learning framework for scientific models across multiple supercomputers, demonstrating practical deployment, analyzing heterogeneity impacts, and fine-tuning a large language model for chemistry applications.

Contribution

It introduces a comprehensive cross-facility federated learning framework tailored for heterogeneous HPC environments, addressing deployment challenges and performance considerations.

Findings

Federated learning across HPC facilities is practically feasible.

Heterogeneity significantly impacts training performance.

Scheduler-aware algorithms are crucial for future HPC federated learning deployments.

Abstract

Artificial Intelligence for scientific applications increasingly requires training large models on data that cannot be centralized due to privacy constraints, data sovereignty, or the sheer volume of data generated. Federated learning (FL) addresses this by enabling collaborative training without centralizing raw data, but scientific applications demand model scales that requires extensive computing resources, typically offered at High Performance Computing (HPC) facilities. Deploying FL experiments across HPC facilities introduces challenges beyond cloud or enterprise settings. We present a comprehensive cross-facility FL framework for heterogeneous HPC environments, built on Advanced Privacy-Preserving Federated Learning (APPFL) framework with Globus Compute and Transfer orchestration, and evaluate it across four U.S. Department of Energy (DOE) leadership-class supercomputers. We demonstrate that FL experiments across HPC facilities are practically achievable, characterize key sources of heterogeneity impacting the training performance, and show that algorithmic choices matter significantly under realistic HPC scheduling conditions. We validate the scientific applicability by fine-tuning a large language model on a chemistry instruction dataset, and identify scheduler-aware algorithm design as a critical open challenge for future deployments.