导读

本次分享的主题是腾讯 PCG 搜广推机器学习框架 GPU 性能优化实践。主要围绕三个问题展开：为什么 GPU 推荐模型训练框架会成为刚需？怎样的设计才能让 GPU 推荐模型训练框架最高效？以及未来还有哪些值得期待的方向？

分享嘉宾｜骆兆楷 腾讯 机器学习框架高级专家

编辑整理｜胡俊琪

内容校对｜李瑶

出品社区｜DataFun

为什么 GPU 推荐模型训练框架是刚需

PCG 算力集群的瓶颈

腾讯 PCG 的推荐模型训练最初全部跑在 CPU 上。但随着业务的深入和深度学习模型的发展，这套方案开始捉襟见肘。简单列几个痛点：

• 系统网络带宽小，而且不稳定。
• 模型规模越来越大，是选择多机多卡还是单机多卡？这直接决定了硬件选型的方向。
• 云上分配的 CPU 型号无法保证——有时是 AMD，有时是英特尔。这种不确定性对参数服务器架构非常不友好，老旧的 CPU 型号会成为性能瓶颈，拖慢整个训练框架。
• 云容器不是独占的，整个机器的 IO 和网络都是共享的。这导致训练框架的稳定性大打折扣。

种种迹象表明，传统的参数服务器已经很难支撑大参数量的推荐模型了。转向 GPU，并且用单机多卡的方式去训练超大模型，成了唯一的选择。

这里还有一层现实考量：如果采用多机多卡，就绕不开 InfiniBand。但 InfiniBand 价格昂贵，而且需要对机房进行改造——换机器、加网卡、改交换机……成本高得吓人。所以最终的技术架构只能锁定在单机多卡。

还有一个隐藏问题：CPU 型号老旧。多机多卡会让每个节点的性能参差不齐。GPU 在容器中是共享的，不能保证每个 GPU 都能绑定到特定的 CPU。一旦涉及多机多卡，K8s 调度集群的压力会剧增，还会产生大量 GPU 碎片。

业务需求与上游生态

结合硬件现状、业务场景和上游生态来看，我们需要明确几个技术目标。首要目标是：用单机多卡来训练一个超大模型，模型规模在几 TB 到 10 TB 之间。

第二个目标：尽量压榨硬件的每一分潜力——不仅要用 SSD、主机内存和 GPU，还要把网卡也用上，以此提高模型训练的上限。

整个架构必须同时支持离线训练和在线训练。离线训练用来追模型，在线训练负责实时更新。所以架构设计还得考虑如何快速上线、快速部署。

GPU 翻跟斗的选型是个麻烦事——供货紧张。架构设计不能只绑定某一家 GPU，还要兼容 XPU。

深度学习框架上，我们选择了 TensorFlow，但框架得能快速切换到 PyTorch，目前也正在做这件事。

另外，必须兼容老参数服务器的技术方案，让 CPU 训练和推理能平滑迁移到 GPU，降低迁移成本。出于这个考虑，架构上要兼容参数服务器。

以上这些，就是我们根据目标进行的技术选型。

GPU 推荐模型训练框架怎么做才最高效

GPU 训练的数据结构

接下来聊聊 GPU 推荐模型训练框架的具体设计。先看数据结构。大目标是单机多卡，所以存储层级必然包含 SSD、主机内存和 GPU。这个三级缓存的思路和传统 GPU 训练框架很像，但我们的设计里有自己的考量——最下层是对象存储，用来存模型和样本输入，空间非常大（几 TB 甚至上百 TB）。因为存储在远端，读写速度只有几百 MB 到几 GB 的量级。所以必须把读写过程和训练过程分离开。

要尽可能实现读写分离，于是需要把样本和模型先落盘到 SSD。SSD 的读写性能大约在几 GB 到十几 GB 每秒，同样要把 SSD 的读写和训练过程分开。

再往上，主机内存到 GPU 的 HBM 带宽已经达到几十 GB 每秒，所以在这一层可以设计快速的数据交换。

GPU 的 HBM 容量在几十 GB 左右，数据可以以整块的方式刷进 HBM，然后 HBM 再以 batch 的形式进行训练。

整个数据结构分成四级，上图已经画得很清楚了。对应不同数据集，数据会被分成多个 group，每个 group 再细分到不同 PASS，每个 PASS 又包含多个 batch。这样的分块设计刚好匹配了模型的读写需求和硬件的差异化特性。

GPU 训练整体架构

训练流程如上图所示。由于硬件的特性，我们需要把训练过程和数据读写过程分离。训练开始时，先预先下载大量数据，把一个 group 的数据拉下来，然后做预处理。处理完一个 group 后，把它放到主机内存里，再以整块的方式 flash 到 GPU 的 HBM。因为 HBM 里已经存了多个 batch，GPU 的流式处理器访问全局内存时速度会很快。所以当 PASS 全部放入 HBM 后，就可以按 batch 进行一轮又一轮的同步训练了。图中的 compute 就是 forward-backward 同步训练。处理完所有 group 后，进行模型导出。

数据下载和预处理等每个阶段都可以并发执行，这样就能实现多级流水线，把硬件资源的带宽打满。每个 compute 过程（forward 和 backward）都会用到 GPU。推荐模型训练的计算任务通常是 memory bound 的，所以每个访存环节都可能成为瓶颈。一次 forward-backward 对应一个 batch，global memory 的访存是瓶颈；预处理阶段涉及 CPU 和 CPU IO 读取；业务上线还有显存访存的限制……这些限制在很大程度上决定了训练框架的性能。后面会详细介绍对应的优化。

