如何设计并实现一个异构参数服务器

1、背景

• 什么是参数服务器？

对于 AI 大模型来说，如果参数量太大（百亿、千亿级别），单机存储不下时，就需要引入多台参数服务器（PS，Parameter Server），用于分布式存储这些参数。当 Trainer 机器训练时，就可以实时从 PS 拉取（Pull）参数，本地反向计算得到的梯度可以推送（Push）到 PS 端，在 PS 端完成参数更新（Sgd、Adam 等优化算法）。

在 CTR（搜、广、推）场景下，稀疏参数量（Sparse）巨大，并行方式为数据并行，参数服务器训练模式尤其适用。

• 为什么需要异构参数服务器？

如果一个集群中都是同构的机器，比如都是 CPU 或者都是 GPU，而且采用纯粹的数据并行方式（也就是每台训练机器上的模型是一模一样的），那么使用传统的参数服务器就足够了。如果是跨集群训练，而且希望在 A 集群上训练模型的前半部分，在 B 集群上训练模型的后半部分，那么就需要引入异构参数服务器，该模式能更灵活、高效地利用硬件资源