作者
metaapp- 举荐广告研发部：臧若舟，朱越，司灵通
1 背景
举荐场景大模型在国内的应用很早，早在 10 年前甚至更早，百度曾经用上了自研的大规模分布式的 parameter server 零碎联合上游自研的 worker 来实现 TB 级别的万亿参数的稠密模型。起初，各家平台也陆续基于这种计划，开发了本人的分布式训练零碎，广泛特点是大量应用 id embedding，因而参数量微小，模型大小也十分夸大。当然，随着开源训练工具 TensorFlow/Pytorch 的风行，应用 TensorFlow/Pytorch 作为 worker，联合自研 ps 的计划也非常风行。究其原因，以 TensorFlow 为例，尽管内置了分布式训练零碎，然而对于大规模 id embedding 的反对却十分蹩脚，无奈作为残缺的平台应用。而应用 TensorFlow+ 自研 ps 的计划也存在不少问题，比方自研 ps 个别对于特色输出都有特定的要求、二次开发老本比拟低等。
一个典型的分布式 worker-ps 架构
2 业务介绍
metaapp- 举荐广告研发部，次要负责 metaapp 拳头产品 233 乐园的首页信息流的举荐和广告零碎，是比拟传统的推广搜组。咱们在 2020 年之前也是采纳了 TensorFlow+ 自研分布式 ps 的计划，模型大小在靠近 TB 级别（业务体量较小），整个计划的迭代和保护老本都比拟高。
在这种背景下，通过多方考量，阿里云机器学习平台 PAI 开源的 DeepRec（脱胎于 PAI-TF），作为反对了淘宝搜寻、猜你喜爱、定向、直通车等外围业务的训练平台，间接基于 TensorFlow 做二次开发，针对稠密模型在分布式、图优化、算子、Runtime 等方面进行了深度的性能优化，并且齐全开源。
而因为咱们公司自身跟阿里云有着深度的单干，阿里云也被动介绍了过后还是外部我的项目的 DeepRec 给咱们尝试。在近 2 年的工作后，DeepRec 曾经全量用于咱们的模型训练和线上 inference，并且获得了显著的性能晋升和老本降落。
3 稠密模型训练
3.1 EmbeddingVariable 多级存储
因为模型参数量大，一些特色的 embedding 大小达到了靠近 TB 级别，齐全基于内存存储对于老本的要求过高，因而自然而然就会想到多级存储：将最热的 embedding 放在显存或者内存里，其余的能够分级放在 PMEM、SSD 等老本较低的存储介质中。而 DeepRec 中 提供了基于 EmbeddingVariable 的 Embedding 多级存储性能。DeepRec 目前对于 embedding 寄存在各种存储介质的反对曾经相当欠缺。
上面介绍下咱们团队降级 DeepRec 在存储这一块的过程和教训：
3.1.1 compaction 的性能问题
咱们本来基于自研的分布式 parameter server，而过后 PMEM 类的存储介质还不遍及，因而咱们抉择了比拟高性能的 SSD 作为多级存储介质。于是咱们自然而然采纳了类 leveldb（rocksdb）的计划作为 SSD 存储计划。但这种计划在模型训练时，因为参数一直减少和更新，后盾会进行频繁的 compaction，此时会有重大的写放大问题导致 ps 的读取工夫大大缩短，从而导致模型训练的瓶颈简直都在 ps 侧。ps：据说 rocksdb 在 2022 年底的 7.5.3 版本大幅改良了 compaction 的性能，在后盾 compaction 时简直不会影响读取的性能。
3.1.2 DeepRec 的计划
而在晚期咱们试用 DeepRec 时，DeepRec 的 EmbeddingVariable 对于 SSD 存储的计划同样是基于 leveldb，因而同样遇到了跟咱们自研的计划相似的问题。后续咱们将此问题的测试后果反馈给了 DeepRec 相干的同学，他们基于此后续推出了基于 SSDHASH 的存储计划，大大晋升了 compaction 时的读取性能，因而模型训练基于不再受困于 ps 的读取性能问题。后续又进一步了基于 SSDHASH 的同步和异步两种 compaction 的形式。应用同步 compaction 时，向 SSD 写入数据和 compaction 将会应用同一个线程，异步时则各应用一个线程。这里也举荐大家应用这种计划。
3.1.3 压缩模型大小
进一步的，如果能把模型大小管制在数十 GB，那 ps 的性能就能够进一步晋升了。因为采纳 DeepRec，自定义各种压缩形式的算子变得十分轻松。咱们调研并实现了了多篇 embedding 压缩方向的 paper，最初采纳了 binary code 的形式实现了 embedding 的 multihash 计划，能够自在管制 embedding 的大小。咱们尝试在最大的特色 uid embedding 上应用了 multihash，把模型大小从 800GB 升高到 40GB 以下，auc 的损失仅在千分之三左右，线上点击率降落了 1.5%；进一步的，咱们通过优化序列举荐模型，更好的通过序列特色建模了用户的个性化，能够发现在序列模型的根底上把 uid embedding 换成 multihash 的计划，对于线上点击率的影响仅有 0.3% 左右，因而能够释怀全量 multihash 计划。咱们也把基于 multihash 的 embedding variable 算子以 pr 的模式提交给了 DeepRec。
3.2 基于 GPU 的分布式训练
在解决了 ps 的性能瓶颈后，模型训练的速度就和模型 Tensor 计算的算力近似线性相关了。而近几年随着序列模型的倒退，搜广推的矩阵计算复杂度也在显著晋升。此时应用 gpu+ 大 batch size 来代替 cpu 作为 worker 的计划，无论在性能还是老本管制上都有微小的劣势。而阿里云机器学习平台 PAI 开源的 HybridBackend 平台就反对了基于 GPU 的分布式训练计划，并且深度反对了 DeepRec。
能够看到应用 hb 的计划在训练速度上比照 TF-PS 原生计划的劣势。
3.2.1 模型参数齐全放在显存里
想要充沛开释 gpu 的算力，缩小因为数据拷贝带来的性能损耗，最好的计划天然是把所有参数都放在 gpu 显存里。下面 2.1.3 提到的压缩模型大小，为这种计划提供了可能性。调大 batch size 则能够进一步提高显卡的利用率。通过测试，在这种状况下，单张 V100 显卡的算力能够超过 20 台 40core worker 节点的算力。
3.2.2 解决了多卡训练失落数据的问题
在单机多卡训练时，咱们发现和单卡训练相比有近 1/3 的数据被抛弃，这是因为 hybridbackend 默认应用所有 worker 依照 row group 平分数据的策略，以进步读取效率。当 group 数目不够均分时，多余的数据会被抛弃，当 parquet 文件较多且比拟小时，该问题尤为重大。咱们通过应用每个 worker 加载所有的 group，再依照 batch 平分数据的策略，极大地缓解了数据失落的状况，读取压力也在可接管范畴内，后续可思考将两策略联合升高 worker 的读取压力。
4 模型 inference
4.1 痛点
在咱们的理论场景里，线上 inference 的痛点大部分来自于保护老本。因为举荐广告业务场景，须要大量尝试各种模型在线上调配流量做 AB test，因而线上存在的模型量级大略是 10 倍的基线模型量级。而每次上线一个模型，都须要给对应的模型调配相应的资源，并且这个资源跟 AB test 的流量正相干；而每次调整 AB test 流量（比方模型成果不错，放大流量察看）的时候，又须要调整该模型调配的资源。这个过程比拟难实现自动化，往往须要算法工程师手动扩缩容。
4.2 基于 Processer 库的 inference 计划解决痛点
下面这个图是咱们线上理论的 inference 计划。
4.2.1 单机器运行所有模型
基于下面的痛点，咱们给出的计划是应用大规格机器（比方 128C，512G 内存）来做线上 inference，而后每台机器都会有线上所有的模型实例。每台机器运行一个 serving-proxy 会主动的治理所有的模型过程，包含模型高低线、模型更新等。这种计划的益处是整个保护老本根本没有了，所有事件根本都自动化实现了。因为线上整体的流量绝对稳固（比方扩充 AB test 模型的流量，天然基线模型流量就缩小了，整体是稳固的），所以各个模型之间资源竞争也不须要重新分配资源。
4.2.2 基于 DeepRec 提供的 Processer 库
DeepRec Serving Processor 是用于线上高性能服务的 Library，能够参考文档。因为自身是一个独立的 so 包，咱们能够很不便的对接到本人的 Serving RPC 框架中。咱们采纳 golang 语言来实现了咱们本人的 serving rpc 我的项目，长处天然是开发成本低并且性能不错。
4.2.3 应用 DeepRec 的 Session Group
间接应用 TensorFlow 提供的 C++ 接口调用 Session::Run，无奈实现多 Session 并发解决 Request，导致单 Session 无奈实现 CPU 的无效利用。如果通过多 Instance 形式（多过程），无奈共享底层的 Variable，导致大量应用内存，并且每个 Instance 各自加载一遍模型，重大影响资源的使用率和模型加载效率。
DeepRec 中 SessionGroup 可配置一组 Session，并且通过 Round Robin (反对用户自定义策略) 形式将用户申请散发到某一个 Session。SessionGroup 对不同 Session 之间的资源进行隔离，每个 Session 领有公有的线程池，并且反对每个线程池绑定底层的 CPU Core(numa-aware)，能够最大水平地防止共享资源导致的锁抵触开销。SessionGroup 中惟一共享的资源是 Variable，所有 Session 共享底层的 Variable，并且模型加载只须要加载一次。
咱们应用 session group 后，实测调整到适合的 group 数量，能够进步 50% 的 inference 性能。
4.2.4 基于 oneDNN 的优化
DeepRec 集成了英特尔开源的跨平台深度学习性能减速库 oneDNN（oneAPI Deep Neural Network Library），并且批改 oneDNN 原有的线程池，对立成 DeepRec 的 Eigen 线程池，缩小了线程池切换开销，防止了不同线程池之间竞争而导致的性能降落问题。oneDNN 曾经针对大量支流算子实现了性能优化，包含 MatMul、BiasAdd、LeakyReLU 等在业务场景中应用到的常见算子，为业务模型提供了强有力的性能反对。更值得一提的是，oneDNN 的算子反对 BF16 数据类型，与搭载 AMX(Advanced Matrix Extensions) 指令集的第四代英特尔® 至强® 可扩大处理器同时应用，可显著晋升模型训练和推理性能。在 DeepRec Serving Processor 编译选项中，只需退出“–config=mkl_threadpool”，便可轻松开启 oneDNN 优化。
4.2.5 子图优化
子图交融是推理性能优化的罕用办法。然而对于本模型中左图所示的子图构造含有 Reshape 算子，原生 tensorflow 并没有对应构造的图优化器以及算子实现，咱们通过手动交融来实现，交融前后的子图形成如下图所示。这样缩小了多余算子的运行开销，缩小了内存拜访，晋升了计算效率。再联合 oneDNN 减速交融算子，最终业务端到端减速了 10%，CPU 利用率降落 10%。
4.2.6 cost model 的设计
因为大机器的 cpu core 数较多，而咱们是一台机器有所有模型的过程，那么所有模型都共享所有 cpu core 显然会造成不必要的资源竞争等。因而给不同模型设计正当的 cost model 就很有必要。咱们目前采纳比较简单的形式，因为基线模型和须要做 AB test 的模型资源差异较大（流量差距大），咱们会给每个基线模型调配对应的 core，而后让所有非基线模型共享一组 core（总体 AB test 的流量有下限）。尽管这个计划很简略，然而获得了十分好的成果，大略有 30% 的性能晋升。
5 后续布局
1、cost model 的优化，显然有更好的计划来动静的调整每个模型须要的 core。咱们打算开发更好的 cost model 并提供给 DeepRec。
2、开源咱们的 inference 架构计划，因为在咱们的业务里，基于 DeepRec processor 设计的 inference 架构带来了微小的便当，并且性能很好，咱们预计在上半年会开源咱们的 inference 架构计划，欢送大家到时关注。