数据下载过程中的优化手段比较零碎，但很实用。首先下载涉及网络，所以用了 DPDK 这样的网络优化库。网络访问可能需要通过 TCP 访问远端数据，因此用到了 Google 的 BBR 拥塞控制算法来优化网络窗口。第三个是 DMA——网卡每下载一批数据，就能用更少的 CPU 完成操作。第四是落盘优化：一台机器可以配多个 SSD 盘，通过 LVM 虚拟化，让多个盘并行访问，加快数据下载速度。第五是 direct IO，因为不需要过多的 cache IO，直接用 direct IO 完成 group 落盘。最后是数据结构优化，比如采用 Parquet 列存，既能减少 Linux 内核的中断操作，也能把数据变成平整的大块，更利于并发处理。

各阶段优化详细介绍

先看预处理阶段的优化。我们会把每个 group 中用到的去重后的 key 先捞出来，将这些 key 里的特征转换成 CSR 格式。这两点非常关键——CSR 格式能减少访存的 cache miss rate；去重后的 keys 可以对所需数据做预处理，从而实现整块 flash 到 GPU memory 的高效操作。

前面提到数据结构整体是三层结构，大小从下往上递减。内部数据结构是这样的：一个 group 先预处理，然后落盘，与训练过程分离。它会把每个 key 需要用到的 embedding vector 和 optimizer state 提取出来落盘，然后变成一个 sample，再变成 CSR 数据，多个 CSR sample 组成一个 CSR batch，放到 PASS 里。每个 CSR 存在 PASS 里，多个 PASS 组成一个 group。为了节省主机内存，会先把数据落盘到 SSD。

开始训练 PASS 之前，需要把该 PASS 需要的所有 embedding vector 整块 flash 到 GPU 显存上的 embedding 表里。上图中红框部分就是准备好数据的状态，绿色的 PASS 对应三级缓存中绿色那一块。整块 flash 到 GPU 的 global memory 后，global memory 存的是 CSR batch 和该 batch 用到的 embedding vector。这里的 embedding 其实是一个 GPU hashtable。我们做了一点特殊处理：传统 GPU hashtable 支持增删查改，而这里我们直接把插入功能去掉了——因为数据已经整块处理好，不需要 hashtable 的 insert 过程。去掉 insert 后，把整块 hashtable 的数据结构准备好，直接 flash 到 global memory，性能非常好。而且每次 embedding 只存一个 group 所需的训练数据，GPU 显存只需关注这一个 group 的 embedding vector 就够了。

只要 SSD 能存得下，单机多卡就能训练多大的模型。举个例子，虽然单机多卡总显存只有 320 GB，但可以训练十几 TB 的模型——这也是我们框架的优势之一。每个计算过程里，我们会把 PASS flash 到 GPU memory，每次 forward-backward 时捞一个 batch 进行 forward，进入 TensorFlow 的计算图，然后 backward 回来。图中黄色箭头表示 TensorFlow 在 forward-backward 时可以用 tape。每个 TensorFlow 计算都是 batch 数据并行的，所以每个计算图会虚拟一个小的块，forward-backward 后返回梯度，梯度在多卡 buffer 里做 reduce 操作，然后 scatter 到不同卡上。因为 key 已经对每张卡做了哈希，所以每个 batch 里哪个 key 在哪张卡上是固定的。

接下来是 Compute 过程中的优化。空间优化方面，embedding 部分采用了 Dynamic embedding——key value 指向一个 embedding vector，这个 vector 不是定长的，长度可以从 0 到最大值，这样很容易实现动态 embedding。第二个优化是 Multiple Hash，一个 hashtable 可以分解成两个或多个小 hashtable，空间也能得到优化。第三是混合精度训练，虽然已经是常规操作，但确实是有效的空间优化手段。

计算优化方面，首先是 Unified Hash Table。不同特征放在不同 embedding table 里，我们通过对 key 做 rehash——key 是 int64，用低 48 位存真实 key，高 16 位存 hashtable 的 ID——把所有特征打散到一个超大 hashtable。查找只需要一个 kernel，backward 也只要一个 kernel 就能更新，节省了 kernel launch 的开销。第二是 kernel merge，这是最近做的新功能。第三是混合精度，不仅能优化空间，在 GPU 上 FP16 可以用向量化访存，还能用 half2 做多数据指令融合。最后是非对称结构的 hash table。

上图是对 unified hash table 的详细介绍。再详细说一下 kernel merge：把相同参数合到一起，这样 pooling 的 kernel 开销就能降低。这些涉及 GPU 优化的细节就不展开了。

业务上需要增量模型和全量模型的分块操作——每天更新一个全量模型，每小时更新一个增量模型，这样才能快速上线。模型更新和上线的优化手段主要包括 INT 量化、异步 IO 以及腾讯云 COS 的协程队列等。

未来展望

最后聊聊未来的工作方向。首先，GPU 卡供应紧张，将来可能没法用 A100 训练了，所以得考虑如何在 A10、T4、P4 甚至非英伟达 GPU 上训练。第二是推荐大模型与 GPT 的结合。第三是更灵活的架构，比如参数服务器和 embedding 是不是一定要在 GPU 上？GPU 训练能不能结合 PS 架构去支撑更大的模型？第四是更大规模的训练——目前只支持单机多卡，模型上限是 TB 级，之后要考虑能否支持 PB 级。最后，希望用更低的硬件配置去训练更大的模型。

以上就是本次分享的内容，谢谢大家。

分享嘉宾

INTRODUCTION

骆兆楷

腾讯 机器学习框架高级专家

2016 年加入 NVIDIA，参与 TensorRT 开发。2018 年加入阿里巴巴达摩院参与自动驾驶计算中台开发。2020 年加入腾讯 PCG，从0到1参与打造 GPU 推荐领域深度学习计算框架（无量）的开发。