Suggested design for multiprocess federated learning

RickZack: <p>Good afternoon,<br>
I have a federated learning codebase that runs the simulation sequentially. I would like to use pytorch multiprocessing or distribute packages to be able to leverage multiple GPUs and (optionally) multiple nodes.</p>
<p><strong>Brief description of my use case (FL)</strong><br>
A typical FL simulation consist of several <em>rounds</em> of training. At each round, a subset of compute entities (<em>clients</em>) run several optimization steps on a private local model and private data. At the end of each round, the gradients of each client are merged (e.g. weighted average) and used to update a global model, which is then used as initialization model for the clients selected at the next round.</p>
<p><strong>The code I start from</strong><br>
The code I use is built around three classes:</p>
<ul>
<li><code>CenterServer</code>: contains the global model, the aggregation logic of local models and testing code</li>
<li><code>Algo</code>: contanins the main loop, in which the global model is sent to client, they perform local training and then the trained are sent to server for aggregation</li>
<li><code>Client</code>: contains a dataloader and a copy of the model.</li>
</ul>
<p>The <code>CenterServer</code> object and all the <code>Client</code> objects are all created inside <code>Algo</code> (e.g. all the datasets and dataloaders).</p>
<p><strong>What I want to do</strong><br>
At the moment the train loop sequentially trains all the clients, I would like to leverage multiple GPUs or, if possible, also multiple training nodes. I am asking for some advice about what is the best way to port my codebase so that I can achieve that.</p>
<ul>
<li>
<p><strong>Solution 1: use processes</strong><br>
I could use processes to make each client perform its local optimization in parallel. However I have some doubts about how the memory is used. Le it be <code>f</code> the function passed to <code>Process</code>:</p>
<ul>
<li>Do the args passed to <code>f</code> get copied? For example, the client object contains a dataloader, and most of my datasets are preloaded on main memory (as they are small, like CIFAR or MNIST).</li>
<li>If in <code>f</code> I set the client’s model as a copy of the server model and then train it, when the functions will terminate will I still have these models as attribute of clients? Do I need to use <code>share_memory()</code> on them?</li>
<li>For this use case I could make use of <code>mp.Pool</code> to avoid creating processes each time, but on the docs I read that the processes created cannot make other process. This is not handy when using datasets that need to load data from disk, because i could not set <code>num_workers</code> &gt;1 in the dataloaders. Any workaround?</li>
</ul>
</li>
<li>
<p><strong>Solution 2: use distributed</strong><br>
The use of distributed package seems most general and would allow to use also more than one node. However it is even less clear how should I change my current code. Also, in the docs I read I cannot run two processes on the same GPU, while it would beneficial to me since the models I use are relatively small and I already experienced speedup in training multiple clients on the same GPU. Is there a way I can do it with distributed?</p>
</li>
</ul>
<p>I hope I made clear my doubts, and that someone will help figuring out what is the best solution for my case.</p>

Gutomeur: <p>Good evening, senior brother.</p>
<p>Reviving this old thread as I’m facing the exact same challenge you described.</p>
<p>I’m particularly struggling with how to efficiently run multiple small client trainings on multiple GPUs without massive process overhead or resorting to sequential execution.</p>
<p>Have you already known the answer to the question you asked? Any insights would be greatly appreciated.</p>
<p>Thanks!</p>