使用 MPI 加速联邦学习的模拟

前言

最近我正在进行实验，发现自己之前搭建的框架用起来有些不舒服，于是决定进行重构。然而新的框架还没有写好，于是我先搭建了一个基于 mpi4py 的简单联邦学习模拟器，用来跑一些简单的实验。

代码放在了这个 GitHub 仓库：https://github.com/beiyuouo/mpi-fedsim

代码结构

代码结构比较简单，主要分为五个部分：

• model.py：模型定义，包括模型的定义，以及模型的初始化
• utils.py：一些工具函数，以及数据集下载划分等
• server.py 和 client.py：服务器和客户端的定义，包括接收和发送模型，以及模型的更新
• main_sync.py：同步联邦学习的主程序
• algor：一些联邦学习的算法，包括 fedavg.py
• config：一些配置和实验参数。

日志模块使用的 loguru 和 tensorboard，模型创建和训练是 PyTorch，虽然比较的简单，但还是 “五脏俱全” 的，哈哈哈哈哈

执行逻辑

逻辑主要存在于主函数中¹，首先如果 rank=0 的话，就是服务器，否则就是训练进程。服务器会先初始化模型，然后等待每个进程上传样本个数和分配到的客户端 id，并记录。其他进程的话，会先根据预先规则分配到客户端 id，然后根据 id 加载对应的数据集，收到模型后开始训练，每个进程训练完之后，会把模型上传到服务器，服务器收到之后，会把模型更新到全局模型，然后把模型发送给所有进程，进程收到之后，更新模型，然后开始下一轮训练。

因为我本身做实验是单机多卡的环境，因此在线程分配的时候可以顺便分配一下指定的 gpu，但 CNN 毕竟很小不会爆显存，所以这块我没有写，有需要的话只需要在 main_sync.py 中修改一下 client_id 与 gpu_id 的映射关系即可。

由于每个进程包含多个 client，但接受模型参数时通讯一次，因此在数据传输上或许会快上一些。

其实，设置多少个线程都是无所谓的，线程的限制源于对显存的限制，因此只要显存够，线程数可以设置的很大，这种方式只是在有限的显存下，使用尽量多的显存来进行加速。

使用

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

# install requirements
pip install -r requirements.txt

# dataset prepare
python utils.py  # check and download the dataset you need

# run the simulation
# for linux
mpirun -np 4 python main_sync.py  # 4 is the number of processes
# for windows
mpiexec -n 4 python main_sync.py  # 4 is the number of processes

# launch the tensorboard
tensorboard --logdir=logs

Benchmark

这里就用了一个两层卷积的 CNN 来做 MNIST 分类，没跑太多。实测的话 10 个客户端，每个客户端 6000 个样本，local epoch 为 1，round 为 50，开了 11 个进程，大概 1 分钟左右就能跑完。

待补充

改进

其实代码写的还是比较粗糙，还有很多部分可以改进，首先就是接收和发送模型的部分，对于同步来说，用广播的方式应该会更好，比较能提高效率，不需要 1 对 1 的 check。