小红书训推异构引擎的设计与应用
导读 今天给大家带来的分享主题是小红书训推异构引擎的设计与应用。
主要内容包括以下五大部分:1. 小红书模型工程面临的挑战
2. 异构弹性引擎的设计与实施
3. 面向未来 HPC 训练框架
4. AI 编译技术
5. 展望未来
分享嘉宾|曾鸣堃 小红书 训推计算引擎负责人
编辑整理|娄政宇
内容校对|李瑶
出品社区|DataFun
小红书模型工程面临的挑战
首先,随着模型的复杂度增加,其处理的数据量也越来越大;
其次,计算流程的需求也会随之增长;
此外,模型的应用场景也越来越广泛,不仅限于传统的搜索广告和推荐,还包括电商和直播等新业务。
02
异构弹性引擎的设计与实践
1. 第一代训练框架
第一个是它能够支持超大规模的稀疏特征。
第二个是区别于传统 TensorFlow 的 Feature Column 的 Embedding Table 实现,我们使用了基于无冲突哈希表的实现,可以确保每个 ID 都能得到充分的学习。
第三是有一个高性能的 Lookup Table 算子实现。从上图中可以看出,我们将整个稀疏 Embedding 操作抽象成了几个算子。与原始的 TensorFlow 相比,我们进行了许多 OP 的融合,从而获得了更高效的性能收益。
第四个是支持了多种的参数的优化器。
2. 异构 GPU 训练框架
3. 第一代 GPU 推理架构
4. 异构 GPU 推理架构
5. 总结
03
面向未来的 HPC 训练框架
首先,无论是异构的训练引擎还是之前的 Larc,都是基于 PS worker 架构。然而,随着节点数量的增加,训练加速比会有显著的衰减。
其次,为了提高训练吞吐量,我们会增加更多训练节点。但由于 PS worker 采用异步训练模式,增加更多节点会导致更大的异步性,从而影响模型的整体收敛效果。
此外,像 A10、A30 这样的设备虽然基于 PCIE 和 TCP 网络通信,但难以高效地支持大规模复杂 Dense 模型的训练。
首先,我们采用了大规模的 pass 粒度聚合,利用样本间数据的局部性原理,进行了大批量的 pass 后进行统一的去重,从而减少了单位训练所需的 ID 数量。
其次,我们进行了 embedding 的置换,因为我们无法将所有模型都放入显存中,因此需要进行批量置换。在置换过程中,如果采用最原始的方式,即与图内模型训练串行进行,会导致整个 GPU 的利用率非常低。因此,我们进行了 embedding 的置换和图内计算的流水线并行优化,使得图内计算和 embedding 的置换之间可以解耦。
第三,我们做了增量式换入换出,因为原始实现中每轮都会全量换入和换出,但 ID 量很大,全量换入换出会花费很多时间,降低模型训练效率。因此,我们利用相邻两个 pass 之间的数据局部性来提高效率。每次只换入增量的部分。也就是说,在第二个 pass 换入时,会查看 GPU 内部是否已经存在该部分,如果不存在,则只换入这部分。其次,当我们在多个 pass 中积累了足够的数据后,会将后面几轮用不到的 ID 替换掉。这样可以提高整个 ID 的换入和换出效率。
最后是 Table Fusion,与之前提到的方法类似,由于小红书的原始 embedding 列非常多,因此我们按照 embedding 维度进行了一层聚合,从而减少了图内算子的数量,提高了整体模型训练的效率。
04
AI 编译技术
1. AI 编译技术
2. 前端优化
第一个是核心的低效子图识别和高效改写;
第二个是对一些核心算子进行了高效的优化和改写。以 TensorFlow 中的 CPU Matmul 模式为例,我们采用 Onnx 的 MLAS 高效计算库进行了替换和改写,从而在核心场景中取得了较大的提升。
3. 中后端优化
4. 项目收益
05
展望未来
分享嘉宾
INTRODUCTION
曾鸣堃
小红书
训推计算引擎负责人
小红书训推计算引擎负责人,17 年毕业于中国科学技术大学,目前负责搜广推场景下的模型训练和推理引擎侧研发工作。
活动推荐
往期推荐
基于 tugraph-analytics 的实时业务数据异常归因诊断
大语言模型在图推荐系统中的融合与优化策略
Data+LLM:金融真实场景的技术创新实践
京东广告稀疏大模型训练与推理 GPU 优化实践
好的数据治理怎么做?
销售易基于 Lakehouse 的实时分析提升用户数据体验实践分享
Velox内存管理深度解析:从基础到高级特性
Apache Hudi 从零到一:全面解读写入索引(四)
Apache Hudi 从零到一:理解写入流程和操作(三)
用最酷的RAG,训最猛的大模型!
点个在看你最好看
SPRING HAS ARRIVED