Federated Learning (FL) is a distributed learning paradigm where edge devices or clients collaboratively train machine learning (ML) models using privately held local data [McMahan et al., 2017, Kairouz et al., 2021]. Clients iteratively update their local models, which are periodically sent to a central server for aggregation. The aggregated model is then sent to the clients to begin the next round of local updates. Since its introduction in [McMahan et al., 2017], various practical and theoretical aspects of FL, in-cluding client selection [Nishio andYonetani, 2019, Cho et al., 2022], communication challenges [Ang et al., 2020, Chellapandi et al., 2023], scalability and fast training [Xie et al., 2019, Wang et al., 2020b], have been extensively studied. However, these works almost exclusively assume that the server aims to learn model(s) for a single task. Some FL frameworks attempt to learn models personalized to each client [Mansour et al., 2020, Li et al., 2021, Tan et al., 2022], but these models are still intended for the same learning task, e.g., next-word prediction on keyboards.
Many practical applications need devices to perform a wide range of learning tasks, which require training of multiple ML models. For instance, our phones need language models for keyboard next-word prediction as well as image recommendation models to highlight images more likely to be shared [McMahan et al., 2017]. Yao et al. [2023] propose training multiple models in federated smart car networks for different tasks, such as pothole detection and maneuver prediction. Another example can be a chat application requiring speech recognition and response text generator models concurrently, while Le et al. [2022] suggest federated learning of multiple models for air quality index forecasting. Thus, in this paper, we seek to answer the following question:
How can we efficiently train models for multiple tasks in a federated setting using a shared pool of clients?
Simple Solutions that Extend FedAvg. A naïve approach to training multiple models is sequential training, where the models corresponding to different tasks are trained one at a time, each utilizing all the clients. The total training runtime of this approach scales linearly with the number of tasks. An alternative is for all the clients to train all tasks at the same time. However, with this approach each client will have to keep all models in memory, which is infeasible for resource-limited edge clients such as smartphones. To preserve memory, clients will have to queue the training requests and process them sequentially, again resulting in the runtime linearly increasing with the number of tasks. On the other hand, parallel or simultaneous training (ST) of all arXiv:2406.00302v1 [cs.
LG] 1 Jun 2024
the models with time-varying subsets of clients assigned to each task can strike a better trade-off between accuracy and runtime. Bhuyan et al. [2023]'s approach assigns a disjoint subset of clients to each model in each round, which significantly improves the time taken to reach a target accuracy as compared to sequential training. However, these federated simultaneous training (FST) approaches leave room for significant improvement. There are two particular drawbacks: 1) straggler delays due to synchronous aggregation, and 2) the lack of adaptation to the training progress of heterogeneous tasks, which we address in this work.
Synchronous Aggregation and Straggler Delays. Conventional FL employs synchronous aggregation, where in each round, the server waits to receive updates from all the participating clients before each aggregation. However, when the clients have diverse hardware and communication capabilities, faster clients must remain idle until slow or straggling clients finish, causing a large wallclock runtime to complete each communication round. This problem is further exacerbated in FL with multiple simultaneous models [Bhuyan et al., 2023, Zhou et al., 2022], where the aggregation is synchronized across tasks as well. Therefore, the server has to wait for the slowest client across all the parallel tasks. Solutions proposed to alleviate the straggler problem in the single-model context include allowing faster clients to run more local steps [Wang et al., 2020b], aggregating only the client updates that arrive before a timeout [Bonawitz et al., 2019], and sub-sampling from the set of available clients [Luo et al., 2022]. Although these approaches perform well when stragglers appear uniformly at random, they do not work well in the simultaneous training setting because some models (e.g., larger ones) are naturally slower to train. When the multiple models have inherently different training times, synchronized global aggregation rounds are bottlenecked by the slowest client assigned to the most computationally intensive model, leading to large idle times.
Issues. Another solution to the straggler problem is asynchronous aggregation at the server, as proposed in AsyncFL [Xie et al., 2019], where the server updates the global model whenever it receives any client update. While asynchronous aggregation has been extensively studied in single-model federated learning [Chen et al., 2020, Wang et al., 2022, Xu et al., 2023, Yu et al., 2023a], it has not been well-explored for simultaneous federated training. Although AsyncFL addresses the straggler issue, it suffers from undesired staleness even in the standard FL setting, since the received client updates are often based on outdated models. To alleviate the staleness problem in single-model FL, Nguyen et al. [2022] proposed storing the incoming client updates in a buffer at the server and aggregating when the buffer is full.
Adaptive Allocation of Clients to Heterogeneous Tasks.
In this work, we employ asynchronous buffered aggregation to overcome the straggler issue while controlling staleness. However, extending single-model FL algorithms [Xie et al., 2019, Nguyen et al., 2022] to the simultaneous training of multiple models is not straightforward -running multiple independent instances of asynchronous FL can be suboptimal. This is because the tasks can have heterogeneous computation complexities and different data heterogeneity that affect both the number of rounds required to achieve a given target accuracy as well as the wall-clock time taken to complete each round. Since a shared set of clients is used to train the models, the training processes are coupled -more resources assigned to one task implies less for the others. Moreover, the optimal resource requirement for each task can change over time according to its data heterogeneity and training progress and may be difficult to predict before training. Therefore, we propose an adaptive algorithm that dynamically reallocates clients across tasks depending on their training progress, and also adapts the buffer size used for asynchronous aggregation of updates.
Our Contributions. We formalize the FST setting in Section 2 and then make the following main contributions:
• We introduce FedAST, a Federated Asynchronous Simultaneous Training algorithmfoot_0 to simultaneously train models for multiple tasks (Section 3). Our work is one of the first to mitigate the straggler problem faced by synchronous FST methods that extend vanilla FedAvg.
• The proposed algorithm addresses the problem of balancing resources across heterogeneous tasks, a unique challenge to the FST framework, using novel dynamic client allocation, and it also dynamically adjusts the buffer size used in asynchronous aggregation to strike the best trade-off between staleness and runtime.
• We provide a theoretical convergence analysis of FedAST (Section 4), which improves previous analyses even in the single-model FL setting. It improves upon [Koloskova et al., 2022] by considering multiple local updates and the buffer, and on [Nguyen et al., 2022] by relaxing the restrictive assumptions.
• We experimentally validate FedAST's performance (Section 5) in terms of its wall-clock training time and model accuracy on multiple real-world datasets compared to synchronous and asynchronous FL baselines.
We conclude and discuss future work in Section 6.
Related Work. Only a few recent works [Zhou et al., 2022, Bhuyan and Moharir, 2022, Siew et al., 2023, Bhuyan et al., 2023] consider federated simultaneous training of multiple models. In [Zhou et al., 2022], clients are selected with either Bayesian optimization or reinforcement learning to minimize training time and unfairness in participation. Bhuyan and Moharir [2022] formulate the client assignment FL as a bandit problem leveraging local training losses as scores. Siew et al. [2023] introduce biased client sampling, favoring the clients with higher local losses. These methods lack convergence guarantees. Bhuyan et al. [2023] assign clients uniformly at random or in a round-robin fashion and analyze the convergence assuming convex objective functions and bounded gradients. While these works only consider synchronous aggregation, Chang et al. [2023] propose a fully asynchronous FST algorithm. Their approach entails solving a non-convex optimization problem to optimize client assignment, which requires information about delays and models that may be difficult to obtain in practice. Also, the obtained bound does not converge to a stationary point in the presence of data heterogeneity and suffers from increased staleness when the number of clients increases. Lastly, Liu et al. [2023] propose an extension of their singlemodel adaptive asynchronous approach to the multi-model setting. However, they do not carefully handle heterogeneous data distributions across clients and the staleness of updates in their theoretical guarantees for a single model, and they lack these guarantees for multiple models. Also, their method under-utilizes client resources, since after sending a model update to the server, the clients are idle until the next training round.
Notations. For a positive integer c, we define [c] ≜ {1, . . . , c}. ∇ denotes stochastic gradients. Bold lowercase letters (e.g., x) denote vectors. |A| denotes the cardinality of set A. ∥ • ∥ denotes the Euclidean norm.
We now formally introduce the federated simultaneous training (FST) setting, where N clients train M models x 1 , . . . , x M corresponding to M independent tasks. For each task m ∈ [M ], our goal is to find the model that solves that following optimization problem:
where f m is the global loss function for task m, and f m,i is the local loss for task m at client i.
First, we examine a simple extension of FedAvg [McMahan et al., 2017] to simultaneous training of models for M tasks. At the start of each round, the server randomly partitions the available set of clients across the tasks [Bhuyan et al., 2023].
The server sends the current models {x (tm) m } M m=1 for all the tasks to the corresponding subset of clients. The clients perform local training (Algorithm 1) and return their updates to the server, which synchronously aggregates the updates for each task. This naïve simultaneous training extension of
FedAvg performs poorly due to stragglers. The time it takes for a client to return its updates depends on its resources and the size of the model assigned. Since the server waits for the slowest update across all the tasks before commencing the next round, the server waits much longer if a large model is assigned to a slow client. We mitigate this problem via asynchronous training in FedAST, discussed next.
Next, we describe FedAST (Algorithm 2), our proposed Federated Asynchronous Simultaneous Training algorithm, illustrated in Figure 1 for M = 2 tasks. For each task m ∈ [M ], the server maintains a round index t m that is initialized to t m = 0, the number of active training requests R Wait until server receives an update ∆ m 7:
end if clients (and consequently, more training requests R m ) to tasks with larger inter-client data heterogeneity. To empirically estimate this heterogeneity, the server stores the last
where ∆ m is the empirical mean of the ∆ m,i 's. 3 Further, in our experiments, we empirically observe that the optimal choice of buffer size b m is proportional to the number of active requests R m . See Appendix C.5 for our extensive experiments. Using (2) and these empirical observations, the optimal resource allocation emerges as the solution to the following constraints.
where the first set of constraints maintains the total computation budget across tasks, and the second set ensures the allocation of a larger number of training requests to clients with higher heterogeneity. We elaborate on the theoretical motivation for the second set of constraints in Section 4, once we establish our convergence results. We also refer the reader to Appendix A for more details on Realloc.
Sending out New Requests to Reach the New Resource Allocation. To transition from one allocation {R Algorithm 3 Realloc(option,c) Next, we discuss the assumptions used in our analysis.
Assumption 1 (Smoothness). The loss functions are Lsmooth, i.e., for all i ∈ [N ], for all m ∈ [M ], and for all
Assumption 2 (Bounded Variance). The stochastic gradient at each client is an unbiased, bounded-variance estimator of the true local gradient, i.e., for all
. Assumption 3 (Bounded Heterogeneity). The local gradients are within bounded distance of the global gradient, such that for all m ∈ [M ] and x ∈ R dm , max
Assumption 4 (Bounded Staleness). The client updates of task m are received within at most γ max m server model updates after the server sends the training request.
These assumptions are standard in the literature. Assumptions 1-3 are commonly used in the synchronous [Wang et al., 2020b, Jhunjhunwala et al., 2022] and asynchronous [Koloskova et al., 2022, Nguyen et al., 2022] FL analyses. Assumption 4 is used in the convergence proof to guarantee that none of the requested client updates takes an arbitrarily large time to return to the server and is also common in asynchronous FL works [Koloskova et al., 2022, Nguyen et al., 2022]. Furthermore, the maximum staleness can be enforced by dropping over-delayed updates in practice during the training.
Theorem 1 (Convergence of FedAST). Suppose that Assumptions 1 -4 hold, and there are R m active local training requests corresponding to task m ∈ [M ], and the server and client learning rates,
Here, b m is the buffer size, and τ m is the number of local training steps. Then, the iterates, {{x
, of Algorithm 2 satisfy:
where
Proof. See Section E in the Appendix.
Comparison with Synchronous FL Analyses. Comparison with Asynchronous FL Analyses. FedBuff [Nguyen et al., 2022] considers buffered asynchronous aggregation for a single model. Still, comparing [Nguyen et al., 2022, Corollary 1] and the bound in (4) for M = 1, their convergence result (i) depends on stronger assumptions (bounded gradient norm and uniform arrivals of client updates), and (ii) has worse asynchronous aggregation error. Moreover, our analysis is more general compared to [Koloskova et al., 2022] as they do not consider multiple local SGD steps and the buffer. Simultaneous asynchronous training is considered by [Chang et al., 2023], but we observe that they do not achieve convergence unless the data distribution across clients is identical (see [Chang et al., 2023, Eq. (19)]). We discuss the comparison of FedAST to single-model and simultaneous asynchronous federated training baselines in more detail in Appendix Section B.
Corollary 1.1 (Asymptotic convergence after setting learning rates).
√ τ m b m , the bound in Theorem 1 reduces to:
Although the given bound in Corollary 1.1 does not seem to depend on the staleness bound γ max (Assumption 4), its effect is implicit in the number of active requests R m and buffer size b m . The maximum staleness is positively correlated with R m and negatively correlated with b m . In our experiments (Appendix C.5), we tune the buffer size to maintain the update staleness at a reasonable level.
Looking at the bounds in (4) or (5), increasing R m makes the bound worse because to reach the same accuracy in (5), we need to run a higher number of server updates T m . However, increasing R m also shortens the duration between two successive server updates, making the algorithm faster in wall-clock time. We illustrate this effect with a wall-clock comparison to FST baselines below.
Impact of R m on Wall-clock Time. Suppose the arrival times of all the client updates (assuming there is no queue on the clients) are distributed as Exp(λ). The expected time to fill the buffer corresponding to task m is b m /(R m λ). Therefore, in FedAST, the expected time to complete one round at the server is inversely proportional to R m . On the other hand, the expected time to finish one round of synchronous simultaneous FedAvg training is
, which increases with R m . Also, the summation over simultaneously trained tasks shows an exacerbated straggler effect since all the clients wait for the slowest client across all the tasks.
Design of Realloc (Algorithm 3). Next, we theoretically justify the dynamic allocation of resources across tasks described in Section 3 (Algorithm 3, with option = D), which adjusts the number of active requests ({R m } M m=1 ). Given the limited number of available clients (which limits the total number of active training requests), to achieve the best possible allocation, we minimize the sum of the most dominant terms in the bounds (FedAvg Error-II in (4)) across tasks. We also use the empirical observation that the optimal choice of buffer-size b m scales linearly with R m (Appendix C.5). The resulting optimization problem is
where R is the budget for the total number of training requests across all tasks in the system depending on the number of available clients. The Realloc function (Algorithm 3) solves the optimization problem (6). See Appendix A for more details.
We outline our experimental setup in Section 5.1, discuss the existing baselines in Section 5.2, and compare the baselines with FedAST under varied settings in Section 5.3.
We consider image classification tasks with the MNIST [Deng, 2012], Fashion-MNIST [Xiao et al., 2017] and CIFAR-10 [Krizhevsky et al., 2009] datasets, and next character prediction with the Shakespeare [Caldas et al., 2019] dataset using the same models as in previous works [Acar et al., 2021, Yu et al., 2023b, Lecun et al., 1998]. We compare the wall-clock time required by different algorithms to reach some predetermined target test accuracy levels (see Table 1). In Appendix D.5, we present experiments with other target accuracy levels to show the consistency of our results. We also validate our results with ResNet-18, a larger model, trained for the CIFAR-100 classification task in Appendix D.4. In all experiments, we conduct three Monte Carlo runs with different random seeds and report the average results. 30% of which are available to accept new training requests, independent of the past.
Modeling Client Delays. As suggested in [Lee et al., 2018, Dutta et al., 2021, Shi et al., 2021, Zhou et al., 2022], we use shifted-exponential (exponential plus constant) random variables to model the time taken by a client to complete a local training request and return the update to the server. We pick the run-time generation parameters of each task according to real measurements on NVIDIA GeForce GTX TITAN X GPUs. To simulate hardware heterogeneity across clients, we divide them into 25% slow, 50% normalspeed, and 25% fast clients [Leconte et al., 2023]. We relegate additional implementation details to the Appendix.
We explain the synchronous and asynchronous baseline methods to which we compare FedAST:
Synchronous Simultaneous Training. The following synchronous methods differ only in client selection. 1. Sync-ST [Bhuyan et al., 2023]: randomly partition the client set across tasks at each round;
2. Sync-Bayes-ST [Zhou et al., 2022]: Bayesian optimization-based assignment of clients to tasks;
3. Sync-UCB-ST [Bhuyan and Moharir, 2022]: client selection as a multi-armed bandit problem.
In Figure 2, we first simultaneously train six CIFAR-10 models and compare the performance of all synchronous baselines and FedAST. As synchronous methods perform poorly due to a severe straggler issue, we augment them with a straggler mitigation method by aggregating only the first k client updates for each task and discarding the rest [Bonawitz et al., 2019] as the default option. We choose k = 30 by validation experiments across datasets in Appendix D.1. This extra augmentation makes the baselines more competitive. In Figure 2, we also add the result of Sync-ST -NoStrag.Mit., which is the Sync-ST Asynchronous Federated Simultaneous Training. To our knowledge, [Chang et al., 2023] is the only other work that mainly studies asynchronous simultaneous FL. However, their client selection scheme requires the knowledge of network-wide staleness and smoothness constants, which are hard to estimate. If the tasks have similar model complexity and task difficulty, their client selection is similar to that of FedAST with a buffer size of 1. We thus include this no-buffer version of FedAST (we call it FedAST-NoBuffer) as a baseline.
We assess the performance of FedAST under various scenarios. In homogeneous-task experiments, where multiple independent copies of the same model are trained simultaneously using the same dataset, we report the average accuracy over time. In heterogeneous-task experiments involving differing tasks and models, efficiently distributing resources to accelerate the completion of all tasks is the main challenge.
For homogeneous tasks, we use FedAST with the static option (S) and uniform client distribution across tasks. In heterogeneous-task experiments, we use dynamic allocation (option = D) to enhance resource allocation efficiency. To show the benefits of dynamic allocation over static allocation, we also explore heterogeneous-task scenarios with option = S. Dynamic allocation reduces overall training time by up to 11.9%, with comprehensive results shown in Appendix D.6.
To quantify the time saved by using FedAST over some competing baseline, we define time gain as
where T Baseline (T FedAST ) is the simulated time for Baseline (FedAST) to reach the target accuracy.
Comparison with All Synchronous FST Methods. First, we compare the synchronous baselines discussed in Section 5.2 on the CIFAR-10 dataset (Figure 2), where we simultaneously train six identical models. We observe that synchronous methods without straggler mitigation converge very slowly. Among the stragglermitigated synchronous variants that we implement, Figure 2 shows that Sync-Bayes-ST, has similar performance to Sync-ST because it struggles due to the large search space of the optimization problem, stemming from the exponential number of possible client schedules. Further, we do not observe any performance gains from using Sync-UCB-ST over Sync-ST. Given that Sync-Bayes-ST and Sync-UCB-ST have similar performance as Sync-ST, in subsequent experiments, we choose Sync-ST as the sole synchronous baseline.
Need for Buffer. As discussed earlier, incorporating the buffer mitigates the negative impact of highly stale updates. Since staleness increases with the number of active clients, asynchronous FL methods without a buffer exhibit limited scalability as the number of clients grows.
To demonstrate this, in Figure 3, we conduct two experiments: 1) training three models for CIFAR-10 simultaneously, and 2) training three models for Fashion-MNIST simultaneously. We plot the final accuracy values varying the ratio of the active clients. We observe that for small active client ratios, FedAST and FedAST-NoBuffer have comparable performance. However, with more active clients, the staleness of the updates increases, resulting in significantly worse performance of the fully asynchronous FedAST-NoBuffer algorithm. Then, in Figure 4, we simultaneously train two models, one each for CIFAR-10 and Fashion-MNIST. Observing the unsteady learning curves of FedAST-NoBuffer, we conclude that the buffer makes the system more robust to stale updates.
Comparison with Sync-ST. Next, we compare FedAST with the chosen synchronous method, Sync-ST. 1) Homogeneous Tasks: We conduct experiments training 2, 4, and 6 identical models for each of MNIST, Fashion-MNIST, CIFAR-10, and Shakespeare datasets. Figure 5 shows the average finish times of the algorithms, and the significant time gains of our algorithm FedAST over synchronous Sync-ST (even after incorporating straggler mitigation). We observe that the gain increases with the number of simultaneously trained tasks because Sync-ST is especially vulnerable to the straggler problem.
2) Heterogeneous Tasks: The heterogeneous experiment trains 4 models simultaneously, one each for the MNIST, Fashion-MNIST, CIFAR-10, and Shakespeare datasets.
Once one model reaches its target accuracy, its training stops, and its clients are reallocated to other tasks. We use FedAST with option = D for dynamic client allocation.
For the synchronous baseline Sync-ST, we ran 30 different client allocation schemes, including our proposed allocation scheme and uniform allocation across tasks. We report the results achieved with the best-performing scheme.
Figure 6 shows learning curves for FedAST and Sync-ST from a single Monte Carlo run. The dashed vertical lines denote the time instants when a model reaches its target accuracy, following which the clients training this model get reallocated to other tasks. Figure 7 shows the average finish times for 4 simultaneously trained models with FedAST and Sync-ST. For example, with FedAST, the model for MNIST dataset hits its target accuracy at 99, after which the clients training this model get reallocated to the other models. At 619, the training of the final model (for CIFAR-10) is complete. Comparing the finish times for the last model, FedAST provides a 40.1% time gain over Sync-ST.
We observe that thanks to dynamic client allocation, FedAST automatically detects which tasks have higher heterogeneity across clients and need a larger buffer. We notice that the Shakespeare task is allocated fewer clients because it is estimated to be less heterogeneous, which is true based on the label distribution of data samples across clients. We also repeat this experiment (Appendix D.6) using FedAST with static option (S) to validate the proposed dynamic client allocation strategy. Dynamic client allocation consistently has time gain up to 11.9% compared to the static version.
In this paper, we present FedAST, a federated learning framework to simultaneously train multiple models using buffered asynchronous aggregations. We theoretically prove
Further, given the limited number of available clients, we cannot increase the total number of active local training requests arbitrarily without increasing the staleness. Thus, we employ M m=1 R m = R where R is a constant of how many active training requests we assign in total depending on the number of available clients in the setting. Since the goal of federated simultaneous training is to minimize the objective functions of all tasks concurrently, we propose to adjust {R m } M m=1 by solving,
The solution of the minimization problem in (7) suggests allocating local training requests, {R m } M m=1 , in proportion to σ g,m √ η s m η c m τ m for each model m. Using this approach, Algorithm 3 adjusts resource allocation. Further, as the update variance across clients may vary in time during training, we employ adaptive periodical reallocation of resources across models (Line 1 in Algorithm 3). Therefore, we use round indices to denote the changing number of active training requests and buffer sizes.
As we do not have access to true data heterogeneity levels, we need to estimate it (Line 2 in Algorithm 3). When FedAST is run with option = D (dynamic client allocation option), the server keeps the latest V updates of each model. This requires constant and small memory space kept in the server. To present how variance estimation (EstimateVariances()) works, assume that {∆ m,1 , ∆ m,2 , . . . , ∆ m,V } M m=1 are the sets of latest received updates of tasks m ∈ [M ] where each ∆ m,k for k ∈ [V ] is the output of k th latest local training (Algorithm 1). As the output of any local training is the average of all calculated stochastic gradients during that local training, we use those outputs as approximations of the gradients calculated on the local data of clients. Algorithm 4 describes EstimateVariances(). It first calculates the mean of the latest updates for each task,
Then, EstimateVariances() returns sample variance multiplied with other terms (η s m η c m τ m ) suggested by ( 7) and normalized by the mean update norm (to prevent large models or models with inherently large weights from dominating others),
Then, Realloc algorithm allocates the number of active training requests proportionally to the square root of these values (Algorithm 3, Line 3). This approach is sensible both theoretically and intuitively. Based on our experimental observations regarding the relationship between the number of active training requests and buffer size, increasing the number of active local training requests necessitates an increase in buffer size. Moreover, a larger buffer proves beneficial in reducing the variance across updates, Algorithm 4 EstimateVariances() Require: The set of latest V updates {∆ m,1 , ∆ m,2 , . . . , ∆ m,V } M m=1 , server-side learning rates {η s m } M m=1 , client-side learning rates {η c m } M m=1 , and the number of local SGD steps of all models {τ m } M m=1 .
▷ Calculate the means of the latest updates
▷ Multiply with other constants suggested by the convergence guarantee (7) 4: Return {σ 2 g,m }
the buffered updates are averaged during the aggregation. Realloc aims to allocate more clients and provide a larger buffer size for tasks with higher heterogeneity. We choose the number of stored latest updates V = 8 and the period of number of total updates from all clients to trigger reallocation in Realloc subroutine c period = 0.75 × M × M m=1 R m in our experiments. The benefits of dynamic allocation (option = D) over static and uniform resource allocation (option = S) are demonstrated when tasks/models are heterogeneous, as shown in Figures 212223242526.
We compare FedAST with single-model FL methods, too. [Nguyen et al., 2022] is the most similar algorithm to FedAST (with single-model). However, even for the single-task case, FedAST differs by employing a uniform client assignment to ensure unbiased participation of clients irrespective of their hardware speeds. This allows us to relax the assumptions to prove the convergence guarantee. [Nguyen et al., 2022] relies on a strong assumption that the server receives updates from clients uniformly at random and that the norm of gradients is bounded. Moreover, compared to [Koloskova et al., 2022], our analysis is more general as FedAST uses multiple SGD steps in local training and a buffer. Some other recent single-model asynchronous FL works, [Zakerinia et al., 2023] and [Leconte et al., 2023], do not have straightforward and efficient simultaneous federated training extensions for multiple models.
[ Chang et al., 2023] is another asynchronous simultaneous federated learning method. However, [Chang et al., 2023] indeed fails to converge to a stationary point asymptotically unless data is homogeneous, and their assumptions include Bounded Gradient Norm and Weak Convexity.
In our study, we explore a simultaneous federated learning (FL) setting for multiple models. We present the details of our experiments in this section.
We simulate the training with PyTorch on NVIDIA GeForce GTX TITAN X graphics processing units (GPUs) of our internal cluster. We build our code upon the public codes of [Wang et al., 2020a, Yu et al., 2023b].
We consider the federated training of M models simultaneously using N clients. N is 1000 in all experiments and M , specified for each experiment explicitly, varies between 2 -6.
We
We consider the data heterogeneity across clients in FL frameworks. We download MNIST, Fashion-MNIST, and CIFAR-10/100 datasets from PyTorch built-in library methods. The train and test splits provided by the library are used without any modifications. To simulate heterogeneous data distribution across clients, we use Dirichlet distribution with α = 0.1 following the approach suggested in Yurochkin et al. [2019]. We ensure that each client has 300 data points for MNIST, Fashion-MNIST, and CIFAR-10/100 tasks by repeating the train set if necessary. We obtain and preprocess the Shakespeare dataset as described in [Caldas et al., 2019]. This dataset has inherently heterogeneous distribution across clients as each client corresponds to a unique role from Shakespeare's plays.
In this section, we explain how we choose the design parameters.
Client dataset sizes, batch sizes, and number of local steps. While distributing CIFAR-10/100, MNIST, and Fashion-MNIST datasets across clients, each client is allocated 300 data points from each dataset. The Shakespeare dataset, however, maintains its original distribution of data points across roles, so clients have different numbers of data samples in the Shakespeare task. For CIFAR-10/100, MNIST, and Fashion-MNIST tasks, we set the batch size to 32 while we employ a batch size of 64 for the Shakespeare task. We fix the number of local steps in local training (τ m parameter in Algorithm 1 in the main text) of clients at 27 for all tasks. This makes 3 epochs for CIFAR-10/100, MNIST, and Fashion-MNIST tasks. As the number of data points varies across clients for the Shakespeare dataset, there is no fixed number of epochs. m ) when we use dynamic client allocation option (D). A higher number of simultaneous local training requests leads to a higher staleness because it increases the global model's update frequency at the server. On the other hand, buffer size is inversely related to staleness, given its opposing effect on the aggregation frequency. Based on our experimental observations, selecting the number of active training requests and the buffer size of model m such that their ratio is fixed and below 37, (R m /b m ≲ 37 or R (tm) m /b tm m ≲ 37), works well. Selecting the buffer size of FedAST based on this observation avoids the detrimental effects of stale updates while benefiting from fast training thanks to the asynchronous algorithm. We show two experimental results in Figures 8 and9. In Figure 8, we train one Fashion-MNIST and one CIFAR-10 models simultaneously by assigning 175 active training requests to each task and observe that buffer size of 5 strikes a balance between high final test accuracy and fast training to achieve the target accuracy for both tasks. In Figure 9, we repeat a similar experiment with MNIST and CIFAR-10 tasks by assigning 105 active training requests to each. This time, we observe that a buffer size of 3 performs the best for both tasks. These experimental results support our buffer size choice. Learning rate and weight decay. We search for the best learning rate and weight decay hyperparameters considering the training speed and final accuracy levels. We seek client-side learning rate within the range of [1 × 10 -3 , 1 × 10], server-side learning rate within [3 × 10 -2 , 3], and weight decays within [1 × 10 -7 , 1 × 10 -2 ]. We observe that client-side learning rates of 6 × 10 -2 and 7 with weight decays of 3 × 10 -4 and 7 × 10 -5 work best respectively for Fashion-MNIST and Shakespeare tasks for all methods. For CIFAR-10 task, a client-side learning rate of 1 × 10 -1 with weight decays of 7×10 -4 and 3×10 -4 perform best for asynchronous and synchronous methods, respectively. For MNIST, we use client-side learning rates of 1 × 10 -1 and 2 × 10 -1 for asynchronous and synchronous methods, respectively, with a weight decay of 3 × 10 -4 . For server-side learning rates, we observe that 1 for synchronous methods (Sync-ST, Sync-Bayes-ST, and Sync-UCB-ST), 0.1 for FedAST, and 0.038 for FedAST-NoBuffer perform well for all tasks.
In our experiments, following [Lee et al., 2018, Shi et al., 2021, Zhou et al., 2022, Dutta et al., 2021], we employ the shifted-exponential random variables to model the duration between when the server sends a local training request to a client, and when it receives the update of the local training. The exponential component of the distribution reflects the stochastic nature of the device speeds, while the shift component accounts for unavoidable delays such as disk I/O operations.
Whenever a client i performs local training for task m, we draw a random number from the distribution with a cumulative distribution function (CDF) of,
where β i,m depends on the speed of client i and the size of the model associated with task m. Then, we multiply this random number by the number of local steps to calculate the simulation time between when the server requests for the local training, and when it receives the update back.
We quantify the effect of the model sizes based on the average time required to calculate one stochastic gradient for each model on the GPUs of our internal cluster. By our measurements, we set,
In our experiments, we also take the heterogeneity in the speed of client devices into consideration. We categorize clients into three speed groups: slow (%25), normal-speed (%50), and fast (%25). The speed rates for these categories are inversely proportional to 1.3, 1, and 0.7, such that,
In this section, we present supplementary experiments.
In our experiments, to mitigate the high straggler effect, the server in synchronous methods (Sync-ST, Sync-Bayes-ST, and Sync-UCB-ST) only aggregates the first k client updates for each task and discards the rest, following [Bonawitz et al., 2019]. To tune parameter k, we run validation experiments with Sync-ST on single CIFAR-10, MNIST, and Fashion-MNIST tasks and evaluated the training performance with respect to simulated time and number of global rounds.
A larger k results in a longer simulated time per round since we wait for more clients. On the other hand, the variance in aggregated updates on each round becomes smaller since we average more updates. Therefore, the target accuracy is attained faster in terms of the number of global rounds. We also observed that keeping k too small yields lower final accuracy. Navigating these trade-offs, we find that k = 30 strikes an effective balance.
We illustrate test loss plots of the experiments in Figures 3 and4 in the main text.
The ratio of active clients
In Figure 15, we provide the average training curves of the homogeneous-task experiment in Figure 5.
To see how FedAST and the competitor Sync-ST work with different target accuracy, we conduct the experiment in Figure 7 with +3% higher and -10% lower target accuracy levels as presented in Table 3. We observe that proposed FedAST reduces the overall training time by 55.9% and 16.3%, respectively for higher and lower target accuracy levels.
We conclude that the advantage of FedAST over Sync-ST increases with the difficulty of the task (i.e., reaching higher accuracy). 3. Dashed vertical lines show times when tasks reach their target accuracy. The setting is the same as the experiment in Figure 7. 3. The setting is the same as the experiment in Figure 7. 3. Dashed vertical lines show times when tasks reach their target accuracy. The setting is the same as the experiment in Figure 7. 3. The setting is the same as the experiment in Figure 7.
We conduct heterogeneous-task experiments to validate the performance gain of dynamic resource allocation (FedAST (D)) over static option (FedAST(S)) with uniform allocation across tasks in heterogeneous settings. For uniform resource allocation, we allocate the same number of active training requests to each task in FedAST(S). To show the consistency of our results, we run experiments at all target accuracy levels in Table 3. We present the results in Figure 22 (higher target accuracy), Figure 24 (the target accuracy in the main text), and Figure 26 (lower target accuracy). We conclude that our dynamic client allocation based on the variance estimates of the updates reduces the total training time compared to the uniform static client allocation. The advantage of dynamic resource allocation becomes more prominent with more difficult tasks (i.e., higher target accuracy level). 3.
In this section, we present the proofs of the mathematical claims made in the paper. First, we define and explain the notations used in this section. After that, we introduce intermediate lemmas used in the main proof (Section E.2). Then, we present the proofs of Theorem 1 and Corollary 1.1 (Section E.3). Finally, we prove intermediate lemmas (Section E.4). ). The setting is the heterogeneous experiment with lower target accuracy levels in Table 3.
FedAST enables us to divide the convergence analyses of simultaneous tasks into individual ones. We focus on the convergence analysis of a single task within a simultaneous multi-model setting and the analysis holds for all tasks trained together. For brevity, we provide the proofs for a single task of multiple models trained simultaneously. Therefore, we drop all model indices in our analysis. We also drop time indices from the number of active requests (R) and buffer size (b) terms as they remain the same during the training with the static option of FedAST. Table 4 summarizes all notation. Please note that the analysis presented here holds for every model m ∈ [M ] simultaneously trained within FedAST framework.
We first revisit the local training and global update rules of FedAST. The notation may vary slightly from those in the main paper due to dropping model indices, but still accurately depicts the same algorithmic procedures, Algorithms 1 and 2 in the main text.
Local update rule. During local training, clients perform τ consecutive local stochastic gradient steps and return the output to the server. When a client receives the t th version of the global model, x (t) , it takes τ mini-batch stochastic gradient descent steps (for k = 1, . . . , τ ) with following rule:
where t) and ∇ denotes stochastic gradients. We define the average of local stochastic gradients as
. Then, the client returns
i to the server. The server stores the updates in a buffer.
Staleness. The server receives the updates of local training requests asynchronously. It means that the received updates may come in a different order than local training requests sent to clients. Therefore, an aggregated update may have been calculated with an older version of the model, and this is called staleness. We quantify the staleness of an update in terms of the number of global rounds passed between the times when the server sends the local training request and receives the update. The staleness is random for each update, depending on client selections for all tasks and all clients' availability, computation, and communication speeds. We denote the staleness of client i's update received at the server at the t th round as γ t i . Recall that Assumption 4 (Bounded Staleness) bounds this random value above at γ max .
Global update rule. On each global round t, when the buffer at the server, B (t) , is full (|B (t) | = b, where b is the buffer size), the server aggregates the updates to proceed to the next global round. Here, B (t) is the set of clients whose updates are received after (t -1) th and before t th aggregation. The aggregation rule can be written as follows:
We utilize the perturbed iterate idea from Koloskova et al. [2022], Mania et al. [2017].
First, let us introduce some helpful sets and notations. Consider A (t) , which represents the set of clients chosen by the server to receive the t th version of the model. Recall that the server in FedAST selects the clients uniformly at random with replacement from all clients. The size of this set, |A (t) |, is always equal to the buffer size, b, (except initialization, t = 0) because b new local training requests are made on each round. For instance, if b is set to 3, and the server selects the 2 nd , 16 th , and 31 st clients during the 4 th aggregation round, then A (4) is {2, 16, 31}. The server sends x (4) to the 2 nd , 16 th , and 31 st clients and requests local training with this model. In practical terms, A (t) is a multiset, allowing multiple occurrences of the same client if a client is selected more than once. Throughout the proof, we consider each occurrence of the same client in multiset as a distinct update calculated on that particular client. While we acknowledge a slight abuse of notation, this does not lead to any mathematical flaw, and we believe that this significantly enhances the clarity and comprehensibility of the proof. Now, let us define C (t) as the set of clients that have incomplete local training requests at the time of the t th aggregation because of the asynchronous nature of FedAST. The size of this set, |C (t) |, is always equal to the number of active local training requests, R, because the server sends a new local training request for every update it receives. For instance, if R is 4, and the server has sent local training requests to the 12 th , 27 th , 41 st , and 55 th clients prior to the 5 th aggregation, yet these clients are still processing their updates, then C (5) would be {12, 27, 41, 55}. Note that C (0) is an empty set, as there are no active local training requests before the algorithm starts. It is worth noting that C (t) is a multiset, allowing multiple occurrences of the same client if a client has more than one active local training request (recall that multiple requests are queued at the client side). Each occurrence of a client within this multiset represents a different local training calculated on that client. We again acknowledge a slight abuse of notation, but this does not lead to any mathematical flaw, and we believe that this makes the flow of proof significantly easier.
Next, we define the virtual sequence z (t) for t = 0, 1, . . . , T as the model that receives local training updates of the global model, x (t) for t = 0, 1, . . . , T , in the correct order. Namely, unlike x (t) , z (t) receives the local training updates in the order in which the server sends those requests. However, it is crucial to note that the local training updates are still calculated with the global model, x (t) . The update rule of the virtual sequence is:
for t = 0, 1, . . . , T -1 where z (0) ≜ x (0) .
Remark 1. Now, we state an observation using the definitions of C (t) , the virtual sequence, and the global model. When the t th aggregation happens at the server, the virtual sequence, z (t) , has received all the updates from all previous local training requests on rounds 0, 1, . . . , t -1. At the same time, the global model, x (t) , has received the same updates except for the updates of clients in C (t) . By the update rules in ( 9) and ( 10), note that each received update at the server contributes to the global model and virtual sequence equally. Therefore, we can express their difference as:
Remark 2. If we count the number of occurrences of the round index y among the model versions of assigned updates in set C (t) over all rounds t = 0, . . . , T -1, we can bound this value as:
where 1 is an indicator function that returns 1 if the statement is true, and returns 0 otherwise. The reasoning for this observation is as follows. On each round, the server selects b clients (since the server selects one new client for each received and buffered update where the buffer size is b) and sends them the up-to-date global model. We also know that all local training requests must be returned to the server within γ max rounds by Assumption 4 (Bounded Staleness). Therefore, over the rounds t = 0, . . . , T -1, any round indices can appear at most bγ max times in the summation in the left-hand side of the inequality in ( 12). We will use this remark later in the proof.
We define some useful variables used in the proof and present the notation used in FedAST in Table 4. Also, we again want to remind the reader that we dropped all model indices in the proof as the theoretical results we present here hold for any of multiple tasks trained simultaneously, satisfying Assumptions 1 -4.
The set of clients which are requested local training, but have not returned their updates to the server yet (Section E.1.2)
We present intermediate lemmas used through the proof.
Lemma 1. For a set of Q vectors, u 1 , . . . , u Q , where Q is a positive integer,
Proof. The lemma is a direct consequence of Jensen's inequality with a convex function ∥•∥ 2 and uniform random distribution over the set of vectors u 1 , . . . , u Q .
Lemma 2. Suppose that f i (•) satisfies Assumption 1 (Smoothness) and Assumption 2 (Bounded Variance) for all i ∈ [N ], and assume that η c ≤ 1 Lτ . Then the iterates of FedAST satisfy,
where D ≜ L 2 η 2 c τ (τ -1). Further, suppose Assumption 3 (Bounded Heterogeneity) holds. Then, the iterates of FedAST satisfy,
Remark 3. The true gradient at any client using the global model is close to the local update of that client.
Lemma 3. The iterates of FedAST and defined virtual sequence satisfy,
Lemma 4. Suppose that f i (•) satisfies Assumption 2 (Bounded Variance and Unbiased Stochastic Gradients) for all i ∈ [N ], then the iterates of FedAST satisfy,
Remark 4. The noisy global update due to stochastic gradients is close to the expected update calculated with full gradients. The buffer and multiple local steps are useful to reduce the variance due to local SGD steps.
Lemma 5. The iterates of FedAST satisfy,
Further, suppose Assumption 3 (Bounded Heterogeneity) holds. Then, the iterates of FedAST also satisfy,
Remark 5. FedAST benefits the global variance reduction thanks to the buffer.
Lemma 6. The virtual sequence and the iterates of FedAST satisfy,
Remark 6. As discussed in Remark 1, although the virtual sequence and global model get updates in a different order, they receive the same updates. Therefore, we can bound their difference.
We present and prove Theorem 1 and Corollary 1.1 here.
First we restate the theorem:
Proof. Using the update rule of the virtual sequence (10) and Assumption (Asm.) 1 (Smoothness), and taking the conditional expectation with respect to z (t) , we have,
Using Lemmas 3 and 4, we can bound T 1 and T 2 . Then, dividing both sides by η s η c :
Using Lemma 5, we can bound T 3 :
where D ≜ L 2 η 2 c τ (τ -1). Using the tower property of conditional expectation, telescoping the inequality over the round indices t = 0, 1, . . . , T -1, and using Lemma 6, we get,
Suppose the learning rates satisfy η s ≤ √ τ b (which also makes η s ≤ τ √ τ b) and η c ≤ min
τ Rγ max , the following inequality holds:
Also, notice that z (0) is equal to x (0) by definitions (Section E.1.2) of these sequences and min x f (x) ≤ f z (T ) .
Define δ ≜ f x (0) -min x f (x). After reducing high-order terms using the assumptions, η s ≤ √ τ b (which also makes
τ Rγ max , and incorporating the constants into the O(•) notation, we have:
This concludes the proof. T . Then, the bound in Theorem 1 reduces to:
Proof of Lemma 2. We borrow the proof technique from [Wang et al., 2020b, C.5].
Using the recursive appearances of T recursive in ( 14) and ( 15):
Arranging the terms, defining D ≜ L 2 η 2 c τ (τ -1), and assuming η c ≤ 1 Lτ which makes D ≤ 1,
This proves the first part of Lemma 2. Now, averaging it across clients:
This concludes the proof of Lemma 2.
Proof of Lemma 3.
Proof of Lemma 4.
where (a) follows that the clients in A (t) are selected uniformly at random with replacement among all clients (see Section E.1.2), and (b) follows that t) . This proves the first part of Lemma 5.
Proof of Lemma 6. We start by using Assumption 1 (Smoothness) and Remark 1.
Telescoping the inequality over t = 0, . . . , T -1: Global update rule and virtual sequence definition. Although the local update rule remains the same, the global update rule slightly changes for dynamic client allocation due to changing buffer size: t) . Note that ( 17) is almost identical to (9), except the varying buffer-size b (t) . Next, we define r i (t) as the index of the global round when a local training request sent to client i in round t returns to the server. Basically, it is the current round index t, added to the future value of staleness that the requested update will have. We need to define a new virtual sequence y (t) , which is different from the z (t) defined earlier.
for t = 0, 1, . . . , T -1 where y (0) ≜ x (0) . Here, A (t) is defined similarly as it was in Section E.1.2. Note that the probability of being in A (t) is equal across clients due to uniform client selection. However, this time, the size of this set does not have to be equal to the buffer size at round t. Due to the new client selection rule (Line 13 in Algorithm 2), the server may assign 0, 1, or 2 clients for each received update. Therefore, we know that 0 < |A (t) | ≤ 2b (t) .
Here, we need a simplifying assumption for the purpose of this proof:
Assumption 5 (b ri(t) Values). We assume that any b ri(t) value is known at the time when a local training request is sent to client i at round t, and these values are independent of any future information including the received updates. We further assume that b ri(t) values are equal (denote b r(t) ) for all clients in A (t) .
Remark 7. When we keep the period of dynamic client allocation long enough, we observe that most of the assigned local training requests at one round fall in the same window before the next dynamic client allocation happens (Line 8 in Algorithm 2). Hence, based on our empirical observations, what assumption implies holds for most of the local training requests. Further, this assumption can be avoided by taking an average of the updates during aggregation weighted inversely with the number of local training requests sent at the same global round. In other words, one may have avoided this assumption by weighting an update from client i with 1/|A (t-γ t i ) | instead of taking average over buffer during aggregation at round t. However, we did not see any practical benefit of this type of weighting in our experiments, and this strange weighting would be just for theoretical purposes. Therefore, we keep the current version.
We first state the theorem showing the convergence of FedAST with dynamic client allocation option.
Theorem 2. (Convergence of FedAST with option = D): Suppose Assumptions 1 -5 hold, and the learning rates satisfy
Then, the iterations of Algorithm 1 (FedAST) with option = D satisfy:
Proof.
We will need one extra lemma corresponding to Lemma 6.
Lemma 7. The new virtual sequence y (t) and the iterates of FedAST satisfy,
Now, using the update rule of the virtual sequence (18) and Assumption 1 (Smoothness), and taking the conditional expectation with respect to y (t) , we have, E f y (t+1) ≤ f y (t) + ∇f y (t) , E y (t+1) -y (t) + L 2 E y (t+1) -y (t) 2 = f y (t) + ∇f y
(Using Assumption 5)
(
Next, using Lemma 3 (with y (t) sequence) and Lemma 4 (with |A (t) |), using 1/ρ b ≤ |A (t) |/b r(t) ≤ 2ρ b , and dividing both sides by η s η c we obtain,
Using Lemma 5, we get,
where D ≜ L 2 η 2 c τ (τ -1). Using tower property of conditional expectation, telescoping the inequality over the round indices t = 0, 1, . . . , T -1, and using Lemma 7, we get,
Suppose the learning rates satisfy η s ≤ ρ τ Rγ max , the following inequality holds:
Also, notice that y (0) is equal to x (0) by definitions (Section F) of these sequences and min x f (x) ≤ f y (T ) .
T -1 t=0 E ∇f x (t) 2 ≤ 11 f x (0) -min x f (x) T η s η c ρ b (Using ( 19))
Define δ ≜ f x (0) -min x f (x). After reducing high-order terms using the assumptions, η s ≤ ρ This concludes the proof.
Proof of Lemma 7: We start by using Assumption 1 (Smoothness) and observing that Remark 1 still holds with y (t) for the dynamic client allocation option.
Telescoping the inequality over t = 0, . . . , T -1:
(Using Remark 2, however, this time,)
(the maximum appearance can be 2γ max b max )
(Using Remark 2)
Our code is provided at https://github.com/ askinb/FedAST.
Processing the requests in parallel would require clients to keep all the M models in local memory, which can be infeasible.