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

Yijiang Li; Zilinghan Li; Kyle Chard; Ian Foster; Todd Munson; Ravi Madduri; Kibaek Kim

cs.LGMar 20, 2026

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

Authors: Yijiang Li, Zilinghan Li, Kyle Chard, Ian Foster, Todd Munson, Ravi Madduri, Kibaek Kim

Organizations: Mathematics and Computer Science Division, Argonne National Laboratory · Data Science and Learning Division, Argonne National Laboratory · Department of Computer Science, University of Chicago

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.

Figures & tables

			Samples
Task Group	HPC	Task	Per Task	Per Group
Name Conversion	Perlmutter	IUPAC to molecular formula	304,490	1,217,627
		IUPAC to SMILES	304,379
		SMILES to molecular formula	304,379
		SMILES to IUPAC	304,379
Property Prediction	Polaris	ESOL(Water Solubility)	1,111	78,319

Table 1: We distribute the 3.3M SMolInstruct training data samples across four facilities and group them by task group as follows. Refer to Yu et al. [ 55 ] for more details on the tasks and samples.

Supercomputer	GPU architecture	GPU memory	DeepSpeed stage	Micro batch size
Polaris	NVIDIA A100	40GB	ZeRO-3	6
Perlmutter	NVIDIA A100	40GB	ZeRO-3	6
Perlmutter	NVIDIA A100	80GB	ZeRO-1	16
Aurora	Intel Max 1550 Series	64GB	ZeRO-1	8
Frontier	AMD MI250X	64GB	ZeRO-1	12

Table 2: Configurations for throughput scaling experiments. GPU memory capacity largely determines the required DeepSpeed ZeRO optimization stage: supercomputers with smaller memory GPUs must employ ZeRO-3, which achieves memory efficiency at the cost of increased inter-GPU communication overhead, while supercomputers with larger memory GPUs can use the less communication-intensive ZeRO-1. Micro batch size is set per GPU and held fixed across node counts.

Figure 1: Throughput scaling with fixed micro-batch size per GPU. The left panel shows throughput (samples per second) as a function of node count, while the right panel shows the same results as a function of total GPU count. The consistent scaling patterns across both panels confirm that memory capacity and the resulting optimization strategy, rather than node or GPU count alone, are the primary drivers of throughput heterogeneity across supercomputers.

Figure 2: FedAvg test loss over eight global aggregation rounds using 64 nodes per supercomputer (client). The global model (purple) consistently achieves the lowest test loss across all 14 tasks, decreasing from 1.39 to 0.37, while local client models show higher losses reflecting specialization to their assigned task groups.

Facility	GPU architecture	Nodes	Total GPUs	Effective batch size
Polaris	NVIDIA A100	63 (1 node failure)	252	$1512$
Perlmutter	NVIDIA A100 40GB	64	256	$1536$
Aurora	Intel Max 1550 Series	64	768	$6144$
Frontier	AMD MI250X	64	512	$4096$

Table 3: Hardware specifications for the 64-node large-scale co-scheduled experiment. Aurora and Frontier’s higher total GPU counts result in substantially larger effective batch sizes compared to Polaris and Perlmutter. Polaris operated with 63 nodes due to a single node failure during the reservation.

Facility	GPU architecture	Nodes	Total GPUs	Effective batch size
Polaris	NVIDIA A100	2	8	$48$
Perlmutter	NVIDIA A100 40GB	2	8	$48$
Aurora	Intel Max 1550 Series	2	24	$192$
Frontier	AMD MI250X	2	16	$128$

Table 4: Hardware specifications for the two-node small-scale experiment. Same as the 64-node setup, Aurora and Frontier’s higher multi-GPU-per-node architectures result in larger effective batch sizes relative to Polaris and Perlmutter despite the same node count.

Algorithm	Total	Polaris	Aurora	Frontier	Perlmutter
FedAvg	40	10	10	10	10
FedAsync	40	11	6	14	9
FedCompass	40	11	9	11	9
FedBuff	40	8	5	17	10

Table 5: Distribution of local training rounds across four clients for each algorithm over 40 total rounds. FedAvg distributes rounds evenly due to its synchronous aggregation requirement. Asynchronous algorithms produce uneven distributions driven by each client’s computational speed and queue behavior.

Figure 3: Test loss progression of four FL algorithms under realistic queueing conditions using two nodes per facility. Each panel shows local client models alongside the global aggregated model within the same wall-clock time budget of 17,000 seconds. FedCompass achieves the lowest final global test loss, while FedAvg produces the most stable client trajectories.

