1
The following paper is the accepted version of Springer copyrighted material
arXiv:2609.20359v1 [cs.DC] 17 Sep 2026
Gianluca Mittone, & Marco Aldinucci (2026). Accelerating Sharded Data Parallelism at Scale with Federated Learning. In Euro-Par 2026: Parallel Processing - 32nd European Conference on Parallel and Distributed Processing, Pisa, Italy, August 24-28, 2026, Proceedings, Part II (pp. 437–451). Springer. presented at the EuroPar’26 conference in Pisa, Italy, and awarded with the Best Paper Award. DOI: https://doi.org/10.1007/978-3-032-35251-4_30
Accelerating Sharded Data Parallelism at Scale with Federated Learning Gianluca Mittone [0000-0002-1887-6911] and Marco Aldinucci[0000-0001-8788-0829] University of Turin, Turin, Italy {gianluca.mittone, marco.aldinucci}@unito.it https://alpha.di.unito.it
Abstract. The symbiotic scaling of artificial intelligence models and high-performance computing systems continually creates algorithmic challenges in their convergence. Foundation models (FMs) are a crucial example, requiring months-long training on thousands of cutting-edge GPUs. Sharded data parallelism (DP) is the dominant strategy to accelerate such computations by splitting data and models across multiple GPUs. However, it incurs prohibitive communication overhead when deployed at scale, particularly on multi-tier interconnects with heterogeneous performance. Inspired by the efficient communication principles of federated learning (FL), this work introduces two hybrid algorithms—FL+FSDP and FL+HSDP —interleaving sharded DP with FedAvg-style aggregations. Such approaches decouple large DP deployments into smaller, loosely-coupled federation groups, requiring minimal inter-group traffic while keeping the global batch size bounded by the groups’ size. Formal analysis of communication costs and experimental validation prove their scalability and flexibility. A Llama3.1 8B pre-training on 512 A100 GPUs shows that, under identical hyperparameters, FL+FSDP and FL+HSDP achieve up to 8.04× faster data processing and 4.48× lower evaluation perplexity than their counterparts, demonstrating superior computational efficiency and improved model quality. These properties stem from reduced communication overhead and the bounded growth of the global batch size relative to the federation group size. Keywords: Distributed Training · Federated Learning · Communication Performance · Data Parallelism · FSDP · HSDP · HPC · LLM
1
Introduction
Contemporary large-scale clusters and High-Performance Computing (HPC) systems are increasingly being designed, built, and operated with artificial intelligence (AI) workloads as their primary target. This shift is reflected in the growing interest in low-precision hardware [4], AI-specific I/O patterns [8], and even network fabrics tuned for AI traffic [5]. At the same time, the hardware demands of AI workloads are expanding dramatically, with model sizes, training datasets, and parallelism scales all growing by orders of magnitude [17].
Accelerating Sharded Data Parallelism at Scale with Federated Learning
3
Foundation Models (FMs) exemplify this trend, with Large Language Models (LLMs) being their most widely known representatives [22,10]. Training such models can saturate the world’s most powerful supercomputers thanks to finely designed training pipelines that efficiently distribute computation across thousands of GPUs [23,14,19]. One of the most widespread approaches to do so is data parallelism (DP). This technique allows multiple copies of the same model to be trained in parallel on different data batches while keeping their gradients strictly synchronised via per-batch collective communication [9]. The current state of the art in this methodology is sharded DP, which enables training even very large FMs (i.e., tens of billions of parameters) that do not fit on a single GPU. This is achieved by "sharding" the FM parameters across many GPUs and re-collecting them when needed, thereby further increasing communication overheads linearly with the number of model replicas [23]. However, interconnect performance often becomes the limiting factor of such an approach at scale. While intra-node GPU-to-GPU links provide T B/s bandwidth (e.g., NVLink 5.0, ≈1.8T B/s), inter-node links are typically one to two orders of magnitude slower (e.g., InfiniBand XDR, ≈200GB/s per link). Moreover, modern datacenters are frequently organised in hierarchical topologies—such as a two-level fat-tree or the Dragonfly+ architecture [20]—where communication across different node groups incurs additional non-uniform network noise, sometimes reducing effective bandwidth by up to 50% [2]. Federated Learning (FL) [11] offers a complementary paradigm to solve such a problem: it enables DP training between loosely coupled devices through sporadic communications, the frequency of which can be customised. Additionally, it does not assume raw data sharing across model replicas, even though such sharing could affect model quality. Recent literature has explored FL for large-scale FM training, demonstrating that it mitigates communication overhead while preserving, or even improving, model quality [15,3]. However, no prior work addresses how FL can be integrated into current sharded DP techniques to enhance their performance. Such coupling introduces an extra degree of scalability, mitigating the constraints imposed by the underlying network infrastructure and improving computational and learning performance. This work introduces FL+FSDP and FL+HSDP, two FL-augmented sharded DP algorithms (see Figure 1) improving FMs’ training and learning performance at scale on modern HPC infrastructures. The contribution is threefold: – Hybrid communication scheme–a design that interoperates two flavours of sharded DP with FL-style aggregation, yielding higher scalability and reduced inter-node traffic at scale, better exploiting heterogeneous interconnection performance. – Theoretical analysis–a formal evaluation of the communication costs of fully and hybrid sharded DP and the relative FL-enriched proposed versions. – Empirical validation–extensive experiments on a top-10 Top500 HPC system demonstrating the proposed methods’ scalability, communication efficiency, and convergence in a large-scale setting (128 nodes - 512 GPUs).
G. Mittone, M. Aldinucci
Standard
FSDP
FL + sharded DP
<latexit sha1_base64="h29sDnsdvsbzPj0XkdMnTgoF8CE=">AAACDHicbVA5TgNBEJzlxlzmCJBIRthIJFi7BIYMIxJCcxiQbGP1jhs88uyhmV4EWi0BDyDmAaQQkSFSeAMBCX8gZ20IAFNRqapb3VVuqKQh2361+voHBoeGR0YzY+MTk1PZ6ZkDE0RaYEUEKtBHLhhU0scKSVJ4FGoEz1V46La3Ov7hGWojA3+fLkKse3DqyxMpgFKpkZ3P1wjPKd5rgW4mjdhJjuPdFSfJN7I5u2B3wXuJ801yGx+PNy/vV5vlRvat1gxE5KFPQoExVccOqR6DJikUJplaZDAE0YZTrKbUBw9NPe4mSPhSZIACHqLmUvGuiD83YvCMufDcdNIDapm/Xkf8z6tGdLJej6UfRoS+6BwiqbB7yAgt02qQN6VGIuh8jlz6XIAGItSSgxCpGKVdZdI+nL/pe8nBasEpFoo7Tq40x74wwhbYIltmDltjJbbNyqzCBLtkt+yO3VvX1oP1aD19jfZZ3zuz7Bes50+wrJ/O</latexit>
1 ShardR 0
1 ShardR 1
AllGather ReduceScatter
<latexit sha1_base64="vgtybdvfxXqMuJhPDDLk+himHw0=">AAACDnicbVC7SgNREL0b3/EVH5U2FxPBxrBroXYGbCx9RQNJDLM3E73k7oN7Z0VZAna2gh8gWNhqZSe2/oKFvyJuHoUmnmI4nDPDzBw3VNKQbX9aqaHhkdGx8Yn05NT0zGxmbv7EBJEWWBSBCnTJBYNK+lgkSQpLoUbwXIWnbnO37Z9eojYy8I/pOsSqB+e+bEgBlEi1zFKuQnhF8dEF6HqrFh+tO62z+DCpuVoma+ftDvggcXoku/N9//FQeLzdr2W+KvVARB76JBQYU3bskKoxaJJCYStdiQyGIJpwjuWE+uChqcadH1p8NTJAAQ9Rc6l4R8TfEzF4xlx7btLpAV2Yfq8t/ueVI2psV2PphxGhL9qLSCrsLDJCyyQc5HWpkQjalyOXPheggQi15CBEIkZJWukkD6f/+0FyspF3NvObB062sMi6GGfLbIWtMYdtsQLbY/usyAS7YU/smb1Yd9ar9Wa9d1tTVm9mgf2B9fEDYAygDw==</latexit>
1 ShardR S 1
Legend Sharding Group
Processes across wich a single model replica is sharded
Replication Group
Processes across which the same model shard is replicated
Federation Group
Processes across which the sharded DP configuration is replicated
1 ShardR 0
0
<latexit sha1_base64="MwEId7hjcGGLrkzOtWhX0NJuLUo=">AAACEHicbVDLSgNBEJyN7/iKj5seRqPgJWHXQ8xBMCCIx4hGhSSG3kknDs4+mOkVZdmLn+BXeNWTN/HqHyj4F36Am+hBjQUNRVU33V1uqKQh2361MkPDI6Nj4xPZyanpmdnc3PyxCSItsCYCFehTFwwq6WONJCk8DTWC5yo8cS92e/7JJWojA/+IrkNsetD1ZUcKoFRq5ZbXzuK9gpM0CK8oPjwH3U5asZ2cpbXWyuXtot0HHyTON8nvfBTc8tv2SrWVe2+0AxF56JNQYEzdsUNqxqBJCoVJthEZDEFcQBfrKfXBQ9OM+18kfD0yQAEPUXOpeF/EnxMxeMZce27a6QGdm79eT/zPq0fUKTdj6YcRoS96i0gq7C8yQss0HuRtqZEIepcjlz4XoIEIteQgRCpGaV7ZNA/n7/eD5Hiz6JSKpQMnX1lkXxhnS2yVbTCHbbEK22dVVmOC3bA7ds8erFvr0Xqynr9aM9b3zAL7BevlEz9pn60=</latexit>
1
Shard00
ReduceScatter
F
<latexit sha1_base64="QQF3FNXPmbyCXLkojjrsXvwo9v8=">AAACEHicbVDLSgNBEJyN7/iKj5seRqPgJWHXQ8xBMCCIx4hGhSSG3kknDs4+mOkVZdmLn+BXeNWTN/HqHyj4F36Am+hBjXUqqrrprnJDJQ3Z9quVGRoeGR0bn8hOTk3PzObm5o9NEGmBNRGoQJ+6YFBJH2skSeFpqBE8V+GJe7Hb808uURsZ+Ed0HWLTg64vO1IApVIrt7x2Fu8VnKRBeEXx4TnodtKK7eQsdpK1Vi5vF+0++CBxvkl+56Pglt+2V6qt3HujHYjIQ5+EAmPqjh1SMwZNUihMso3IYAjiArpYT6kPHppm3E+R8PXIAAU8RM2l4n0Rf27E4Blz7bnppAd0bv56PfE/rx5Rp9yMpR9GhL7oHSKpsH/ICC3TepC3pUYi6H2OXPpcgAYi1JKDEKkYpX1l0z6cv+kHyfFm0SkVSwdOvrLIvjDOltgq22AO22IVts+qrMYEu2F37J49WLfWo/VkPX+NZqzvnQX2C9bLJ0D6n64=</latexit>
1
1
Shard10
<latexit sha1_base64="fR3VigR7T+0kARk3eJEyhJHEkoo=">AAACEnicbVDJSgNBEO2Je9zichOkNQpeEmY8xBwEA4J4jEtUSGKo6VRiY89Cd40ow9z8BL/Cq568iVd/QMG/8AOcRA9u7/R4r4qq99xQSUO2/WJlBgaHhkdGx7LjE5NT07mZ2SMTRFpgTQQq0CcuGFTSxxpJUngSagTPVXjsnm/3/OML1EYG/iFdhdj0oOvLjhRAqdTKLa6cxjsFJ2kQXlJ8cAa6nbRiOzmN91N1pZXL20W7D/6XOF8kv/VecMuvm0vVVu6t0Q5E5KFPQoExdccOqRmDJikUJtlGZDAEcQ5drKfUBw9NM+7nSPhqZIACHqLmUvG+iN83YvCMufLcdNIDOjO/vZ74n1ePqFNuxtIPI0Jf9A6RVNg/ZISWaUHI21IjEfQ+Ry59LkADEWrJQYhUjNLGsmkfzu/0f8nRetEpFUt7Tr4yzz4xyhbYMltjDttgFbbLqqzGBLtmt+yO3Vs31oP1aD19jmasr5059gPW8wdmZaBB</latexit>
F
1
1 ShardR 0
Replication Group 0
F
1 ShardR S 1
<latexit sha1_base64="nrbc6YB2iw/q4q51eXnziWw0NBw=">AAACEnicbVC7SgNREL2rRmN8xUcnyMUo2Bh2LdTOgI1lfEQDeTF7M+rFuw/unRXDsp3gByj4D7Za2YmtP2Dhb9gJbqKFGk8xHM6ZYWaOGyppyLZfrYHBoczwSHY0NzY+MTmVn545NEGkBVZEoAJddcGgkj5WSJLCaqgRPFfhkXu23fWPzlEbGfgH1Amx4cGJL4+lAEqlVn5hqRnbSZ3wguL9U9DtpBXvrzpJM95L61IrX7CLdg+8nzjfpLD1kbm6frt9L7fyb/V2ICIPfRIKjKk5dkiNGDRJoTDJ1SODIYgzOMFaSn3w0DTi3h8JX44MUMBD1Fwq3hPx50QMnjEdz007PaBT89friv95tYiONxux9MOI0BfdRSQV9hYZoWUaEPK21EgE3cuRS58L0ECEWnIQIhWjNLFcmofz9/t+crhWdNaL67tOoTTHvpBl82yRrTCHbbAS22FlVmGCXbI7ds8erBvr0Xqynr9aB6zvmVn2C9bLJ64goeI=</latexit>
AllReduce
Shard11
ReduceScatter
1
1
AllGather ReduceScatter
<latexit sha1_base64="yNJc3fV/kjPstKTKqWygjRdwGcA=">AAACEnicbVDJSgNBEO2Je9zichOkNQpeEmY8xBwEA4J4jEtUSGKo6VRiY89Cd40ow9z8BL/Cq568iVd/QMG/8AOcRA9u7/R4r4qq99xQSUO2/WJlBgaHhkdGx7LjE5NT07mZ2SMTRFpgTQQq0CcuGFTSxxpJUngSagTPVXjsnm/3/OML1EYG/iFdhdj0oOvLjhRAqdTKLa6cxjsFJ2kQXlJ8cAa6nbRiJzmN91N1pZXL20W7D/6XOF8kv/VecMuvm0vVVu6t0Q5E5KFPQoExdccOqRmDJikUJtlGZDAEcQ5drKfUBw9NM+7nSPhqZIACHqLmUvG+iN83YvCMufLcdNIDOjO/vZ74n1ePqFNuxtIPI0Jf9A6RVNg/ZISWaUHI21IjEfQ+Ry59LkADEWrJQYhUjNLGsmkfzu/0f8nRetEpFUt7Tr4yzz4xyhbYMltjDttgFbbLqqzGBLtmt+yO3Vs31oP1aD19jmasr5059gPW8wdn/KBC</latexit>
F
0
<latexit sha1_base64="F32vninDZ+uoRXZzEkqxahu3tZ8=">AAACEHicbVC7TsNAEDyHd3iFRwUUJwISDZFNAXQg0VAGhYRISYjWl4Wccn7obo1Alhs+gYZfSAsVHaLlC6CAD6HCCRQQmGo0s6vdGTdU0pBtv1qZoeGR0bHxiezk1PTMbG5uvmKCSAssi0AFuuqCQSV9LJMkhdVQI3iuwhO3c9DzTy5QGxn4x3QVYsODc1+eSQGUSs3cytppbCd1wkuKS23QraQZlzad5DR2krVmLm8X7D74X+J8k/zex8vSe/dWFJu5t3orEJGHPgkFxtQcO6RGDJqkUJhk65HBEEQHzrGWUh88NI24nyLh65EBCniImkvF+yL+3IjBM+bKc9NJD6htBr2e+J9Xi+hstxFLP4wIfdE7RFJh/5ARWqb1IG9JjUTQ+xy59LkADUSoJQchUjFK+8qmfTiD6f+SylbB2S5sHzn5/UX2hXG2zFbZBnPYDttnh6zIykywa9Zld+zeurEerEfr6Ws0Y33vLLBfsJ4/AWeroTU=</latexit>
AllGather
Shard01
<latexit sha1_base64="eWcMK9cqMz0WeYVB4zKTov+TuIc=">AAACEHicbVDLSgNBEJyN7/iKj5seRqPgJWHXQ8xBMCCIx4hGhSSG3kknDs4+mOkVZdmLn+BXeNWTN/HqHyj4F36Am+hBjQUNRVU33V1uqKQh2361MkPDI6Nj4xPZyanpmdnc3PyxCSItsCYCFehTFwwq6WONJCk8DTWC5yo8cS92e/7JJWojA/+IrkNsetD1ZUcKoFRq5ZbXzuK9gpM0CK8oPjwH3U5asZOcpbXWyuXtot0HHyTON8nvfBTc8tv2SrWVe2+0AxF56JNQYEzdsUNqxqBJCoVJthEZDEFcQBfrKfXBQ9OM+18kfD0yQAEPUXOpeF/EnxMxeMZce27a6QGdm79eT/zPq0fUKTdj6YcRoS96i0gq7C8yQss0HuRtqZEIepcjlz4XoIEIteQgRCpGaV7ZNA/n7/eD5Hiz6JSKpQMnX1lkXxhnS2yVbTCHbbEK22dVVmOC3bA7ds8erFvr0Xqynr9aM9b3zAL7BevlE0KPn68=</latexit>
F
Shard1S 1
AllReduce
1 ShardR 1
<latexit sha1_base64="lljc3MCPG6bnrDHgX/doJhuRJbc=">AAACEHicbVDLSgNBEJyN7/iKj5seRqPgJWHXQ8xBMCCIx4hGhSSG3kknDs4+mOkVZdmLn+BXeNWTN/HqHyj4F36Am+hBjXUqqrrprnJDJQ3Z9quVGRoeGR0bn8hOTk3PzObm5o9NEGmBNRGoQJ+6YFBJH2skSeFpqBE8V+GJe7Hb808uURsZ+Ed0HWLTg64vO1IApVIrt7x2Fu8VnKRBeEXx4TnodtKKneQstpO1Vi5vF+0++CBxvkl+56Pglt+2V6qt3HujHYjIQ5+EAmPqjh1SMwZNUihMso3IYAjiArpYT6kPHppm3E+R8PXIAAU8RM2l4n0Rf27E4Blz7bnppAd0bv56PfE/rx5Rp9yMpR9GhL7oHSKpsH/ICC3TepC3pUYi6H2OXPpcgAYi1JKDEKkYpX1l0z6cv+kHyfFm0SkVSwdOvrLIvjDOltgq22AO22IVts+qrMYEu2F37J49WLfWo/VkPX+NZqzvnQX2C9bLJ0D+n64=</latexit>
F
0
1 ShardR 1
Replication Group 1
Broadcast
AllReduce
ReduceScatter
<latexit sha1_base64="wBVRBxwqrCY6aMO2kfqGBhKTw4Y=">AAACEHicbVA5TgNBEJw1t7nMkUEwYCOR2NolAAIkkEgIuYwt2cbqHbfNiNlDM70ItNqEJ/AKUojIECk/AIlf8ADWxgFgKipVdau7yg2VNGTbb1ZmaHhkdGx8Ijs5NT0zm5ubPzNBpAWWRaACXXXBoJI+lkmSwmqoETxXYcW93O/6lSvURgb+Kd2E2PCg48u2FECp1MwtF85jO6kTXlN8cgG6lTRjJzmPj4tOUmjm8nbJ7oEPEqdP8rufRXf7fWflsJn7qLcCEXnok1BgTM2xQ2rEoEkKhUm2HhkMQVxCB2sp9cFD04h7KRK+FhmggIeouVS8J+LPjRg8Y248N530gC7MX68r/ufVImpvN2LphxGhL7qHSCrsHTJCy7Qe5C2pkQi6nyOXPheggQi15CBEKkZpX9m0D+dv+kFytlFyNkubR05+b5F9Y5wtsVW2zhy2xfbYATtkZSbYLbtnD+zRurOerGfr5Xs0Y/V3FtgvWK9fWwSfug==</latexit>
0
Shard0S 1
Broadcast
AllReduce
Broadcast
<latexit sha1_base64="T6e5gXf7+P22K8ag5dQDnciotaI=">AAACEHicbVA5TgNBEJw1t7nMkUEwYCOR2NolAAIkkEgIuYwt2cbqHbfNiNlDM70ItNqEJ/AKUojIECk/AIlf8ADWxgFgKqqu6lZ3lxsqaci236zM0PDI6Nj4RHZyanpmNjc3f2aCSAssi0AFuuqCQSV9LJMkhdVQI3iuwop7ud/1K1eojQz8U7oJseFBx5dtKYBSqZlbLpzHdlInvKb45AJ0K2mm9Xl8XHSSQjOXt0t2D3yQOH2S3/0sutvvOyuHzdxHvRWIyEOfhAJjao4dUiMGTVIoTLL1yGAI4hI6WEupDx6aRtz7IuFrkQEKeIiaS8V7Iv6ciMEz5sZz004P6ML89brif14tovZ2I5Z+GBH6oruIpMLeIiO0TONB3pIaiaB7OXLpcwEaiFBLDkKkYpTmlU3zcP5+P0jONkrOZmnzyMnvLbJvjLMltsrWmcO22B47YIeszAS7ZffsgT1ad9aT9Wy9fLdmrP7MAvsF6/ULWW2fuQ==</latexit>
AllGather
Shard11
0
0
AllGather
<latexit sha1_base64="w2hMgi/YaqC8ZO43v7Mlx8hNBOI=">AAACEnicbVC7SgNBFJ31GeMrPioFGYyCjWHXItopCGKpxEQhieHu5GqGzD6YuSuGZTs/QfAXxFYrO7H1A7TQD7FyEy18nepwzr3ce44bKmnItl+svv6BwaHhzEh2dGx8YjI3NV0xQaQFlkWgAn3kgkElfSyTJIVHoUbwXIWHbnu76x+eoTYy8A+oE2Ldg1NfnkgBlEqN3MLScbyz6iQ1wnOKSy3QzaQRl1LlOLaTpUYubxfsHvhf4nyR/Ob789zb9ZXYa+Rea81ARB76JBQYU3XskOoxaJJCYZKtRQZDEG04xWpKffDQ1ONejoQvRwYo4CFqLhXvifh9IwbPmI7nppMeUMv89rrif141opONeiz9MCL0RfcQSYW9Q0ZomRaEvCk1EkH3c+TS5wI0EKGWHIRIxShtLJv24fxO/5dU1gpOsVDcd/Jbs+wTGTbPFtkKc9g622K7bI+VmWAX7Ibdsjvr0rq3HqzHz9E+62tnhv2A9fQBcxKhvA==</latexit>
F
1
Shard0S 1
F
1
Shard1S 1
F
1
1 ShardR S 1
<latexit sha1_base64="i4jDiDlqGU4hHFLqnGYq8FZWI5U=">AAACEnicbVC7SgNBFJ31GeMrPioFGYyCjWHXItopCGKpxEQhieHu5GqGzD6YuSuGZTs/QfAXxFYrO7H1A7TQD7FyEy18nepwzr3ce44bKmnItl+svv6BwaHhzEh2dGx8YjI3NV0xQaQFlkWgAn3kgkElfSyTJIVHoUbwXIWHbnu76x+eoTYy8A+oE2Ldg1NfnkgBlEqN3MLScbyz6iQ1wnOKSy3QzaQRl1LlOHaSpUYubxfsHvhf4nyR/Ob789zb9ZXYa+Rea81ARB76JBQYU3XskOoxaJJCYZKtRQZDEG04xWpKffDQ1ONejoQvRwYo4CFqLhXvifh9IwbPmI7nppMeUMv89rrif141opONeiz9MCL0RfcQSYW9Q0ZomRaEvCk1EkH3c+TS5wI0EKGWHIRIxShtLJv24fxO/5dU1gpOsVDcd/Jbs+wTGTbPFtkKc9g622K7bI+VmWAX7Ibdsjvr0rq3HqzHz9E+62tnhv2A9fQBdKOhvQ==</latexit>
AllReduce
AllReduce <latexit sha1_base64="JosrybJj9hgvXN1Mqq8X23WgX/Y=">AAACDHicbVA5TgNBEJzlxlzmCJBIRthIJFi7BIYMIxJCcxiQbGP1jhs88uyhmV4EWi0BDyDmAaQQkSFSeAMBCX8gZ20IAFNRqapb3VVuqKQh2361+voHBoeGR0YzY+MTk1PZ6ZkDE0RaYEUEKtBHLhhU0scKSVJ4FGoEz1V46La3Ov7hGWojA3+fLkKse3DqyxMpgFKpkZ3P1wjPKd5rgW4mjdhOjuPdFSfJN7I5u2B3wXuJ801yGx+PNy/vV5vlRvat1gxE5KFPQoExVccOqR6DJikUJplaZDAE0YZTrKbUBw9NPe4mSPhSZIACHqLmUvGuiD83YvCMufDcdNIDapm/Xkf8z6tGdLJej6UfRoS+6BwiqbB7yAgt02qQN6VGIuh8jlz6XIAGItSSgxCpGKVdZdI+nL/pe8nBasEpFoo7Tq40x74wwhbYIltmDltjJbbNyqzCBLtkt+yO3VvX1oP1aD19jfZZ3zuz7Bes50+vFZ/N</latexit>
Shard10
<latexit sha1_base64="dhPjYkUNJn3kdE+akO552XplpMw=">AAACEHicbVC7TsNAEDzzDOEVHhVQnEiQaIhsikBHJBpKEASQkhCtLwuccn7obo2ILDd8Ag2/QAsVHaLlC6CAD6HCTih4TTWa2dXujBsqaci2X62BwaHhkdHcWH58YnJqujAze2iCSAusiUAF+tgFg0r6WCNJCo9DjeC5Co/cznbmH12gNjLwD6gbYtODM1+eSgGUSq3CUukktpMG4SXF++eg20kr3l9zkkwttQpFu2z3wP8S54sUtz5eFt5vb8Ruq/DWaAci8tAnocCYumOH1IxBkxQKk3wjMhiC6MAZ1lPqg4emGfdSJHwlMkABD1FzqXhPxO8bMXjGdD03nfSAzs1vLxP/8+oRnW42Y+mHEaEvskMkFfYOGaFlWg/yttRIBNnnyKXPBWggQi05CJGKUdpXPu3D+Z3+LzlcLzuVcmXPKVbnWR85tsiW2Spz2Aarsh22y2pMsCt2y+7YvXVtPViP1lN/dMD62pljP2A9fwJmGqE0</latexit>
AllGather ReduceScatter
Broadcast
<latexit sha1_base64="/5DFCB+5JVo+6dlMpiG9+cdGp7o=">AAACDHicbVC7SgNBFJ31GeMrPgrFZjAKNoZdi2hnwMZS0aiQxHB3ctUhsw9m7ophWcEfsLe31cpObC38Awu/QxvBTUyh0VMdzrmXe89xQyUN2far1dc/MDg0nBnJjo6NT0zmpqYPTBBpgWURqEAfuWBQSR/LJEnhUagRPFfhodvcavuH56iNDPx9aoVY8+DUlydSAKVSPTe3VCW8oHjvDHQjqcd7q05yHDvJUj2Xtwt2B/wvcbokv/k5//7xcnWzU8+9VRuBiDz0SSgwpuLYIdVi0CSFwiRbjQyGIJpwipWU+uChqcWdBAlfjgxQwEPUXCreEfHnRgyeMS3PTSc9oDPT67XF/7xKRCcbtVj6YUToi/Yhkgo7h4zQMq0GeUNqJIL258ilzwVoIEItOQiRilHaVTbtw+lN/5ccrBWcYqG46+RLs+wbGbbAFtkKc9g6K7FttsPKTLBLdsvu2L11bT1Yj9bT92if1d2ZYb9gPX8B+PWgBA==</latexit>
Shard1S 1
<latexit sha1_base64="ifuHYTSwmKyl/2XiLTV5ihqCcBU=">AAACDnicbVDLSgNBEJyNrxhf8XHSy2gieDHselAPgoIXj4pGA0kMvZPWDJl9MNMrhmXBT/ArvOrJm3j1FxT8Cz/ATeJBEwsaiqpuurvcUElDtv1uZUZGx8YnspO5qemZ2bn8/MK5CSItsCwCFeiKCwaV9LFMkhRWQo3guQov3PZh17+4QW1k4J9RJ8S6B9e+vJICKJUa+eXiZWwnNcJbik9boJtJI3aSy7SKjXzBLtk98GHi/JDC/temu/uxt3rcyH/WmoGIPPRJKDCm6tgh1WPQJIXCJFeLDIYg2nCN1ZT64KGpx70fEr4eGaCAh6i5VLwn4u+JGDxjOp6bdnpALTPodcX/vGpEV7v1WPphROiL7iKSCnuLjNAyDQd5U2okgu7lyKXPBWggQi05CJGKUZpWLs3DGfx+mJxvlZzt0vaJUzhYYn1k2QpbYxvMYTvsgB2xY1Zmgt2xB/bInqx769l6sV77rRnrZ2aR/YH19g02h58n</latexit>
0
AllGather ReduceScatter
AllReduce
AllGather ReduceScatter
AllReduce
Shard11
<latexit sha1_base64="zPm8okJj89nc0ccp4eOl8xE0/dU=">AAACDnicbVC7SgNREL0bXzG+4qPS5moi2Bh2LdRCULCxVDQaSGKYvRnNJXcf3DsrhmXBT/ArbLWyE1t/QcG/8APcJBaaeKoz58wwM8cNlTRk2+9WZmR0bHwiO5mbmp6ZncvPL5ybINICyyJQga64YFBJH8skSWEl1Aieq/DCbR92/Ysb1EYG/hl1Qqx7cO3LKymAUqmRXy5exnZSI7yl+LQFupk00voydpJiI1+wS3YPfJg4P6Sw/7Xp7n7srR438p+1ZiAiD30SCoypOnZI9Rg0SaEwydUigyGINlxjNaU+eGjqce+HhK9HBijgIWouFe+J+HsiBs+YjuemnR5Qywx6XfE/rxrR1W49ln4YEfqiu4ikwt4iI7RMw0HelBqJoHs5culzARqIUEsOQqRilKaVS/NwBr8fJudbJWe7tH3iFA6WWB9ZtsLW2AZz2A47YEfsmJWZYHfsgT2yJ+veerZerNd+a8b6mVlkf2C9fQM08p8m</latexit>
AllReduce
<latexit sha1_base64="B5ZX/vPU9FDSSLHZbCtHihtAErQ=">AAACCnicbVDLTgJBEJz1ifjCR7x4mQgmnsiuB/UmxotHjIIkgKR3aGHi7CMzvUay8gd+hVc9aWJivPoRevAv/AAX8KBgJZ1UqrrT3eWGShqy7Q9rbHxicmo6NZOenZtfWMwsLZdNEGmBJRGoQFdcMKikjyWSpLASagTPVXjmXh72/LMr1EYG/il1Qqx70PLlhRRAidTIrOVqhNcUn7RBN7uN2OmeJ5VrZLJ23u6DjxLnh2T3vw6etp7fg2Ij81lrBiLy0CehwJiqY4dUj0GTFAq76VpkMARxCS2sJtQHD0097t/f5ZuRAQp4iJpLxfsi/p6IwTOm47lJpwfUNsNeT/zPq0Z0sVePpR9GhL7oLSKpsL/ICC2TYJA3pUYi6F2OXPpcgAYi1JKDEIkYJUmlkzyc4e9HSXk77+zkd46dbGGVDZBi62yDbTGH7bICO2JFVmKC3bA7ds8erFvr0XqxXgetY9bPzAr7A+vtG6XBnpQ=</latexit>
Shard10
<latexit sha1_base64="zXuvF1LP+sMLUCkk7gcQ+CIQa3A=">AAACDHicbVC7SgNBFJ31GeMrPgrFZjAKNoZdi2hnwMZS0aiQxHB3ctUhsw9m7ophWcEfsLe31cpObC38Awu/QxvBTUyh0VMdzrmXe89xQyUN2far1dc/MDg0nBnJjo6NT0zmpqYPTBBpgWURqEAfuWBQSR/LJEnhUagRPFfhodvcavuH56iNDPx9aoVY8+DUlydSAKVSPTe3VCW8oHjvDHQjqcd7q05yHNvJUj2Xtwt2B/wvcbokv/k5//7xcnWzU8+9VRuBiDz0SSgwpuLYIdVi0CSFwiRbjQyGIJpwipWU+uChqcWdBAlfjgxQwEPUXCreEfHnRgyeMS3PTSc9oDPT67XF/7xKRCcbtVj6YUToi/Yhkgo7h4zQMq0GeUNqJIL258ilzwVoIEItOQiRilHaVTbtw+lN/5ccrBWcYqG46+RLs+wbGbbAFtkKc9g6K7FttsPKTLBLdsvu2L11bT1Yj9bT92if1d2ZYb9gPX8B92SgAw==</latexit>
Shard0S 1
Shard01
Replication Group S-1
Broadcast
<latexit sha1_base64="idZTbTPvM1BLZ7UV1APPaSKsc3k=">AAACCnicbVA5TgNBEJzlNOYyh0hIRhgkImuXwJBhREJoBDZItrF6x4098uyhmV6EtewPeAUpRCAhIVIeAQG/4AGsbQKuikpV3erqckMlDdn2mzUyOjY+MZmZyk7PzM7N5xYWqyaItMCKCFSgT10wqKSPFZKk8DTUCJ6r8MTt7vf9kwvURgb+MfVCbHjQ9uW5FECp1MytrNcJLyk+6oBuJc3YTs5iJ1lv5vJ2wR6A/yXOF8nvfuw9bD6+BuVm7r3eCkTkoU9CgTE1xw6pEYMmKRQm2XpkMATRhTbWUuqDh6YRD/InfCMyQAEPUXOp+EDE7xsxeMb0PDed9IA65rfXF//zahGd7zRi6YcRoS/6h0gqHBwyQsu0GOQtqZEI+smRS58L0ECEWnIQIhWjtKls2ofz+/u/pLpVcIqF4qGTLy2zITJsla2xTeawbVZiB6zMKkywK3bDbtmddW3dW0/W83B0xPraWWI/YL18AqQsnpM=</latexit>
AllGather ReduceScatter
0
AllReduce
Sharding Group 1
Shard01
Replication Group S-1
AllReduce
<latexit sha1_base64="jKDLH/niI2vJytAYlPz77cj4w3A=">AAACCnicbVA5TgNBEJzlNOYyh0hIRhgkImuXwJBhREJoBDZItrF6x4098uyhmV6EtewPeAUpRCAhIVIeAQG/4AGsbQKuikpV3erqckMlDdn2mzUyOjY+MZmZyk7PzM7N5xYWqyaItMCKCFSgT10wqKSPFZKk8DTUCJ6r8MTt7vf9kwvURgb+MfVCbHjQ9uW5FECp1MytrNcJLyk+6oBuJc3YSc5iO1lv5vJ2wR6A/yXOF8nvfuw9bD6+BuVm7r3eCkTkoU9CgTE1xw6pEYMmKRQm2XpkMATRhTbWUuqDh6YRD/InfCMyQAEPUXOp+EDE7xsxeMb0PDed9IA65rfXF//zahGd7zRi6YcRoS/6h0gqHBwyQsu0GOQtqZEI+smRS58L0ECEWnIQIhWjtKls2ofz+/u/pLpVcIqF4qGTLy2zITJsla2xTeawbVZiB6zMKkywK3bDbtmddW3dW0/W83B0xPraWWI/YL18AqQwnpM=</latexit>
Shard00
Shard00
Broadcast
<latexit sha1_base64="hmj4v/fbZt1S6NFPyxAU/dMcpUM=">AAACCnicbVDJSgNBEO1xN25xwYuXxijkFGY8RG9GvHhUNImQxFDTqWiTnoXuGjGM+QO/wqueFATx6kfowb/wA5xJcnB7UPB4r4qqem6opCHbfrdGRsfGJyanpjMzs3PzC9nFpYoJIi2wLAIV6FMXDCrpY5kkKTwNNYLnKqy6nf3Ur16iNjLwT6gbYsODc1+2pQBKpGZ2daNOeEXx8QXoVq8Z272zpDaa2ZxdsPvgf4kzJLndz73H/NNbcNjMftRbgYg89EkoMKbm2CE1YtAkhcJeph4ZDEF04BxrCfXBQ9OI+/f3+GZkgAIeouZS8b6I3ydi8Izpem7S6QFdmN9eKv7n1SJq7zRi6YcRoS/SRSQV9hcZoWUSDPKW1EgE6eXIpc8FaCBCLTkIkYhRklQmycP5/f1fUtkqOMVC8cjJlVbYAFNsja2zPHPYNiuxA3bIykywa3bL7ti9dWM9WM/Wy6B1xBrOLLMfsF6/AKKbnpI=</latexit>
<latexit sha1_base64="gWr2hkhLLu6oz3rgKCwkYGy3h1Y=">AAACDnicbVC7TsNAEDzzDOEVHhU0BwGJhsimAAokItFQBkFIpCRE68sCJ84P3a0RkWWJT+AraKGiQ7T8Akj8BR+AnVBAYKrRzK52Z9xQSUO2/WYNDY+Mjo3nJvKTU9Mzs4W5+VMTRFpgVQQq0HUXDCrpY5UkKayHGsFzFdbcq4PMr12jNjLwT6gbYsuDC1+eSwGUSu3C0tpZbCdNwhuKjy9Bd5J27CSZttYuFO2S3QP/S5xvUtz/3HR33/dWKu3CR7MTiMhDn4QCYxqOHVIrBk1SKEzyzchgCOIKLrCRUh88NK24lyHh65EBCniImkvFeyL+3IjBM6bruemkB3RpBr1M/M9rRHS+24qlH0aEvsgOkVTYO2SElmk5yDtSIxFknyOXPheggQi15CBEKkZpW/m0D2cw/V9yulVytkvbR06xvMj6yLFltso2mMN2WJkdsgqrMsFu2T17YI/WnfVkPVsv/dEh63tngf2C9foFNPafJg==</latexit>
0
Sharding Group 1
Replication Group 1
Replication Group 1
<latexit sha1_base64="PSS5TnCYnZQJoRNgTFjCnijHS8A=">AAACDnicbVC7SgNREL3rM8ZXfFTaXE0EG8OuhVoICjaWisYEkhhmbybJJXcf3DsrhiXgJ/gVtlrZia2/oOBf+AHuRgs1nurMOTPMzHFDJQ3Z9qs1Mjo2PjGZmcpOz8zOzecWFi9MEGmBJRGoQFdcMKikjyWSpLASagTPVVh2u0epX75CbWTgn1MvxLoHbV+2pABKpEZupXAZ2/0a4TXFZx3QzX4jqVOt0Mjl7aI9AB8mzjfJH3xsuXtv+2snjdx7rRmIyEOfhAJjqo4dUj0GTVIo7GdrkcEQRBfaWE2oDx6aejz4oc83IgMU8BA1l4oPRPw5EYNnTM9zk04PqGP+eqn4n1eNqLVXj6UfRoS+SBeRVDhYZISWSTjIm1IjEaSXI5c+F6CBCLXkIEQiRkla2SQP5+/3w+Riu+jsFHdOnfzhMvtChq2ydbbJHLbLDtkxO2ElJtgNu2P37MG6tR6tJ+v5q3XE+p5ZYr9gvXwCM2GfJQ==</latexit>
AllReduce
Sharding Group 0
Replication Group 0
Replication Group 0
<latexit sha1_base64="yvsmwAutf1W5g9NPkoE7srAubmM=">AAACCHicbVC7TgJBFJ31ifgCobOZCCY2kl0LtJPExhKDgAkQcne46MTZR2buGnVD7P0KW63sjK2JH2HhV/gDLmCh4KlOzrk399zjhkoasu0Pa2Z2bn5hMbWUXl5ZXVvPZDcaJoi0wLoIVKDPXDCopI91kqTwLNQInquw6V4eDf3mFWojA/+UbkLseHDuy74UQInUzeSLbcJrimsXoHuDblzbdQbFbqZgl+wR+DRxfkjh8Cubzt0dvFe7mc92LxCRhz4JBca0HDukTgyapFA4SLcjgyGISzjHVkJ98NB04lH6Ad+ODFDAQ9RcKj4S8fdGDJ4xN56bTHpAF2bSG4r/ea2I+gedWPphROiL4SGSCkeHjNAyqQV5T2okgmFy5NLnAjQQoZYchEjEKOkpnfThTH4/TRp7JadcKp84hUqejZFim2yL7TCH7bMKO2ZVVmeC3bIH9sierHvr2XqxXsejM9bPTo79gfX2DbGMnAk=</latexit>
ShardS 1
Sharding Group R-1
AllGather ReduceScatter
Sharding Group 0
Shard1
Sharding Group 0
<latexit sha1_base64="PsRMbKtJSEnyG6tZeSk4HxyB9nk=">AAACBnicbVC7TgJBFJ1FRcQXop3NRDCxIrsWaKeJjSVGeSRAyN3hohNmH5m5azAber/CwkYrO2Nh4wf4A8bY+xkuYKHiqU7OuTf33OOGShqy7TcrNTM7l57PLGQXl5ZXVnNr+ZoJIi2wKgIV6IYLBpX0sUqSFDZCjeC5Cutu/2jk1y9RGxn4Z3QVYtuDc1/2pABKpE4uX2wRDig+vQDdHXZiZ1js5Ap2yR6DTxPnmxQOPl9v3+E5XenkPlrdQEQe+iQUGNN07JDaMWiSQuEw24oMhiD6cI7NhPrgoWnH4+xDvh0ZoICHqLlUfCziz40YPGOuPDeZ9IAuzF9vJP7nNSPq7bdj6YcRoS9Gh0gqHB8yQsukFORdqZEIRsmRS58L0ECEWnIQIhGjpKVs0ofz9/tpUtstOeVS+cQpHG6wCTJsk22xHeawPXbIjlmFVZlgA3bD7ti9dW09WI/W02Q0ZX3vrLNfsF6+APR4nSs=</latexit>
Shard0
Sharding Group 1
Sharding Group
Federation Group 0 <latexit sha1_base64="fm8X4c4ZdTYnMmQW9x/8Ial/f9s=">AAACBnicbVC7TgJBFJ3FF+JrRTubiWhiRXYt0E4SG0uNIiZAyN3hAhNnH5m5ayAber/CwkYrO2Nh4wf4A8bY+xkuSKHoqU7OuTf33ONFShpynDcrMzU9MzuXnc8tLC4tr9ir+XMTxlpgRYQq1BceGFQywApJUngRaQTfU1j1Lg+HfvUKtZFhcEb9CBs+dALZlgIolZp2fqtO2KPktAu6NWgmzmCraRecojMC/0vcMSkcfL7evsPz7HHT/qi3QhH7GJBQYEzNdSJqJKBJCoWDXD02GIG4hA7WUhqAj6aRjLIP+HZsgEIeoeZS8ZGIPzcS8I3p+1466QN1zaQ3FP/zajG19xuJDKKYMBDDQyQVjg4ZoWVaCvKW1EgEw+TIZcAFaCBCLTkIkYpx2lIu7cOd/P4vOd8tuqVi6cQtlNfZN7Jsg22yHeayPVZmR+yYVZhgPXbD7ti9dW09WI/W0/doxhrvrLFfsF6+APLnnSo=</latexit>
HSDP
Sharding Group R-1
Proposed
Sharding Group R-1
4
<latexit sha1_base64="wwPB/OTk0BojAL2XjtjUuiX9UqU=">AAACFHicbVC7SkNBEN3rM8ZXfHQ2i1GwMdxroXYKgljGRx6QxDB3M+ri3ge7c0W5pNUPEPQnbLWyE1t7C3/DTnATLdR4iuFwzgwzc/xYSUOu++r09Q8MDg1nRrKjY+MTk7mp6bKJEi2wJCIV6aoPBpUMsUSSFFZjjRD4Civ+6VbHr5yhNjIKD+gixkYAx6E8kgLISs0cXzhMt5e9dp3wnNL9E9CtdjPdt8phumfrQjOXdwtuF7yXeN8kv/ExeHX9dvtebObe6q1IJAGGJBQYU/PcmBopaJJCYTtbTwzGIE7hGGuWhhCgaaTdT9p8MTFAEY9Rc6l4V8SfEykExlwEvu0MgE7MX68j/ufVEjpab6QyjBPCUHQWkVTYXWSEljYi5C2pkQg6lyOXIReggQi15CCEFRObWdbm4f39vpeUVwreamF118tvzrIvZNgcm2dLzGNrbJPtsCIrMcEu2R27Zw/OjfPoPDnPX619zvfMDPsF5+UTvAiiag==</latexit>
Replication Group S-1
Federation Group F-1
Fig. 1. Standard sharded DP techniques (left) vs the proposed FL-augmented (right)
2
Background
Distributed Learning State-of-the-art distributed training techniques leverage several orthogonal parallelisation dimensions, such as Model (GPipe [6]), Tensor /Sequence (Megatron-LM [19]), experts [18] parallelism, but data parallelism (ZeRO [14]) remains the most common base. Standard DP implies training multiple replicas of the same model deployed on different computing devices simultaneously. Each model replica is fed different data and, thanks to per-batch gradient synchronisation, the replicas’ parameters evolve in perfect alignment. This also has the collateral effect of increasing the global batch size of training in proportion to the number of replicas, establishing a strong relation between the level of parallelism and the final model’s quality. However, DP requires that each model replica fit into each computing device’s memory, which is not a valid assumption for modern FMs. In such a case, sharded DP extends this technique, enabling model partitioning and reconstruction across multiple devices. Fully Sharded Data Parallelism (FSDP ) [23] (top-left of Figure 1) organises the processes into a one-dimensional mesh of size S. Each model’s layer is split equally over the mesh in shards of size M : from this, the definition fully sharded.
Accelerating Sharded Data Parallelism at Scale with Federated Learning
5
The processes sharing a sharded model replica are called sharding group. During the forward and backward passes, each layer must be reconstructed on each process to perform computation, implying an all-gather over the sharding group for each of the L layers. Similarly, during the backward pass, full gradients must be reconstructed for each layer. Thus, L reduce-scatters over the S gradient shards (also of size M ) are required to maintain replicas aligned during training. The communication cost of FSDP is formalised in Section 4.1. Hybrid Sharded Data Parallelism (HSDP ) [23] (bottom-left of Figure 1) extends FSDP by organizing the available processes into a two-dimensional R × S mesh: 1) R sharding groups, each constituting an independent FSDP instance, and 2) S replication groups, each connecting corresponding shards across different sharding groups. Each sharding group keeps internal model replicas synchronised through the FSDP communication schema. However, additional L all-reduces within each replication group are required to synchronise corresponding gradient shards among the replication dimension. Such all-reduces are mutually independent between replica groups, and can therefore be executed in parallel. The communication cost of HSDP is formalised in Section 4.1. Federated Learning Federated Learning (FL) [11] is a widely adopted paradigm for training AI models across loosely coupled devices that have access to local, possibly private, data. Unlike classic DP training, FL does not require frequent collective communications: synchronisation, called aggregation, occurs only every A training steps, termed rounds, with a frequency that can range from once every few local mini-batches to once every several local epochs, depending on the training constraints. Such a property makes FL very adaptable to the most diverse network conditions, while allowing periodic divergence among the model replicas. The standard FL aggregation strategy is Federated Averaging (FedAvg) [11], an iterative algorithm which computes and redistributes back a weighted average of all locally trained models at the beginning of each federated round.
3
Related Works
Recent related work also explores FL as a means to harness loosely coupled, even geographically dispersed, compute resources for FM training. In DiLoCo [3], Google pioneers an FL-inspired training pipeline for FMs with 60M–400M parameters on a modest cluster (up to eight NVIDIA A100 GPUs). The authors employ a FedOpt variant of FedAvg, using AdamW as the inner optimiser and Nesterov momentum for the outer aggregation step. DiLoCo demonstrate faster convergence and an order-of-magnitude reduction in communicated data compared with a purely centralised training baseline. On the other hand, Photon [15] extended this idea to a truly geo-distributed environment. FMs ranging from 125M to 7B parameters are trained across up to five sites on different continents (each equipped with up to two NVIDIA H100 GPUs) using a FedOpt configuration similar to DiLoCo (AdamW as inner optimiser and Nesterov’s accelerated
6
G. Mittone, M. Aldinucci
gradient as outer, i.e., FedMom [7]). Photon confirms that loosely-coupled compute can be aggregated without sacrificing – and sometimes even improving – convergence relative to centralised training. Both DiLoCo and Photon demonstrate that FL can be an effective alternative to conventional centralised training. However, neither envisages FL as an intrasilo (i.e., single-cluster) DP training strategy, nor addresses the communication patterns that arise from such a combination at scale. The proposed contribution differs from these works in three key respects: – Intra-silo focus–this work targets the hierarchical networks typical of modern HPC clusters, where intra-node bandwidth is orders of magnitude higher than intra-rack bandwidth, which is higher than inter-rack bandwidth. – FL as a higher-level DP primitive–FL is employed not as a replacement for DP but as an additional, orthogonal, coarser-grained parallelism layer that coordinates multiple DP deployments through focused communications. – Scalable communication model –the proposed three-dimensional communication topology (sharding, replication, and federation groups) reflects the physical interconnects of modern HPC systems, allowing scaling of FM training more efficiently than the performance offered by standard DP algorithms.
4
Enhancing FSDP and HSDP Performance Through FL
This research work augments sharded DP training algorithms with FL-inspired techniques to dramatically reduce the frequency of heavyweight collectives over slow, more contended, and noisy links, while exploiting faster, more local interconnects. This is achieved by introducing an additional, orthogonal dimension to the process topology used by sharded DP algorithms–the federation group dimension (right side of Figure 1). Training processes are first grouped into conventional sharding and replication groups, and then further partitioned into federation groups that communicate only through lightweight FL-style aggregations. By confining expensive collective communications to intra-federated-group levels and exploiting lightweight FedAvg-style inter-federated-group synchronisation, the proposed algorithms can leverage the multi-tier structure of modern HPC interconnects, thereby reducing overall communication overhead. Before delving into the proposed FL+FSDP and FL+HSDP training schemes, the communication costs of traditional state-of-the-art training techniques are presented. 4.1
Communication Cost of SOTA Techniques
This work focuses on communication time; other overheads, such as startup time and latency, are not modelled, as their impact, although tangible, falls outside the scope of the proposed discussion. Table 1 summarises the used notation. Collective communications cost formulas are based on NVIDIA NCCL 1 [13]. 1
https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md)
Accelerating Sharded Data Parallelism at Scale with Federated Learning
7
Table 1. Table of symbols used in communication cost formulas Symbol
Description
t Bn
Time [s] Bottleneck bandwidth among n processes
L S R F A M
Number of sharded model layers Number of shards per model replica Number of model replicas Number of federated groups Aggregation Frequency [batches] Layer shard size [bits]
CAG (D, n)
Cost of an all-gatheramong n processes with message size D Cost of a reduce-scatteramong n processes with message size D Cost of an all-reduceamong n processes with message size D Cost of a broadcastamong n processes with message size D
CRS (D, n) CAR (D, n) CBC (D, n)
Value
L = 1 if not sharded S = 1 if not sharded R = 1 if FSDP F = 1 if not federated A = 1 if not federated n−1 t = D nB n n−1 t = D nB n n−1 t = 2D nB n
t = BDn
FSDP As described in Section 2, FSDP and HSDP exhibit different collective communication patterns, which determine distinct cost models. FSDP arranges the S available processes into a one-dimensional sharding group; consequently, every collective operation—whether an all-gather, reduce-scatter, or all-reduce—is executed over the same mesh. The communication cost (in terms of time) of processing one data batch through FSDP can be modelled as:
S−1 CFSDP = 2 CAG (M, S) L + CRS (M, S) L = M L 3 SBS
(1)
where BS denotes the smallest bandwidth available inside the sharding group.
HSDP HSDP organises the processes into a two-dimensional R × S mesh consisting of R sharding groups and S replication groups. Each replication group acts as an independent FSDP deployment, incurring the same communication schema (two all-gathers and the reduce-scatter per layer) and cost (the S−1 term of Equation 1). In addition, HSDP requires an all-reduce of the 3 SB S gradient shards across the replication groups; this operation is less frequent (once per layer during the backward pass) and can be executed in parallel over the sharding dimension. The HSDP communication cost can thus be modelled as:
8
G. Mittone, M. Aldinucci
CHSDP = 2CAG (M, S) L + CRS (M, S) L + CAR (M, R) L S−1 R−1 = ML 3 +2 SBS RBR
(2)
where BS and BR are, respectively, the minimal bandwidth within the sharding and replication groups. FL Conversely, assuming a FedAvg aggregation strategy, FL requires calculating the weighted average of all F models only every A training steps. Such an approach is usually applied in client-server architectures, but can be ported to fully decentralised topologies, as in standard DP, by arranging the F processes into a one-dimensional mesh and running an all-reduce over the model parameters every A training steps. It can be further adapted to standard DP by running an all-reduce for each model layer, rather than a single one across the entire model. Given such assumptions, the average communication cost for processing one data batch with FL through decentralised FedAvg is: 1 F −1 1 2 (3) CFL = CAR (M, F ) L = M L A A F BF where BF denotes the minimal bandwidth between the federation participants. This work focuses on FedAvg as a base aggregation strategy, enabling easy integration of more complex ones, since most of these are built on it. 4.2
Augmenting Sharded DP Through FL
The proposed approaches share a fundamental assumption: the global processes pool is partitioned into F federation groups, and each federation group applies the same sharding policy. This guarantees one-to-one correspondence between shards across different federation groups, enabling correct aggregation. Violating such an alignment constraint has severe consequences. In the best-case scenario, the mismatch in shard sizes would cause a runtime error, halting training. In the worst-case scenario, shards with different semantic values could be aggregated, leading to silently corrupted models and catastrophic learning divergence. FL+FSDP Multiple instances of the same FSDP training can be loosely synchronised through FL by a decentralised, shard-by-shard FedAvg implementation. Each of the F parallel FSDP trainings is modelled as a federation group containing a single sharding group of size S. Each federation group proceeds independently of the others, thus retaining the same CFSDP as detailed in Equation 1 and performing FedAvg aggregation every A training steps. Based on the decentralised FedAvg technique discussed in Section 4.1, such aggregation can be performed by running S independent, parallel all-reduces for every set of F corresponding shards across the federation groups for all L layers. In this
Accelerating Sharded Data Parallelism at Scale with Federated Learning
9
way, after each aggregation, every shard is equal to the average of its corresponding shards in all the other federation groups, making the sharding groups synchronised and ready to continue training. The average communication cost for processing one data batch with FL+FSDP can be modelled as: 1 S−1 1 F −1 CFL+FSDP = CFSDP + CAR (M, F ) L = M L 3 + 2 (4) A SBS A F BF FL+HSDP As in the FL+FSDP case, assume F parallel HSDP trainings with dimensions R × S, each constituting a federation group with communication cost CHSDP detailed in Equation 2. If R = 1, then this scenario falls back to FL+FSDP, implying the same communication cost CFL+FSDP (Equation 4). If R > 1, then it is possible to apply the same shard-by-shard decentralised FedAvg aggregation discussed in Section 4.2 across corresponding replication groups between federation groups. Such an approach works because all replication groups within the same federation group are already synchronised with one another via HSDP, as discussed in Section 2. It is possible to save further network traffic and achieve better interconnection exploitation by running the FedAvg aggregation only between one replica per federation group and then having the updated replica broadcast its parameters to the other replicas in its federation group. In this way, only the minimal amount of information necessary is communicated between inter-federation groups, exploiting more the intra-federation-group interconnect, which is expected to be more efficient (i.e., BR ≫ BF ). The average communication cost for processing one data batch with FL+HSDP is thus: 1 1 CFL+HSDP = CHSDP + CAR (M, F ) L + CBC (M, R) L A A S−1 R−1 1 F −1 1 = ML 3 +2 + 2 + SBS RBR A F BF BR 4.3
(5)
Communication Cost Comparison
Table 2 compares the theoretical communication cost of the base sharded DP algorithms and their FL-enhanced versions. As can be seen, the FSDP performance is bounded by the minimum bandwidth BS available between any pair of processes in the sharding group, as captured by the term 3 SS−1 BS , which dominates the communication cost. HSDP, in turn, contributes the 2 RR−1 BR term to the FSDP cost, outperforming FSDP only when the intra-sharding group network is significantly faster than the inter-replication group one, i.e., BS ≫ BR . In practice, this situation arises when each sharding group is confined to a single physical node, benefitting from the high bandwidth of NVLink or PCIe, while replication groups span multiple nodes, relying less on InfiniBand, Ethernet, etc. FL communication cost is always lower than FSDP, assuming the same message size and same process groups minimal bandwidth (i.e., S = F, BS = BF ), and
10
G. Mittone, M. Aldinucci
Table 2. Communication cost of FSDP, HSDP, and their FL-augmented versions
FSDP HSDP
Distributed S−1 M L 3 SB S S−1 R−1 M L 3 SB + 2 RB S R
Federated i S−1 1 −1 M L 3 SB +A 2 FFB S F h i S−1 R−1 1 −1 M L 3 SB + 2 RB +A 2 FFB + B1R S R F h
also HSDP, assuming a square mesh (S = R = F ) and that aggregation happens on the slowest bandwidth (BF = BR , BR ≪ BS ), for any value of A. Passing on the FL-enriched algorithms, it is evident that FL+FSDP is very similar, from a communication pattern and overhead perspective, to HSDP. HSDP all-reduces the sharded gradients before updating the model’s parameters during each training step, while FL+FSDP all-reduces the sharded model’s parameters after the local update every A training steps. Models trained with HSDP thus never diverge from each other, whereas models trained with FL+FSDP constantly diverge after each aggregation. FL+FSDP is advantageous over standard FSDP only if BS ≫ BF , as HSDP, which always holds when the number of processes exceeds the number of GPUs available on a single compute node. Conversely, FL+HSDP models three tiers of interconnection performance, i.e., BS ≫ BR ≫ BF , allowing better exploitation of the fastest and most efficient links while proportionally relying less on the slowest, more contended ones. It also allows, as FL+FSDP, to tune the aggregation frequency through the A parameter based on the specific deployment. These characteristics make FL+HSDP an extremely flexible, general, and powerful sharded DP approach, opening up new possibilities in large-scale FM training.
5
Experimental Evaluation
Experimental Setup Experimental FL+FSDP and FL+HSDP implementation exploiting PyTorch v2.10.0 and NVIDIA NCCL v2.27.5 are built into crossFacility Federated Learning (xFFL)2 [1], an extensive, open-source, researchoriented Python framework integrating tools for deploying large-scale FM trainings, both on single- or cross-facility scenarios. A brief pre-training of Llama3.1 8B [10] on ≈150M tokens from the clean_mc4_it dataset [16] is selected as benchmark. All runs share the same hyper-parameter setup: bfloat16 precision, 65,536 training samples, 4,096 test samples, batch size 2, AdamW optimiser, learning rate 0.0003, aggregation every 8 steps, cosine-decay scheduler with 10% warm-up steps, and same RNG configuration. Experimental evaluation is carried out on CINECA’s Leonardo—currently the tenth most powerful Top500 HPC [21]—using 128 nodes. Each node comprises a single-socket 32-core Intel Xeon Platinum 8358 CPU, 8×64 GB DDR4 RAM, 4× custom NVIDIA Ampere A100 GPUs (64 GB), 4× NVLink 3.0 links 2
https://github.com/alpha-unito/xffl/tree/FL+DP
Accelerating Sharded Data Parallelism at Scale with Federated Learning
105
11
FSDP HSDP FL+FSDP FL+HSDP
Perplexity
104
103
102
0
1
2
3
4
5
6 7 Time (minutes)
8
9
10
11
12
Fig. 2. Training perplexity of Llama3.1-8B on ≈150M tokens from the clean_mc4_it dataset, measured on 512 GPUs under a common hyper-parameter configuration
(4×200 Gbit/s per link) connecting the GPUs, and 2× dual-port HDR100 NICs (400 Gbit/s). Nodes are linked via an NVIDIA Mellanox DragonFly+ fabric that provides 200 Gb/s bandwidth between any pair of nodes. A more extensive set of results—omitted here for space constraints—is publicly available on WandB3 . 5.1
Learning Results
Figure 2 and Table 3 report the obtained perplexity for FSDP, HSDP, FL+FSDP, and FL+HSDP. FSDP, assumed as a baseline, achieves 2.32× lower train perplexity than HSDP, despite being theoretically equivalent. The different communication patterns can explain this result: HSDP averages gradients via a two-tier all-reduce, whereas FSDP performs a single, global one. The resulting change in floating-point operation order, amplified by the limited precision of bfloat16, can lead to divergent numerical trajectories, especially for large models. Conversely, both FL-augmented approaches converge markedly better than their counterparts. FL+FSDP improves train perplexity by 1.30× over FSDP and 2.86× over HSDP. FL+HSDP achieves the lowest perplexity of all, obtaining better results than FSDP, HSDP and FL+FSDP by factors of 1.95×, 4.53×, and 1.59×, respectively. The same is observed on the test set: FL+HSDP outperforms FSDP, HSDP, and FL+FSDP by 1.83×, 4.48×, and 1.71×, confirming the learning improvement. These gains can stem from the reduced effective global batch size implied by the FL-augmented schemes. By partitioning the workload into loosely-coupled federation groups, the global batch grows only with the number 3
https://wandb.ai/alpha-unito/FL+DP/workspace?nw=fe1n0p5w1r7
12
G. Mittone, M. Aldinucci
Table 3. Llama3.1 8B train and evaluation results over ≈150M tokens over 128 Leonardo nodes (512 GPUs) - aggregation every 8 steps Training Perplexity Time [s] FSDP HSDP FL+FSDP FL+HSDP
496.00 1151.47 403.29 254.31
770.54 117.37 95.87 101.04
Evaluation Perplexity Time [s] 486.64 1145.78 436.65 255.81
15.26 1.83 1.80 1.89
of groups—not with the total number of processes—resulting in smaller, more stable batches. Periodic FedAvg aggregations then stabilise convergence at the global level, yielding higher learning performance. 5.2
Computational Performance
Figure 2 and Table 3 show that FSDP is the slowest method at this scale among those evaluated. In contrast, HSDP, FL+FSDP and FL+HSDP exhibit comparable runtimes. FL+FSDP achieves the best throughput, processing the same amount of data 8.04×, 1.22×, and 1.05× faster than FSDP, HSDP, and FL+HSDP, respectively. The same holds for the evaluation phase, where FL+FSDP is 8.48×, 1.02×, and 1.05× faster than FSDP, HSDP, and FL+HSDP. These performance gains can be traced back to the different communication patterns involved. FSDP relies on a flat communication structure, incurring frequent, large-scale collectives across many GPUs that are severely slowed by the higher-level, noisy inter-node links. By contrast, HSDP and the FL-augmented schemes employ multi-tier communication patterns, exploiting intra-node, internode, and even higher-level connections to keep most traffic on the local, faster links, while relying proportionally less on the slower ones. Table 4 breaks down the training runtime into each step sub-phase: forward, backward, aggregation, and optimization. The forward pass performance of FSDP is the slowest, up to 12.24× slower than the counterparts, which are instead almost aligned. This advantage stems from the better locality of communication in these algorithms, which exploit the higher bandwidth of local interconnects. However, the HSDP backward pass performance is the one that suffers more, in proportion to the forward pass, mainly due to the switch to cluster-wise gradient reduction collectives. Conversely, FL-augmented approaches avoid such slowdowns by restricting gradient reduction to federation groups, dramatically reducing the volume of data sent over the slowest links, and are thus up to 8.43× and 1.61× faster than FSDP and HSDP. The optimisation step is practically 0 for FSDP due to the high level of sharding in the proposed setup; in contrast, in the other algorithm, it is constant because of the fixed-size replica group. However, FL introduces the aggregation step: communication is run over the slowest, busiest links in the cluster, and as such exhibits high variance. Still,
Accelerating Sharded Data Parallelism at Scale with Federated Learning
13
Table 4. Llama3.1 8B training time detail over 128 Leonardo nodes (512 GPUs) ± standard deviation (seconds) - aggregation every 8 steps
FSDP HSDP FL+FSDP FL+HSDP
Forward
Backward
Optimizer
Aggregation
Overall
4.53±0.16 0.37±0.00 0.37±0.00 0.38±0.02
7.50±0.20 1.43±0.15 0.89±0.00 0.99±0.01
0.00±0.00 0.03±0.00 0.03±0.00 0.03±0.00
0.00±0.00 0.00±0.00 0.17±0.55 0.14±0.54
12.04±0.26 1.84±0.15 1.45±0.54 1.53±0.54
aggregation occurs only once every A training steps, so its frequency and perbatch impact remain limited. Overall, FL+FSDP processes a mini-batch up to 8.30× faster than FSDP, while matching HSDP’s performance (the exact ratio depends on the chosen aggregation frequency). FL-augmented approaches thus offer up to 8.30× and 1.27× better per-batch processing performance than FSDP and HSDP, while also implying higher variances due to differences in execution time between training steps that require or do not require aggregation.
6
Discussion
The FL-augmented sharded DP algorithms reduce processing time by up to 8.04× and improve convergence by up to 4.48× compared to their standard counterparts in the proposed scenario. These results lead to two observations: 1) standard FSDP does not scale to hundreds or thousands of GPUs due to its flat communication topology, forcing every collective to involve all processes and becoming a bottleneck; and 2) HSDP mitigates such problem through a two-tier communication hierarchy, but the gain in raw throughput can be offset by noticeable degradation in learning performance. From a computational standpoint, the proposed FL-augmented algorithms successfully model a twotier (FL+FSDP ) and a three-tier (FL+HSDP ) communication topology, giving fine-grained control over the communications pattern and minimising traffic on heavily contended links. From a learning perspective, FL-augmented approaches prevent the global batch size from growing linearly with the number of processes. Instead, they bound it to the federation group size, yielding markedly better convergence and reducing dependence on aggressive hyper-parameter tuning. Finally, experimental results confirm that FL+FSDP ’s runtime aligns with HSDP’s, as predicted by the communication-cost model discussed in 4.2; the aggregation frequency governs the remaining difference. FL+FSDP delivers superior computational performance to FL+HSDP. Although the three-tier topology modelled by the latter would be expected to perform better, the observed advantage of FL+FSDP can be explained by two factors: 1) the bandwidth difference between intra- and inter-federation groups is insufficient to compensate for the extra broadcastcommunication step introduced by FL+HSDP, and 2) SLURM’s nondeterministic node allocation splits
14
G. Mittone, M. Aldinucci
federation groups across different interconnection groups, breaking the assumed logical hierarchy. Even so, the performance gap is small compared with the disparity in final perplexity: FL+FSDP is only ≈5.39% faster thanFL+HSDP, but FL+HSDP achieves a ≈70.69% improvement in perplexity.
7
Conclusions and Future Works
This work introduces FL+FSDP and FL+HSDP, two novel FL-augmented DP training algorithms that combine state-of-the-art sharded DP methods (FSDP, HSDP) with FL techniques (FedAvg). By modelling a three-tier communication hierarchy—sharding, replication, and federation groups—the schemes dramatically reduce inter-node traffic and the overall communication overhead. Both approaches are evaluated through a formal theoretical analysis of their communication costs and large-scale experiments on a LLM pre-training task. The results show up to 8.04× faster data processing and 4.48× lower evaluation perplexity compared with the baseline methods. These findings confirm that FL can be considered a “higher-level” DP primitive, useful not only for cross-cluster training but also for scaling state-of-the-art training of FMs within a single one. Future work will extend the methodology to other parallelism axes—MP, TP/SP, and expert EP—thereby treating FL as a fifth orthogonal dimension of FM parallelism. Further research will also automate the selection of sharding, replication, and federation-group sizes based on the underlying interconnect topology, eliminating manual tuning while maximising performance. Acknowledgements and Artifact Availability This research work is funded by the Spoke “FutureHPC & BigData” of the ICSC - Centro Nazionale di Ricerca in “High Performance Computing, Big Data and Quantum Computing”, by the European Union - NextGenerationEU, by the EuroHPC-JU funding under grant No. 101093441, with support from the Horizon-EuroHPC-JU-2021-COE-01 (SPACE CoE) and by the DYMAN project, funded by the European Union - European Innovation Council under G.A. n. 101161930. The authors also acknowledge the computational support kindly offered by the Abdus Salam International Centre for Theoretical Physics (ICTP), Trieste, Italy. The artifact is available in the Zenodo repository [12]. Disclosure of Interests. The authors have no competing interests to declare that are relevant to the content of this article.
References 1. Colonnelli, I., et al.: Cross-facility federated learning. Procedia Computer Science 240, 3–12 (2024). https://doi.org/10.1016/j.procs.2024.07.003, proceedings of the First EuroHPC user day 2. De Sensi, D., et al.: Exploring gpu-to-gpu communication: Insights into supercomputer interconnects. In: Proc. of the Int. Conf. for HPC, Networking, Storage, and Analysis, SC. p. 33. IEEE (2024). https://doi.org/10.1109/SC41406.2024.00039
Accelerating Sharded Data Parallelism at Scale with Federated Learning
15
3. Douillard, A., et al.: Diloco: Distributed low-communication training of language models. CoRR abs/2311.08105 (2023) 4. Hao, Z., et al.: Low-precision training of large language models: Methods, challenges, and opportunities. CoRR abs/2505.01043 (2025) 5. Hoefler, T., et al.: Hammingmesh: A network topology for large-scale deep learning. Commun. ACM 67(12), 97–105 (2024). https://doi.org/10.1145/3623490 6. Huang, Y., et al.: Gpipe: Efficient training of giant neural networks using pipeline parallelism. In: Adv. in Neural Information Processing Systems 32: Annual Conf. on Neural Information Processing Systems, NeurIPS. pp. 103–112 (2019) 7. Huo, Z., et al.: Faster on-device training using new federated momentum algorithm. CoRR abs/2002.02090 (2020) 8. Lewis, N., et al.: I/O in machine learning applications on HPC systems: A 360degree survey. ACM Comput. Surv. 57(10), 256:1–256:41 (2025). https://doi.org/ 10.1145/3722215 9. Li, S., et al.: Pytorch distributed: Experiences on accelerating data parallel training. Proc. VLDB Endow. 13(12), 3005–3018 (2020). https://doi.org/10.14778/ 3415478.3415530 10. Llama Team: The llama 3 herd of models. CoRR abs/2407.21783 (2024) 11. McMahan, B., et al.: Communication-efficient learning of deep networks from decentralized data. In: Proc. of the 20th Int. Conf. on Artificial Intelligence and Statistics, AISTATS. vol. 54, pp. 1273–1282. PMLR (2017) 12. Mittone, G., et al.: Artifact of the paper: Accelerating sharded data parallelism at scale with federated learning (Jun 2026). https://doi.org/10.5281/zenodo.20528934 13. NVIDIA: NVIDIA NCCL documentation. https://docs.nvidia.com/deeplearning/ nccl/user-guide/docs/index.html (2025), accessed: 2025-09-18 14. Rajbhandari, S., et al.: Zero: memory optimizations toward training trillion parameter models. In: Proc. of the Int. Conf. for HPC, Networking, Storage and Analysis, SC. p. 20. IEEE/ACM (2020). https://doi.org/10.1109/SC41405.2020.00024 15. Sani, L., et al.: Photon: Federated LLM pre-training. In: Proc. of the Eighth Conf. on Machine Learning and Systems, MLSys. OpenReview (2025) 16. Sarti, G., et al.: IT5: text-to-text pretraining for italian language understanding and generation. In: Proc. of the 2024 Joint Int. Conf. on Computational Linguistics, Language Resources and Evaluation. pp. 9422–9433. ELRA/ICCL (2024) 17. Sevilla, J., et al.: Compute trends across three eras of machine learning. In: Int. Joint Conf. on Neural Networks, IJCNN. pp. 1–8. IEEE (2022). https://doi.org/ 10.1109/IJCNN55064.2022.9891914 18. Shazeer, N., et al.: Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In: Int. Conf. on Learning Representations, ICLR (2017) 19. Shoeybi, M., et al.: Megatron-lm: Training multi-billion parameter language models using model parallelism. CoRR abs/1909.08053 (2019) 20. Shpiner, A., et al.: Dragonfly+: Low cost topology for scaling datacenters. In: Int. Work. on High-Performance Interconnection Networks in the Exascale and BigData Era. pp. 1–8. IEEE (2017). https://doi.org/10.1109/HIPINEB.2017.11 21. Top500: TOP500 list - november 2025. https://top500.org/lists/top500/list/2025/ 11/ (2026), accessed: 2026-02-26 22. Zhang, S., et al.: OPT: open pre-trained transformer language models. CoRR abs/2205.01068 (2022) 23. Zhao, Y., Gu, A., Varma, R., et al.: Pytorch FSDP: experiences on scaling fully sharded data parallel. Proc. VLDB Endow. 16(12), 3848–3860 (2023). https://doi. org/10.14778/3611540.3611569