Shard the Gradient, Scale the Model: Serverless Federated Aggregation via Gradient Partitioning

Abstract: Federated learning (FL) aggregation on serverless platforms faces a hard scalability ceiling: existing architectures (lambda-FL, LIFL) partition clients across aggregators, but every aggregator must hold the complete model gradient in memory. When gradients exceed the per-function memory limit (e.g., 10 GB on AWS Lambda), aggregation becomes infeasible regardless of tree depth or branching factor. We propose GradsSharding, which instead partitions the gradient tensor into M shards, each averaged independently by a serverless function that receives contributions from all clients. Because FedAvg averaging is element-wise, this produces bit-identical results to tree-based approaches, so model accuracy is invariant by construction. Per-function memory is bounded at O(|θ|/M), independent of client count, enabling aggregation of arbitrarily large models. We evaluate GradsSharding against lambda-FL and LIFL through HPC experiments and real AWS Lambda deployments across model sizes from 43 MB to 5 GB. Results show a cost crossover at approximately 500 MB gradient size, 2.7x cost reduction at VGG-16 scale, and that GradsSharding is the only architecture that remains deployable beyond the serverless memory ceiling.

• The paper introduces GradsSharding, which partitions gradient tensors by parameter rather than by client to decouple aggregator memory constraints from model size.
• It demonstrates that sharding enables up to 16× speedup in aggregation latency and significantly lowers cost by reducing per-shard memory and compute overhead.
• The study underscores that serverless federated aggregation using gradient sharding is essential to scale models beyond the 10GB memory cap of typical FaaS platforms.

Shard the Gradient, Scale the Model: Serverless Federated Aggregation via Gradient Partitioning

Introduction and Problem Scope

The expansion of model sizes in Federated Learning (FL) directly exposes a critical bottleneck in serverless aggregation infrastructures: memory ceilings imposed by Functions-as-a-Service (FaaS) platforms such as AWS Lambda. Conventional serverless FL aggregation architectures, notably $\lambda$-FL and LIFL, partition by clients but each aggregator is required to buffer the full model gradient, which makes these approaches fundamentally unscalable as gradient dimensions approach tens of gigabytes. The paper "Shard the Gradient, Scale the Model: Serverless Federated Aggregation via Gradient Partitioning" (2604.22072) addresses this core scalability constraint with the GradsSharding architecture, which partitions by gradient parameters and not by client, thereby decoupling aggregator memory from both model dimensionality and client count.

Aggregation Architectures: Model, Limitations, and GradsSharding

FL’s classic FedAvg protocol requires that a central aggregator efficiently performs element-wise averaging of client-uploaded gradients. For small- and medium-scale models, this procedure is tractable on FaaS platforms; however, with modern architectures (e.g., VGG-16, GPT-2 Large), single gradient objects often surpass Lambda’s $10$GB memory constraint, regardless of the hierarchical depth of the aggregation tree. Both $\lambda$-FL and LIFL suffer from this architectural barrier, as every tree node (aggregator) in their job DAG requires full-gradient buffering for combining updates. LIFL optimizes for shared-memory event transfer but is thwarted on platforms with strict container isolation.

GradsSharding resolves this by partitioning each gradient tensor along the parameter axis into $M$ disjoint shards. Each shard aggregator independently accumulates $|\theta|/M$ elements from all $N$ clients, and the full average is simply the concatenation of these $M$ independently-averaged shards. This sharding directly limits the memory per Lambda function to $\mathcal{O}(|\theta|/M)$, thus supporting arbitrarily large models by simply increasing $M$. This approach does not alter the semantics of FedAvg—bitwise identical aggregates are guaranteed since averaging is element-wise and perfectly parallelizable.

The functional architecture and contrasts between $\lambda$-FL, LIFL, and GradsSharding are depicted schematically in the system model section.

Figure 2: Round time breakdown across model scales. Training time (blue) dominates, resulting in high parameter server idle ratios for large models.

Quantitative Evaluation: Latency, Cost, and Memory Scaling

Parameter Server Idle Time

A core empirical result is the quantification of the server's idle-to-active ratio during each FL round. For models spanning ResNet-18 ($10$0MB) to GPT-2 Large ($10$1GB), the aggregator is idle for 79.8% to 99.6% of the time, waiting for client training, which fortifies the rationale for ephemeral, FaaS-based aggregation rather than persistent server deployments.

