联邦学习系列
引言
本节介绍从理论上推理联邦系统的带宽使用情况，以及如何在实践中使用 Flower 来测量带宽消耗。
理论分析
需要考虑到：发送个客户端的模型或者客户端发送来的更新模型可以压缩；每次聚合不对所有的客户端进行聚合而是采用抽样的方式；可以有以下计算公式。
带宽消耗 = （发送给客户端的模型大小 + 接收客户端发送来的更新模型大小）× 客户端数量 × 每个轮次中选择客户端的比例 × 训练轮次。
本课程使用的语言模型是 EleutherAI 的 pythia-14m，它拥有 1400 万个参数，这个模型的大小是 53 MB。
本小节所用到的资源都放在这里
实验验证
模型准备
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from flwr.client.mod import parameters_size_mod
from utils5 import *
# 初始化模型
model = AutoModelForCausalLM.from_pretrained(
"EleutherAI/pythia-14m",
cache_dir="./pythia-14m/cache",
)
# 打印模型信息
vals = model.state_dict().values()
total_size_bytes = sum(p.element_size() * p.numel() for p in vals)
total_size_mb = int(total_size_bytes / (1024**2))
log(INFO, "Model size is: {} MB".format(total_size_mb))
# 定义客户端
class FlowerClient(NumPyClient):
def __init__(self, net):
self.net = net
def fit(self, parameters, config):
set_weights(self.net, parameters)
# No actual training here
return get_weights(self.net), 1, {}
def evaluate(self, parameters, config):
set_weights(self.net, parameters)
# No actual evaluation here
return float(0), int(1), {"accuracy": 0}
def client_fn(context: Context) -> FlowerClient:
return FlowerClient(model).to_client()
client = ClientApp(
client_fn,
mods=[parameters_size_mod],
)
# 自定义策略跟踪发送的模型大小
bandwidth_sizes = []
class BandwidthTrackingFedAvg(FedAvg):
def aggregate_fit(self, server_round, results, failures):
if not results:
return None, {}
# Track sizes of models received
for _, res in results:
ndas = parameters_to_ndarrays(res.parameters)
size = int(sum(n.nbytes for n in ndas) / (1024**2))
log(INFO, f"Server receiving model size: {size} MB")
bandwidth_sizes.append(size)
# Call FedAvg for actual aggregation
return super().aggregate_fit(server_round, results, failures)
def configure_fit(self, server_round, parameters, client_manager):
# Call FedAvg for actual configuration
instructions = super().configure_fit(
server_round, parameters, client_manager
)
# Track sizes of models to be sent
for _, ins in instructions:
ndas = parameters_to_ndarrays(ins.parameters)
size = int(sum(n.nbytes for n in ndas) / (1024**2))
log(INFO, f"Server sending model size: {size} MB")
bandwidth_sizes.append(size)
return instructions
# 定义服务端
params = ndarrays_to_parameters(get_weights(model))
def server_fn(context: Context):
strategy = BandwidthTrackingFedAvg(
fraction_evaluate=0.0,
initial_parameters=params,
)
config = ServerConfig(num_rounds=1)
return ServerAppComponents(
strategy=strategy,
config=config,
)
server = ServerApp(server_fn=server_fn)
模型运行
1
2
3
4
5
run_simulation(server_app=server,
client_app=client,
num_supernodes=2,
backend_config=backend_setup
)
打印结果
1
log(INFO, "Total bandwidth used: {} MB".format(sum(bandwidth_sizes)))
结论
通过实验发现即使是只训练一轮所消耗的带宽也非常的大。在联邦学习中有许多减少带宽使用的方法，两类改进是减少更新大小和减少通信次数。
减少更新大小：稀疏化和量化。对于 top case（典型场景）进行稀疏化，如果要通信的梯度低于一定阈值可以将其作为零进行通信而不是实际通信它们。在训练接近尾声的时候这种情况很有可能发生，此时梯度中的更多元素的幅度会更小。量化可以减少客户端和服务端交换的更新的大小，它使用预训练模型而不传输所有层。
减少通信次数：在本地训练一定轮次之后再进行更新。但是当本地训练轮次太多会导致模型聚合效果不好，听说还有可能会过拟合。
