FL_PyTorch: optimization research simulator for federated learning

Federated Learning (FL) has emerged as a promising technique for edge devices to collaboratively learn a shared machine learning model while keeping training data locally on the device, thereby removing the need to store and access the full data in the cloud. However, FL is difficult to implement, test and deploy in practice considering heterogeneity in common edge device settings, making it fundamentally hard for researchers to efficiently prototype and test their optimization algorithms. In this work, our aim is to alleviate this problem by introducing FL_PyTorch : a suite of open-source software written in python that builds on top of one the most popular research Deep Learning (DL) framework PyTorch. We built FL_PyTorch as a research simulator for FL to enable fast development, prototyping and experimenting with new and existing FL optimization algorithms. Our system supports abstractions that provide researchers with a sufficient level of flexibility to experiment with existing and novel approaches to advance the state-of-the-art. Furthermore, FL_PyTorch is a simple to use console system, allows to run several clients simultaneously using local CPUs or GPU(s), and even remote compute devices without the need for any distributed implementation provided by the user. FL_PyTorch also offers a Graphical User Interface. For new methods, researchers only provide the centralized implementation of their algorithm. To showcase the possibilities and usefulness of our system, we experiment with several well-known state-of-the-art FL algorithms and a few of the most common FL datasets.


INTRODUCTION
Over the past few years, the continual development of DL capabilities has revolutionised the way we interact with everyday devices. Much of this success depends on the availability of largescale training infrastructures such as large GPU clusters and the ever-increasing demand for vast amounts of training data. In contrary to this, users and providers are becoming increasingly aware of their privacy leakage because of this centralised data collection, leading to the creation of various privacy-preserving initiatives by industry service providers [2] or government regulators [10].
Recently, a viable solution that has the potential to address the aforementioned issue is Federated Learning (FL). FL a term was initially proposed in [26] as an approach to solving learning tasks by a loose federation of mobile devices. However, the underlying concept of training models without data being available in a single location is applicable beyond the originally considered scenario and it turns out to be useful in other practical use-cases. For example, learning from institutional data silos such as hospitals or banks which cannot share data due to confidentiality or legal constraints, or applications in edge devices [18,42].
The main goal of FL is to provide strong privacy protection which is obtained by storing data locally rather than transferring it to the central storage. To solve the underlying machine learning objective each client provides focused updates intended for immediate aggregation. Stronger privacy properties are possible when FL is combined with other technologies such as differential privacy [9] and secure multiparty computation (SMPC) protocols such as secure aggregation [4].
To enable research in federated learning, several frameworks have been proposed including LEAF [7], FedML [15], Flower [5]), (PySyft [37], TensorFlow-Federated (TFF) [19], FATE [42], Clara [29], PaddleFL [25], Open FL [20]. These frameworks are mainly built with a focus to be deployed on real world systems while also providing user with an ability to run experiments with the same code. This desired property often comes with the price that the entry bar for researchers to extend or experiment with these frameworks is limited in a sense that they either need to have extensive experience with distributed systems or require assistance from experts on given framework.
In our work, we decided to take one step back and focus on the construction of a framework that, although not aimed for being deployed to edge devices as primary goal, can provide a useful simulation environment for researchers with the following goals: • Low Entry Bar/ Simplicity: We aim our tool to be as simple as possible for the user while providing all necessary functionalities.
• Extensibility: It is easy to bring your own algorithm or dataset or extend existing ones. We aim to achieve this by providing universal abstractions with a sufficient level of flexibility to experiment with existing and novel approaches to advance the state-of-the-art. • Hardware Utilisation: Cross-device FL experiments are usually of much smaller scale comparing to the centralized setting. This is mainly because of limited device capabilities. Running such experiments on GPU can lead to the under-utilisation of available hardware. We aim to resolve this via the ability to parallelise clients' computation. • Easy Debugging: Debugging multi-process or multi-thread systems is hard. We only require user to provide a single thread implementation which is automatically adjusted to multi-GPU and multi-node setup.
To the best of our knowledge, there is no such tool that could provide a sufficient level of freedom and is simple to use and therefore, we design FL_PyTorch to achieve the aforementioned goal.
FL_PyTorch is an optimization research simulator for FL implemented in Python based on PyTorch [32]. FL_PyTorch is a simple to use tool with its own Graphical User Interface (GUI) implemented in PyQt [33]. During the simulation process, the selected local CPUs/GPUs are accessed in a parallel way. In addition, there is a possibility to use remote CPUs/GPUs. Remote devices are required to have a TCP transport layer for communicating with the master. Regarding supported devices, we target server and desktop stations running on Linux, macOS, or Windows OS for efficient simulations.
For the paper organization, we introduce the general FL minimization objective in Section 2. Subsequently, we provide a deep dive into our FL_PyTorch system in Section 3 and, finally, we demonstrate FL_PyTorch capabilities by concluding several experiments on multiple FL baselines and the most used FL dataset in Section 4.

FL OBJECTIVE AND FEDAVG
In this section, we introduce the FL objective in its general form The global objective is an expectation over local objectives over the randomness inherited from the client distribution P, and the local objectives { } have the form of an expectation over the local datasets {D }. The { } has an additional regularization term that is useful for incorporating prior knowledge of parameters to find .
In the case of the finite number of clients and local data points, both global objectives and local losses { } can be written as simple weighted averages in the empirical risk minimization form (ERM). The most common algorithm to solve (1) is federated averaging FedAVG [26]. This algorithm divides the training process into communication rounds. At the beginning of the -th round ( ≥ 0), the server broadcasts the current global model ( ) to a random subset of clients S ( ) (often uniformly sampled without replacement in simulations). Then, each sampled client performs local SGD updates on its own local dataset and sends its local model update Δ ( ) = ( , ) − ( ) to the server. Finally, the server uses the aggregated model updates to obtain the new global model as follows: where is the relative weight of client . The above procedure will repeat until the algorithm converges. In the cross-silo setting where all clients participate in the training at every round, we have S ( ) = {1,2, . . . , }.

FL OPTIMIZATION SIMULATOR
FL_PyTorch is a system that has been built using Python programming language, and it is based on the DL framework PyTorch. We have made the code repository publicly available 1 . Its backbone consists of a general form of FedAVG displayed in Algorithm 1 partially inspired by Algorithm 1 in [35]. Our proposed general form preserves the standard structure of federated optimization where in each round , subset S ( ) of all available clients is selected. As a next step, the master generates its state which is broadcasted together with the current copy of the global solution ( ) to the selected subset of clients S ( ) . Afterwards, each participating client initializes its local solution ( ,0) to be a copy of the global solution ( ) . Each of these clients then performs steps using the local optimizer ClientOpt with gradient estimated by the LocalGradient function. After this step, each client computes the model and state update, which are then sent back to the master which estimates the global update direction using the ServerGradient function. This estimate is used in ServerOpt that updates the global solution to its new value ( +1) . Lastly, the server global state is updated. In Section 3.2, we show that this general scheme captures all standard algorithms, thus our scheme is sufficient and we believe that it gives researchers a sufficient level of freedom to develop and experiment with each component of our general scheme to push both practical and theoretical FL state-of-the-art.
Our current implementation is fully determined and allows us to configure up to 51 parameters which can be specified either through our easy to use GUI tool or directly via the command line that we discuss in the next subsection. These parameters can be grouped into 4 categories based on their function: • Server Optimizer: the number of communication rounds ; the number of sampled clients per round; global initial local learning rate, global learning rate schedule, and global optimizer with its parameters such as momentum, the name of the algorithm to execute. • Local Optimizer: the number of local epochs or local iterations ; batch size for data loading; local optimizer with its parameters. • Model and Data: model's and dataset's names.
• System Setup: directory to store run metadata, target compute devices, usage of remote compute devices, logging level, number of workers for loading dataset, random seed, thread pool sizes, experiments grouping, user defined comment, enable usage of NVIDIA Ampere GPU Tensor Cores, optional cleanup of PyTorch GPU cache at the end of each round, device for store client state 2 .

Frontend
As we mention previously, FL_PyTorch supports two modes for running federated optimization-Graphical User Interface and Command Line Interface. Below, we discuss both of these in detail.

Graphical User Interface (GUI).
We implemented our Graphical User Interface using the PyQt5 GUI framework [33]. This GUI framework supports all the standard desktop operating systems such as Windows, Linux, or macOS. In addition, the GUI part of FL_PyTorch has built-in VNC server support. If the target machine has not a native windows manager system one can use this mode and connect to GUI part via VNC Viewer software.
To get a better picture of our GUI, we display its main menu provided in Figure 1,  Some extra configurable displayed information from GUI includes a number of currently executing experiments, showing the available size of physical DRAM and swap memory in OS.
Some optimization algorithms require a client to store in its internal state-specific information about algorithm-specific shifts that have a notion about gradient estimator. The analysis is very often irrelevant for the specific strategy of initializing shifts, but in practice, different initialization policies may have consequences in convergence speed. We provide two policies with initializing shift by zero or by full gradient at (0) .
As an example, we show how Algorithm 1 is adjusted to FedAVG and SCAFFOLD. For FedAVG, both global state is an empty dictionary. Locally, we run iteration, which is usually set to be a constant for the theoretical and size of the local dataset for the experiments. LocalGradient returns the unbiased stochastic gradient at ( , ) and ClientOpt is Stochastic Gradient Descent (SGD) with given step-size . Similarly to the global state, local step update is none. ServerGradient is a (weighted) average of local updates and ServerOpt is SGD with fixed step-size 1. For the SCAFFOLD, there are 2 changes comparing to FedAVG. Firstly, the global state ( ) is a non-empty vector of the same dimension as ( ) initialized as zero and each ClientOpt has its own local ( ) , which is set to be a stochastic local gradient at ( ) . LocalGradient then returns the unbiased stochastic gradient at ( , ) shifted by ( ) − ( ) . Secondly, the local model update is sent to master together with the local state update, which is set to the difference between current and previous communication round ( ) = ( ) − ( −1) . The ServerGlobalState is then updated using the average of ( ) 's.

Supported Compressors
In federated learning (especially in the cross-device setting), clients may experience severe network latency and bandwidth limitations. Therefore, practical FL algorithms generally use a communication reduction mechanism to speed up the training procedure. Three common methods to reduce the communication cost are: (1) to reduce the communication frequency by allowing local updates; (2) to reduce communication traffic at the server by limiting the participating clients per round; (3) to reduce communication volume by compressing messages.
We naturally support the first and the second option and we also added support for the compression.
FL_PyTorch allows compressing messages both from the server to local clients and vice-versa. We support several unbiased and biased compressors. The supported unbiased compressor are: Identical compressor (no compression), Lazy or Bernoulli compressor (update is communicated with probability rescaled for the update to be unbiased), Rand-(only random coordinates are preserved uniformly at random rescaled for the update to be unbiased), Natural compressor [17], Standard dithering [1], Natural Dithering [17], TernGrad [41] and QSGD compressor [1]. The supported biased compressor are: Top- [6] and Rank- [38]. For Top-and Rank-compressors, parameter can be specified as a percentage in [0,100] from or as an absolute number greater than 1. Also, we provide means to construct new compressors via function composition and probabilistic switching from existing.   Besides DL models, we provide support for simple quadratic loss function to enable users to explore algorithms in this simplistic regime or debug their implementation. We synthetically generate our objective such that local loss is of the form

Supported models(or pattern) structures
Users are able to specify: dimensionality , strong convexity parameter and smoothness constant for (3) and the number of samples per each client. In addition, we consider two settings wherein the first case all clients optimizes the same objective, i.e., iid setting. In the second scenario, each client has a different objective to optimise, i.e., non-iid setting.

Tracking metrics and experiments analysis
An important feature to evaluate experiments is the ability to track all important quantities during the run of the algorithms. FL_PyTorch allows users to track the number of communication rounds, loss, accuracy, the norm of computed gradient, the norm of the full objective gradient, function values, number of gradient oracle calls, used GPU memory, number of bits which would be sent from workers to master in an actual real-world system, number of clients per round, wall-clock execution time. Furthermore, we provide a dedicated visualization tab in our GUI tool which allows user to load and visualize their experiments in an interactive fashion in several plots. You can observe such quantities during experiment execution via two-dimensional plots. Each experiment has a lot of associated information. We provide means to configure the view of experimental results via selecting needed attributes, sorting experiments by some attribute in the analysis tab, and reordering experiments manually.

The internals of the system
To reach one of our goals of making FL_PyTorch computationally effective, we analyzed various aspects of the PyTorch infrastructure, including the place of PyTorch at NVIDIA computation stack, PyTorch initialization, its forward and backward computation, and speed of typical buses used in a local node setup. This analysis led us to the system design, which asynchronously exploits available hardware.
The Algorithm 1 has several independent parts, and instead of sequential execution of these parts in a simulated environment, independent parts can be partition across independent thread-pools of workers. In our implementation, each worker within a thread pool is a separate CPU thread that lives within a Python interpreter process launched in the operation system (OS) if the system's resources allow it. Since PyTorch and therefore FL_PyTorch allow launching of the computations in a GPUs, for providing separate CPU threads ability to work independently, each thread has its separate GPU CUDA stream to submit its computation work for GPU independently. By design, each worker with the same purpose is assigned to its task-specific thread pool with a thread-safe tasks queue. There are three specific types of tasks -deferred execution, process request for finish execution, and wait on the completion of all submitted work.
Finally, we stress that all this is happening "behind the curtain, " and user interaction, including implementation of the new methods, is agnostic to the aforementioned implementation details.

Bringing custom Algorithm/ Model/ Data
As described in Algorithm 1, each algorithm supported by FL_PyTorch requires implementation of InitializeServerState, ClientState, LocalGradient, ClientOpt, LocalState, ServerOpt, ServerGradient, ServerGlobalState. These are standard centralizedlike PyTorch functions that need to be provided to FL_PyTorch in order to run simulations.
One example of adding new algorithm might be SCAFFOLD assuming that we are given implementation of FedAvg in the required Algorithm 1 format. As described in Section 3.2, one needs to set ClientState to return global shift as proposed in [22] and update LocalGradient to account for both global and local shifts. One further needs to define LocalState to return local shift update, see [22] for the detailed formula. The last modification is in ServerGlobalState, which updates global shift using the local shift updates.
For the datasets, we require users to provide it in a form that is loadable via PyTorch's DataLoaders and for the models, the required format is PyTorch's nn.Module.

EXPERIMENTS
In this section, we provide several experiments to showcase potential use cases of FL_PyTorch . In general, FL_PyTorch can serve as a simulator for researchers to test and compare different algorithms, extend them and analyze each algorithmic component in a plug-and-play fashion as described in Section 3 and Algorithm 1.
For instance, one might be interested in how to reduce the communication burden of a given algorithm. In the literature, three approaches have been proposed. For the first approach, one performs local updates to reduce the frequency of the global model updates, thus also reducing communication frequency. By this strategy, communication cost per client model update can be effectively reduced by a number of local steps. However, while often in practice this leads to improvement, it is not always guaranteed, as it was theoretically shown that there is a tradeoff between communication reduction and convergence [8]. Another possibility is to employ compression techniques that reduce the bits transmitted between the clients and the server. This compression operator should be selected carefully as it comes with an extra variance and it might

SCAFFOLD with Compression
In the first example, we used our framework to analyze whether SCAFFOLD can work with compression, i.e., whether both (Δ ( ) , ( ) ) can be compressed, without a significant decrease in performance. This experiment was carried on a synthetically generated quadratic minimization problem 3. We set the dimensionality of the problem to be 20. Both features 's and responses are generated using uniform distribution on the interval [0, 1]. After this step, we update the data matrix such that the objective is L-smooth with = 2 and strongly convex with = 1 using singular value decomposition (SVD). We generate 10 clients and consider a full participation scenario. The number of communications rounds is chosen to be 100. We set the global learning rate to 0.5 and the local learning rate to 1.0. For the compressor, we choose RAND-, with 3 values of : 100%(no compression), 40% and 20%. We provide results of this experiment in Figure 2. One can observe that dropping 60% coordinates at random has minimal effect on the convergence with respect to iterates. Dropping 80% brings a visible slow-down in per iterate convergence, but it is comparable with respect to the number of communicated bits. This experiment demonstrates that it may be worthwhile to consider an extension of SCAFFOLD by adding compression.

Benefits of Local Updates
For the second experiment, we consider a similar setup as for the previous experiment. We set smoothness constant to 5 and strong convexity parameter = 1. As a compressor we select Lazy/Bernoulli compressor described below4 with = 0.8.
We run MARINA and DIANA algorithms with full gradient estimation for 100 rounds with 10 clients and full participation. These methods were not analyzed in the setting with several local iterations that motivates us to consider these methods for this experiment. Looking into Figure 3, we can observe that both MARINA and DIANA do not benefit from extra local updates. We also experimented with different problem condition numbers ( / ), but it seems that local iterations do not speed up the convergence in the strongly convex regime. This might suggest that a naive combination of these algorithms with local updates will not lead to theoretical benefits, too.

Stochastic Updates and Compression
For the third example, we run MARINA and choose the model to be ResNet18 and for the dataset, federated version of CIFAR10 split uniformly at random (u.a.t.) among 100 clients. We run 200 communications rounds and in each round, we randomly select 25 clients. The local learning rate is 1.0 and the global learning rate is 0.001. We perform one local iteration, i.e., = 1.
We investigate the effect of the combined effect of compression and stochastic gradient estimation since both of these approaches introduce extra variance that might negatively impact the quality of the obtained solution.
For gradient estimation, we employ SGD-NICE sampling strategy that estimates gradient via selecting uniformly at random a subset of samples of fixed size.
in the experiments reflects the relative subset size with respect to the total size of the local dataset. The compressor operator is Bernoulli compression. Looking into Figure 4, one can see that having Lazy compression with = 0.1 combined with SGD-NICE with = 10% hurts the convergence. On the other hand, having a less aggressive compressor or using a bigger sample size for SGD-NICE does not deteriorate the performance.

CONCLUSIONS
In this work, we have introduced FL_PyTorch , an efficient FL simulator based on PyTorch that enables FL researchers to experiment with optimization algorithms to advance current state-ofthe-art (SOTA). FL_PyTorch is an easy to use tool that supports SOTA FL algorithms and the most used image classification FL datasets. In addition, FL_PyTorch uses several levels of parallelism for efficient execution while being simple to extend using only standard PyTorch abstractions. In its basic version, it does not require any programming and all its pre-implemented components are accessible through its GUI which allows to set up, run, monitor and evaluate different FL methods.
For the future, we plan to provide users with more pre-implemented algorithms, datasets, and models. We also aim to extend our visualization tool. After open-sourcing FL_PyTorch , we expect a significant increase in the number of supported algorithms, datasets, and models via inputs from the FL research community.  Step-1: Master -Initialize Model -Prepare shared Datasets for each thread in a training Thread Pool -Send request for local optimization handled with a local thread pool Step-3: Master Does not perform waiting of workers completion Step-4: Master -Compute weight for aggregation based on data statistics -Get the next ready model for aggregation -Aggregate Step-4: Master Step-5: Completion