Lambda Aggregation and S3 Overhead

The paper presents measurements from real AWS Lambda deployments, demonstrating that S3 read time constitutes 91--99% of total aggregation duration. Aggregation compute is negligible (1.3–2.9s across models) compared to I/O (e.g., reading $10$2GB from S3 spans $10$3s per large-shard aggregator).

Figure 1: Lambda aggregation time breakdown; S3 reads are blue and dominate the wall clock, while FedAvg compute remains a small orange constant across scales.

Shard Ablation: Memory, Latency, and Cost Scaling with $10$4

The GradsSharding shard ablation confirms the anticipated asymptotic scaling law: $10$5-fold sharding reduces aggregator memory by $10$6, making previously infeasible models tenable within Lambda's memory constraints. Cumulative aggregation latency increases linearly in a sequential environment (if all $10$7 aggregators run serially), but under Lambda’s concurrent model, wall-clock latency drops inversely: a 16-shard configuration yields a $10$8 speedup in time-to-aggregate compared to a single-shard design.

S3 I/O cost grows linearly in $10$9, but Lambda compute cost drops due to reductions in per-shard memory and execution time. For example, VGG-16 aggregation at $\lambda$0 is completed in 11.2s—16.2× faster than $\lambda$1—for a minimal 19% cost increase.

Figure 3: Cost per round decomposed into Lambda compute and S3 I/O as a function of the number of shards $\lambda$2.

Figure 4: VGG-16 shard sweep on AWS Lambda: (a) Aggregation time is strongly reduced with increased sharding, (b) S3 cost rises with $\lambda$3, but Lambda compute cost drops.

Figure 5: Latency and cost tradeoff; $\lambda$4 is cost-optimal, $\lambda$5 latency-optimal. S3 I/O linear in $\lambda$6, with per-function throughput stable.

Cross-Architecture Analysis

A comprehensive Lambda deployment shows that for small models ($\lambda$7MB), $\lambda$8-FL is marginally cheaper due to fewer S3 operations. As model size increases, GradsSharding quickly dominates: at VGG-16 scale, it achieves 2.7$\lambda$9 lower cost and 1.7$M$0 lower latency. Beyond $M$13GB gradients, $M$2-FL and LIFL exceed Lambda's hard memory cap, becoming infeasible, whereas GradsSharding remains deployable for arbitrarily large models.

Figure 6: Cross-architecture comparison of aggregation time and cost per 1,000 rounds across model scales; only GradsSharding is viable at multi-GB scales.

Practical and Theoretical Implications

Scaling Federated Learning with FaaS

This study demonstrates that parameter-axis sharding at aggregation, rather than classic client-topology partitioning, is crucial for compatibility with stateless serverless infrastructures under realistic public cloud constraints. GradsSharding enables FL aggregation for extremely large models within Lambda, with performance and cost characteristics dominated by S3 I/O rather than compute. This shifts the focus of system optimization away from parallel arithmetic design towards minimizing cloud object store overheads (e.g., through compression, S3 batch operations, or hybrid direct invocation for small shards).

Guidelines for $M$3 Selection

Latency-sensitive applications should maximize $M$4; cost-constrained applications may reduce $M$5 to the lowest value that still fits within Lambda's $M$6GB memory ceiling. As concurrency is unconstrained on Lambda, wall-clock is determined by the largest shard's slowest download, not the sum of per-shard compute.

Directions for AI System Design

This work reveals the importance of system-architecture–aware model scaling in the cloud. For federated learning at the billion-parameter level, it is no longer client compute/comm or global convergence that bottlenecks, but the resource model of aggregation infrastructure. Hence, future developments in FL system design for cloud-native settings will need to integrate gradient compression prior to sharding, adaptive $M$7 tuning, and even architecture-aware sharding to account for non-uniform parameter distributions in models such as MoEs or transformer hybrids.

Conclusion

The paper substantiates that GradsSharding is the only viable aggregation topology for serverless FL at modern model scales. It imposes a strict and provable memory bound per function, yields elementwise-accurate aggregation, and is shown to outperform legacy client-based partitioning in cost and latency for anything beyond trivial FL models. These results suggest broader opportunities for integrating parameter-sharded, concurrent aggregation as the standard for future FL services on major serverless platforms, particularly as models continue to scale and cloud-native federated deployment becomes the norm.