闲话 - 前
这是一篇 ATC'25 论文 的中文简要版本.
然而这篇文章在学术创新上实在乏善可陈,正好多说一点闲话了。这也是为什么没有在文章发表前就把博客写出来,不然万一被撤稿了可就搞笑了。
写这篇博客是 laekov 一直以来坚持要把自己发表的每篇文章都翻译成中文,因为曾经的 laekov 觉得凭什么论文都要用英文写呀,而且八股文读起来好麻烦呀,用中文简洁明了地把文章讲清楚不是更好吗。
然后天真的 laekov 读博五年一共也就中了两篇文章。这是第二篇,中稿时间是毕业前三个月。
HypeReca 这个名字是 Hybridly Partitioned Embeddings for Recommendation Model Training Acceleration 的各种胡乱拼凑里找了个听起来好听一点的。
背景
基于深度学习的推荐模型是现在主流的推荐算法,它分为学习元素特征的稀疏部分和学习元素间关系的稠密部分。 稀疏部分主要是非常稀疏的嵌入表 embedding tables,它们可能会有 TB 级别甚至更大,所以得要很好地分布式地存下来,并且支持高吞吐地读取和更新。 稠密部分通常由传统的 DNN 构成,属于计算密集型任务(比如 CNN / MLP),所以得用 GPU 来加速。
在做分布式训练的时候,这两部分由于特性不同所以采用了不同的并行方式:稀疏部分用模型并行(切分存储),稠密部分用数据并行(小模型,大 batch size)。 这两部分之间通过 all-to-all 集合通信来传输所需的数据和用于更新的梯度。 然后就会发现在 scale up 的时候两部分之间的通信是瓶颈,可能会占到迭代时间的 90% 以上。
解决方案
这篇文章的核心想法其实很直白:embedding data 有冷有热。 把频繁访问的热数据找出来,也用数据并行的方法来存,重复地放到所有 GPU 上。
接下来要解决两个问题:
首先,挑多少热数据比较合适? 热数据越多,可以省掉的 all-to-all 通信就越多,但同步热数据所需的 all-reduce 量就越大。 这里正好是一个在单峰函数上找最优的 trade-off。结合实际数据建模观察一下即可。
然后,对于一个 batch,怎么高效地查出每条所需的数据在哪? 于是搞了一些分布式哈希表,为了减少锁开销搞了一些分片+流水线的机制,从而使得在 CPU 上查找位置+插入数据的吞吐量足够高,可以 match 到 GPU 的速度。
具体的系统实现就是把 huge embedding table 包装成一个 kv store 的形式来用,在 kv store (或者说数据库)内部自己用二部并行策略来维护不同的数据。
实验评估
没有闭源数据集,也没有大规模 GPU 集群。
在最大的开源数据集 Criteo Terabytes 上跑,最多用了 32 块 A100 或 V100,最快比 NVIDIA 的 HugeCTR 快了几倍,因为原版的 all-to-all 着实是慢。
闲话 - 后
用中文写出来之后发现这篇文章真的是太简单太直白了。能中真的是个奇迹。
这个工作最初是2021年底被派去和某公司进行校企合作,获得了一个场景。但公司内部代码仓库很复杂,难以使用,而 laekov 的导师则建议在开源数据集和软件栈上进行科研。 于是在最初这个简单的想法上进行了反复的折腾和开发,历时一年(消失的2022年)终于写出了一个能跑且性能能打过 HugeCTR 的版本,于是在那个所有人轮流 covid 的 12 月里决定总结一下创新点开始投稿。 然后就有了在之后前前后后两年多的时间一共投了六次才中的经历。
- ATC'23 初投,因为时间紧所以写了 5 页短文,胡乱讲了讲,被 early reject。
- SC'23 第一次扩成正经长文,讲了一个 scalability 的故事,但是实际只有 32 卡,缺乏说服力,文章写得也废话太多。
- PPoPP'24 继续投 SC 版本,着重讲性能模型,没中但幸运地获得了一个 poster。
- ASPLOS'24 秋季接着投,故事改成了 sparse communication,因为和相关工作没讲清楚而被拒。考虑到 ppopp 有了一个 poster,以及后来开始做其它工作了,所以这时候几乎就弃疗了。
- ICDE‘25 在还有半年多毕业的时候想着怎么也得投出去吧,于是找了一个当时正好 DDL 的隔壁领域的会,把故事改成了 database,但因为没注意到他们的会是单盲要在 pdf 上写名字而被 desk reject。
- ATC'25 其实是给另一篇文章陪跑,顺便就把 ICDE desk rejected 的版本改了改模板交上去了。唯一修改是为了符合数据库的故事而把 lookup 提到了前面,two-part parallel 推到了后面。没想到竟然有审稿人觉得这个文章讲得又清晰又 solid,于是很幸运地中了。(又或许我仍未知道真正的中稿原因)
总结起来这篇文章的投稿周期里,每次都是对故事小修小改,实验自从2022年写完之后,就没有再进行过大的更新。 大家常用的卡已经从 V100 变成了 A100 再到 H100,期间还出现了如 SOSP 的 BagPipe 等一系列相关文章,把 laekov 这篇文章按在地上摩擦。
laekov 一度已经很绝望了,不再想碰这篇文章。 不难想象,工业界早就掌握了比这更高超的将 embedding vectors 分开并行处理的方式,并很好地解决了训练中的 communication bottleneck。 学术界如果没有和工业界合作,并不会有兴趣训练这样的纯 profit-making 的模型,更不要谈它的优化过程了。 来自学术界的审稿人只会觉得这篇文章无聊,而来自工业界的审稿人会发现这些事他们早就玩烂了,只不过没有发文章而已,因而这篇投稿实在是缺乏新意。 实在难以想象它是在何种场景下受到了赏识。
更令人唏嘘的是同场投稿的另一篇 laekov 认为不论是问题还是解决方案都更具创新性甚至引领性的文章在同一次审稿中被拒了。 事实上那篇文章的 arxiv 版本的引用量已经是这篇文章 (ppopp poster) 的十倍有余,足以见得这两者的高下了,但是为什么明明是倍受尊敬的 peer review 审稿过程却无法保序地判断两者的高下呢?
laekov 只能说自己虽然博士毕业了,但其实并没有完全学会学术罢。