Figure 4: Globus Transfer communication efficiency across clients for five models ranging from OPT-125m to Llama2-13b. Panel a shows the linear relationship between model parameters and storage size in BF16 format. Panel b shows transfer speed versus model size, with Aurora and Polaris achieving the highest speeds due to co-location with the server at Argonne. Panel c summarizes average transfer speeds over all five models per destination facility.

Destination	Min	Avg	Max	Mean Dev.
Aurora (ALCF)	0.122	0.266	0.432	0.094
Polaris (ALCF)	0.112	0.210	0.342	0.074
Frontier (OLCF)	17.236	17.281	17.321	0.167
Perlmutter (NERSC)	45.178	45.205	45.264	0.286

Table 6: Round-trip latency (ms) between the server and each client, measured via ping over 20 rounds. Aurora and Polaris exhibit sub-millisecond latency due to co-location with the server at Argonne National Laboratory, while Frontier and Perlmutter face substantially higher wide-area network latency.

Figure 5: Polaris queue time with varying number of nodes and walltime requested. Solid lines represent the mean queue wait time and shaded regions represent one standard deviation. Different queues are designed to handle jobs with different characteristics in order to maximize supercomputer utilization. Each queue imposes its own constraints on the maximum number of nodes requestable, the maximum walltime allowed, and the allocation of available node-hours.

Figure 6: Overall implementation of the cross-facility FL framework. The FL server (blue) comprises a Server Agent handling aggregation and scheduling, and a Server Communicator using Globus Compute and ProxyStore for model exchange. The FL client (yellow) comprises a Client Agent running the Trainer and DeepSpeed via a Globus Compute MPI Engine, and a Client Communicator mirroring the server-side communication components. Model parameters are transferred between server and clients via Globus Transfer.

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 943 papers

FedQueue: Queue-Aware Federated Learning for Cross-Facility HPC Training
May 4, 2026Yijiang Li, Emon Dey, Zilinghan Li +3Federated Learning Aggregation High-Performance Computing
Mathematics and Computer Science Division, Argonne National Laboratory, Lemont, IL, USA · Data Science and Learning Division, Argonne National Laboratory, Lemont, IL, USA
Accelerating Sharded Data Parallelism at Scale with Federated Learning
Sep 17, 2026Gianluca Mittone, Marco AldinucciCommunication-Efficient Distributed Training High-Performance Computing
University of Turin, Turin, Italy
FedACT: Concurrent Federated Intelligence across Heterogeneous Data Sources
Mar 11, 2026Md Sirajul Islam, Isabelle G Chapman, N I Md Ashafuddula +4Online Federated Learning Federated Learning
School of Computing and Informatics, University of Louisiana at Lafayette, USA · Department of Computer and Information Sciences, University of Delaware, USA · Department of Computer Science, University of Illinois at Urbana-Champaign, USA

May 4, 2026J/K move · Enter open · S save

FedQueue: Queue-Aware Federated Learning for Cross-Facility HPC Training

Yijiang Li, Emon Dey, Zilinghan Li, Krishnan Raghavan+2

Mathematics and Computer Science Division, Argonne National Laboratory, Lemont, IL, USA · Data Science and Learning Division, Argonne National Laboratory, Lemont, IL, USA

Federated Learning Aggregation High-Performance Computing Asynchronous Federated Learning Federated Learning Heterogeneous Federated Learning

Open paper AskPDF arXiv

Federated learning (FL) across multiple HPC facilities faces stochastic admission delays from batch schedulers that dominate wall-clock time. Synchronous FL suffers from severe stragglers, while asynchronous FL accumulates stale updates when queues spike. We propose FedQueue, a queue-aware FL protocol that incorporates scheduler delays directly into training and aggregation, which (i) predicts per-facility queue delays online to budget local work, (ii) applies cutoff-based admission that buffers late arrivals to bound staleness, and (iii) performs staleness-aware aggregation to stabilize heterogeneous local workloads. We prove the convergence for non-convex objectives at rate

\mathcal{O}(1/\sqrt{R})

under bounded staleness, and show that the admission controls yield bounded staleness with high probability under queue-prediction error. Real-world cross-facility deployment of FedQueue shows 20.5% improvement over baseline algorithms. Controlled queue simulations demonstrate robust improvement over the baselines; in particular, up to 60% reduction in time to reach a target accuracy level under high queue variance and non-IID partitions.

Abstract

Federated Learning

Explore similar work

FedQueue: Queue-Aware Federated Learning for Cross-Facility HPC Training

Accelerating Sharded Data Parallelism at Scale with Federated Learning

FedACT: Concurrent Federated Intelligence across Heterogeneous Data Sources