研究摘要
深度学习推荐模型（DLRM）在个性化推荐、广告投放及电子商务等领域发挥着重要作用。然而，DLRM 的训练过程受限于嵌入表的高内存占用以及分布式训练中的通信开销，导致计算效率低下。现有方法如张量训练（TT）分解尽管能够有效压缩嵌入表，但引入了额外的计算开销，同时传统的分布式训练框架亦面临数据传输瓶颈。针对上述问题，本研究提出 EcoRec 框架，该框架结合 TT 分解与分布式训练，通过优化计算模式减少冗余计算，并引入基于排序索引的微批量技术，以降低存储需求而不增加计算开销。此外，EcoRec 采用嵌入层流水线训练系统，实现数据均衡分配并优化通信效率。实验结果表明，相较于现有方法，EcoRec 可实现最高 3.1倍 的加速效果，并降低 38.5% 的内存使用量。本研究的成果为高性能 DLRM 训练提供了新的解决方案。
研究背景
近年来，深度学习推荐模型（DLRM）凭借其强大的特征交互建模能力，在互联网个性化推荐等应用中获得了广泛关注。DLRM 主要依赖于嵌入表存储稀疏特征，并利用多层感知机（MLP）处理密集特征。然而，随着数据规模的增长，嵌入表的存储需求呈指数级增长，单个 GPU 的高带宽存储难以满足大规模模型的训练需求。此外，分布式训练过程中，由于嵌入向量及梯度交换的通信开销较高，进一步加剧了训练效率的瓶颈。
如图 1 所示，DLRM 训练架构主要包括嵌入层和多层感知机（MLP）层，其中嵌入层存储并查询稀疏特征的嵌入向量，而 MLP 负责密集特征的计算。图中标注的通信路径表明，嵌入层的分布式存储与梯度同步需要大量数据传输，成为训练过程中的主要瓶颈。
因此，如何在保证模型精度的前提下，提高 DLRM 训练的计算与通信效率，成为当前研究的关键问题。
图 1 DLRM 模型架构示意图
研究现状
当前的研究主要集中在以下两类方法：（1）基于嵌入表压缩的垂直方法，例如 TT-Rec 和 EL-Rec 通过张量分解技术减少嵌入表存储需求（图1），但计算复杂度较高，影响训练效率；（2）基于分布式训练的水平方法，例如 Nvidia HugeCTR 和 Neo 采用混合并行策略，在多个 GPU 之间分布存储嵌入表，以缓解单个 GPU 的存储压力，如图2所示。然而，这些方法依赖 AlltoAll 通信操作进行数据交换，在大规模训练时通信开销较高，难以满足 DLRM 高效训练的需求。因此，亟需一种综合考虑嵌入表压缩与高效并行计算的解决方案，以提升 DLRM 训练的整体性能。
图2 分布式模型并行示意图
研究方法
针对现有方法的不足，本研究提出 EcoRec 框架，以优化 DLRM 训练过程的计算和通信效率。EcoRec 采用了一种新的 TT 计算模式，称为收缩-乘法（C-M）与乘法-收缩（M-C）计算模式。这种模式的核心思想是通过合并冗余计算操作，减少张量乘法的计算负担，如图 3 所示。传统 TT 计算模式会对相同的嵌入索引进行多次计算，而 C-M/M-C 模式能够有效识别并消除这些冗余计算，从而提升计算效率。
图3 基线方法与C-M/M-C模式对比
此外，EcoRec 还引入了基于排序索引的微批量技术。传统的批量训练方式容易导致计算负载不均衡，而排序索引技术能够根据索引访问频率进行重新排列，使得同一批次内的数据更加均匀地分布在不同计算设备之间。这一方法在降低存储需求的同时，不会增加计算开销。
为了进一步优化通信效率，EcoRec 采用了一种新的嵌入层流水线训练系统，如图 4 所示。传统的 AlltoAll 通信方式通常在计算任务完成后再进行通信，而 EcoRec 的流水线系统能够在计算过程中逐步进行通信，实现计算与数据传输的并行化，从而减少通信延迟，提高分布式 DLRM 训练的吞吐量。
在分布式训练环境下，不同计算设备的负载不均衡问题也会影响训练效率。因此，EcoRec 还采用了一种基于统计分析的负载均衡策略，以确保计算任务在各个 GPU 之间均匀分配。通过动态调整嵌入表的划分方式，该策略能够有效避免单个计算设备成为性能瓶颈，提高整体资源利用率。
图 4 流水并行示意图
研究成果
在 32 GPU 集群环境下，我们对 EcoRec 进行了全面评测，并与现有 DLRM 训练框架进行对比。实验结果表明，EcoRec 在多个方面实现了显著的优化效果。如图5所示，首先，在训练加速比方面，EcoRec 相较于 EL-Rec 实现了最高 3.1x 的训练加速，相较于 TT-Rec 则提升了 6.82x。这一显著的性能提升表明，EcoRec 的计算优化方法能够有效减少计算冗余，提高训练吞吐量。其次，在内存优化方面，EcoRec 通过微批量优化方法，使得嵌入表的峰值内存使用量降低了 38.5%。这一优化使得单个 GPU 可以训练更大规模的模型，从而提升了硬件资源的利用效率。本研究通过 EcoRec 框架的提出，成功解决了 DLRM 训练中的计算、存储及通信瓶颈，为大规模深度学习推荐模型的高效训练提供了重要的理论和实践支持，并在未来可进一步推广至其他深度学习任务的优化中。
图 5 模型对比实验。左图表示相对加速比，右图表示显存占用量
相关研究成果已收录于 CCF-A 类顶级会议 SC 24（The International Conference for High Performance Computing, Networking, Storage, and Analysis），论文题为 Accelerating Distributed DLRM Training with Optimized TT Decomposition and Micro-Batching。