起源 | 机器之心
个性化举荐已成为人们获取信息的次要模式。以往，人们更多通过被动搜查本人感兴趣的信息，而当初，基于算法举荐技术的信息散发平台会自动识别用户趣味，疾速筛选信息，推送用户所感兴趣的信息。
一方面，举荐零碎大幅晋升了用户体验，另一方面，个性化散发信息更精准、高效，能够帮忙平台更精确地匹配用户和信息，大大提高流量变现效率，基于举荐技术的流量变现引擎甚至成就了万亿市值的宏大商业帝国。
从短视频信息流举荐、广告搜寻到线上购物，这些利用都构筑于精准的举荐零碎之上，背地的外围功臣就是深度学习模型。
不过，随着海量数据的积攒以及更加频繁的用户数据迭代，底层零碎可扩展性和训练速度面临严厉的挑战。人们发现，通用深度学习框架都不能间接满足工业级举荐零碎的需要，而是必须基于通用深度学习框架做深度定制，甚至于要开发专门的零碎才行。
针对古代举荐零碎的种种痛点，OneFlow 团队推出了一款高性能、可扩大、灵便度高的举荐零碎组件 OneEmbedding。它的应用形式和通用深度学习框架一样简略，性能却远超通用框架，甚至超过了 NVIDIA HugeCTR 这样为举荐场景定制开发的零碎。
具体而言，在 DCN、DeepFM 两个模型上，无论是 FP32 还是混合精度（automatic mixed-precision, AMP）训练，OneEmbedding 的性能大幅超过 HugeCTR，而在 HugeCTR 深度优化以至于有点“过拟合”的 DLRM 模型上，OneEmbedding 性能与 HugeCTR 根本持平。
（以上测试环境均为：CPU Intel(R) Xeon(R) Platinum 8336C CPU @ 2.30GHz 2；CPU Memory 1920GB；GPU NVIDIA A100-SXM-80GB 8；SSD Intel SSD D7P5510 Series 3.84TB 4）*
当用户应用 OneFlow 搭建举荐模型时，只需应用以下数行代码对 Embedding 词表进行配置即可训练含有 TB 级别词表的举荐模型：
# self.embedding = nn.Embedding(vocab_size, embedding_vec_size)
self.embedding = flow.one_embedding.MultiTableEmbedding(
"sparse_embedding",
embedding_dim=embedding_vec_size,
dtype=flow.float,
key_type=flow.int64,
tables=tables,
store_options=store_options,
)基于 OneEmbedding 搭建的常见搜寻举荐广告模型案例地址：https://github.com/Oneflow-In…
1
大规模举荐零碎的挑战
一般而言，举荐零碎须要应用相似性别、年龄、行为等方面的离散特色（sparse feature），在一个 Embedding 词表中用特色 ID 进行查表（lookup），获得对应的 Embedding 向量并送到上游应用。
罕用的公开数据集 Criteo1T 中大略蕴含十亿个特色 ID，如果 embedding_dims 配置为 128，那总共须要 512 GB 空间来包容 Embedding 参数，如果应用 Adam 优化器，因为须要保留额定的两个状态变量 m 和 v，所需存储容量就减少到 1536 GB。理论利用场景中，数据规模比 Criteo 还要高出几个数量级，模型的容量就更大了。
大规模举荐零碎的外围问题就是，如何高效经济地反对大规模 Embedding 的查问和更新 。衡量规模、老本和效率，呈现了如下三种常见的解决方案。
最常见也是最早呈现的一种解决方案是将 Embedding 全副部署在 CPU 上，利用 CPU 内存容量大、成本低的特点扩大参数规模，长处是模型规模简直能够无限大。不过，其毛病也很显著，无论是计算性能还是带宽，CPU 都远低于 GPU，导致 Embedding 局部成为显著瓶颈，往往须要数十乃至上百台 CPU 服务器能力撑持一个工业级的举荐零碎。
鉴于 GPU 在浓密计算中得天独厚的劣势，也有人倡议用 GPU 来训练大型 Embedding 模型。问题是，GPU 很贵且显存容量无限，如果应用显存容量为 40GB 的 A100 来基于 Criteo 数据训练 128 维嵌入向量，至多须要 13 张显卡能力放下 512GB 的 Embedding 词表。每张卡只有 40GB 显存容量，分布式 Embedding 须要应用所谓的模型并行技术，现实状况下，为了解决更大规模的模型只须要减少 GPU 的数量即可。
事实是，GPU 绝对于 CPU 的老本十分昂扬，并且举荐零碎中模型主体计算局部不大，模型并行在扩大过程中只是解决了 Embedding 规模的问题，训练速度的收益比拟无限，甚至会因为多设施之间引入通信导致训练速度降落，因而通常只实用于小规模集群。
为了缓解 GPU 之间传输带宽的问题，业界倒退出比以太网带宽更高的 NVSwitch、Infiniband 网络等互联技术。一方面，这意味着额定的老本，另一方面，很多用户的基础设施不具备相应革新、降级的条件。
那么，有没有鱼和熊掌兼得的计划？
针对上述计划存在的问题，OneFlow 团队设计了 OneEmbedding， 通过分层存储让单卡也能反对 TB 级模型的训练，通过横向扩大让模型容量没有天花板，通过 OneFlow 的主动流水线机制、算子优化和通信量化压缩等技术实现极致性能，在用法像 PyTorch 一样简略的前提下，OneEmbedding 在 DLRM 模型上性能是 TorchRec 的 3 倍以上，开启 TorchRec 没有反对的混合精度后，OneEmbedding 的性能更是 TorchRec 的 7 倍以上。
（TorchRec 性能数据参考 8 卡 A100 测试后果：https://github.com/facebookre…）
2
OneEmbedding 的外围劣势
分层存储：单卡也能反对 TB 级模型训练
利用数据的空间局部性和工夫局部性，多级缓存能够很好地实现性能和老本的折衷。OneEmbedding 也基于这个思维实现了多级缓存，即便用户只有一张 GPU 也能够训练 TB 级别的模型。
用户既能够把 Embedding 部署到 GPU 显存上，也能够把 Embedding 部署到 CPU 内存甚至是 SSD 上。这种计划能够施展 CPU 内存或者 SSD 更低的老本劣势，以对 Embedding 参数规模进行扩大，还能够利用 GPU 显存作为高速缓存设施，以实现高性能成果。
OneEmbedding 会动静地将最近频繁拜访的条目缓存到 GPU 显存上，同时将最近拜访频率较低的条目逐出到 CPU 内存或者 SSD 等底层存储中。在数据遵循幂律散布这一前提下，基于无效的 Cache 治理算法，OneEmbedding 能够使 GPU 缓存的命中率始终维持在较高的程度。
值得强调的是，OneEmbedding 只是将 CPU 内存和 SSD 作为存储设备，所有计算都在 GPU 上执行。目前，OneEmbedding 提供三种预置存储计划：
- 应用 GPU 显存存储模型所有参数
- 将 CPU 内存作为 Embedding 参数存储设备，并搭配应用 GPU 作为高速缓存
- 将 SSD 作为 Embedding 参数存储设备，并搭配应用 GPU 作为高速缓存
# 应用 SSD 作为存储设备，搭配 GPU 作为高速缓存
store_options = flow.one_embedding.make_cached_ssd_store_options(
cache_budget_mb=cache_memory_budget_mb,
persistent_path=persistent_path,
capacity=vocab_size,
)用户能够依据理论应用时的硬件设施状况，只需用短短数行代码进行配置，即可一箭三雕实现规模、效率、老本的最优化。
为了覆盖 CPU 和 SSD 取数据的提早，OneEmbedding 引入流水线、数据预取等技术手段，使得在以 CPU 内存 和 SSD 作为存储后端的同时，效率仍然能够和应用纯 GPU 训练那样放弃在同样的程度。
别离对三种存储计划进行测试。其中，测试用例与 MLPerf 的 DLRM 模型统一，参数规模约为 90GB。在应用 SSD 和 CPU 内存作为存储设备时，咱们配置的 GPU 缓存大小为每个 GPU 12GB，相比于 90 GB 的总参数量，只能有一部分参数放弃在 GPU 显存内，其余的参数则是保留在 CPU 内存或者 SSD 上，随着训练过程动静的换入到 GPU 缓存中来，测试后果如下图。
（测试环境：CPU Intel(R) Xeon(R) Silver 4214R CPU @ 2.40GHz 2；CPU Memory 512GB；GPU NVIDIA A100-PCIE-40GB 4；SSD Intel SSD D7P5510 Series 7.68TB 4）*
从测试后果中能够看到：
（1）纯 GPU 显存计划性能最佳，但因为 GPU 显存只有 4x40GB，实践上最大只可能训练 160 GB 模型；
（2）相比纯 GPU 显存计划，GPU 缓存 + CPU 存储的计划性能只有渺小损失，但能够把参数规模的天花板扩大到 CPU 内存容量，往往是数百 GB～数 TB；
（3）更进一步，如果能承受更大的性能损失，GPU 缓存 + SSD 存储的计划能将参数规模的天花板扩大到 SSD 的容量，模型规模可达数十 TB，甚至更大。
如果咱们想在只有一个 NVIDIA A30-24GB GPU 的服务器上进行上述 DLRM 模型的残缺训练，24G 的显存显然无奈间接训练 90GB 规模的模型。借助分层存储，应用 CPU 内存作为存储设备，GPU 显存作为高速缓存，就能够反对比 90GB 还大的模型。
横向扩大：多卡线性减速，突破模型天花板
应用分层存储技术，OneEmbedding 晋升了单卡状况下的 Embedding 参数规模极限，只有内存空间够大甚至可能训练 TB 级别大小的模型。如果模型的容量进一步扩充到甚至大大超过 CPU 内存的容量，用户还能够在多级存储的根底上借助 OneFlow 的并行能力轻松地横向拓展到多机多卡，以训练更大的模型。
在举荐零碎中，模型主体参数相较 Embedding 则小的多。因而咱们个别将 Embedding 局部设置为模型并行，模型主体设置为数据并行。通过应用多机多卡，可进一步晋升 Embedding 大小。
具体到实现细节，每个 Rank 各自负责一部分 Embedding 的存储，特色 ID 进入到各个 Rank，可能存在反复 ID 的状况，首先要进行去重（即下图的 ID Shuffle）；各个 Rank 拿着去重后的 ID 去查问 Embedding，失去对应的部分数据，所有 Rank 数据合并后各 Rank 失去残缺的 Embedding 数据（即下图的 Embedding Shuffle）；最初，各 Rank 以数据并行的形式实现整个模型训练过程。
下图展现了 OneEmbedding 采取纯 GPU 显存的策略训练 DLRM 模型时，FP32 和 AMP 配置下，不同 GPU 个数下模型吞吐量。
（测试环境：CPU Intel(R) Xeon(R) Platinum 8336C CPU @ 2.30GHz 2；CPU Memory 1920GB；GPU NVIDIA A100-SXM-80GB 8；SSD Intel SSD D7P5510 Series 3.84TB 4）*
能够看到，随着 GPU 设施数的减少，模型吞吐量均能显著减少，在混合精度状况下，单张 GPU 能有 600 万的吞吐量，当扩大到 8 张 GPU 时能有近 4000 万的吞吐量。
流水线机制：主动重叠计算和数据传输
在 DLRM 模型中，Embedding 中的 Dense Feature 会进入到 Bottom MLP 中，而 Sparse Feature 通过 Embedding 查问失去对应特色。两者进入 Interaction 进行特色穿插，最初进入到 Top MLP。
Embedding 相干操作蕴含查表（Embedding Lookup）、更新（Embedding Update）。因为 OneEmbedding 应用的是分层存储的机制，可能会遇到特色 ID 没有命中高速缓存的状况，此时，数据拉取耗时较长，会影响训练速度。
为防止这个有余，OneEmbedding 退出数据预取（Embedding Prefetch）操作，以保障查表和更新操作均能在 GPU 上执行。因为前后迭代之间的数据预取不存在依赖关系，在以后迭代计算的同时，能够预取下一个迭代须要的 Embedding 数据，实现计算和预取的重叠。
在 Embedding 数据查问替换的过程中，与 Embedding 操作无关的 Dense Feature 能够进入到 Bottom MLP 进行计算，在工夫上进行重叠。残缺的重叠执行时序如下图所示。
如此简单的数据流水线管制在传统深度学习框架里是一个很挑战的问题。不仅如此，在理论举荐场景中，用户的数据在一直变动，这要求流水线机制还要能应答动态数据。
而 OneFlow 的 Actor 机制让这所有问题变得很简略，每个 Actor 都通过本人外部的状态机和音讯机制实现了分布式协同工作。通过为每个 Actor 赋予多份存储块，不同的 Actor 能够同时工作，重叠各自的工作工夫，从而实现 Actor 之间的流水线。咱们只须要将 Embedding 操作调配到独自的一个 stream 上，即可让零碎自发地造成流水线。
算子优化：迫近 GPU 极限性能
OneFlow 团队不仅对通用算子进行了深度优化，还针对风行的举荐零碎模型特点，减少了多个高性能 CUDA 算子实现。
对于 DLRM、DCN 模型中的特色穿插局部，OneFlow 别离实现了 FusedDotFeatureInteraction 和 FusedCrossFeatureInteraction 算子。
(FusedCrossFeatureInteraction 算子，图片出自《Deep & Cross Network for Ad Click Predictions》)
对于模型中多个全连贯层局部，OneFlow 基于 cublasLt 矩阵运算库实现了 FusedMLP 算子。
而对于带 Dropout 操作的全连贯层，OneFlow 深度定制了其中的 ReluDropout 操作，应用 bitmask 模式存储前向产生的 mask，在反向流传中，通过设置 cublasLt 矩阵乘的参数 alpha=dropout_scale 以实现反向算子交融。
量化压缩：压迫通信效率
在模型训练的通信过程中，近期也有不少工作对数据进行量化压缩以节俭通信量，进步通信效率，OneEmbedding 也反对这个个性。
并行训练中，各个 Rank 之间须要通信以替换 Embedding 数据，咱们先将浮点类型数据量化成 int8 类型，替换完后再反量化复原。
下图以 DLRM 模型为例展现了抉择纯 GPU 显存存储配置，别离测试在 FP32 和混合精度条件下量化前后模型吞吐量。
量化前后模型精度比照（AUC）：
（测试环境：CPU Intel(R) Xeon(R) Silver 4214R CPU @ 2.40GHz 2；CPU Memory 512GB；GPU NVIDIA A100-PCIE-40GB 4；SSD Intel SSD D7P5510 Series 7.68TB 4）*
测试结果表明，在不影响模型精度的前提下，量化通信在 FP32 的状况下相比默认通信模式有 64% 的晋升，在混合精度的状况下有 13% 的晋升。
易用性：构建大规模举荐模型就像应用 PyTorch 一样简略
OneEmbedding 作为 OneFlow 的一个外部扩大组件，意味着用户能够在应用 OneEmbedding 的高级个性同时，还能享受 OneFlow 通用框架的灵活性构建本人的举荐模型。
class DLRMModule(nn.Module):
def __init__(self, args):
super(DLRMModule, self).__init__()
self.bottom_mlp = FusedMLP(...)
self.embedding = OneEmbedding(...)
self.interaction = FusedDotInteraction(...)
self.top_mlp = FusedMLP(...)
def forward(self, sparse_feature, dense_feature):
dense_fields = self.bottom_mlp(dense_feature)
embedding = self.embedding(sparse_feature)
features = self.interaction(dense_fields, embedding)
return self.top_mlp(features)最初，值得一提的是，OneEmbedding 通过内置的编码机制对特色 ID 进行编码，反对动静插入新增数据。用户不须要提前布局 Embedding 容量，也无需对数据集中的特色 ID 进行非凡解决。这种动静机制人造地反对增量训练的场景，同时也缩小应用累赘。
目前 OneFlow 的 models 仓库下提供了基于 OneEmbedding 搭建的一系列模型，如 DLRM, DeepFM, xDeepFM, DCN, PNN, MMoE，后续也会补充更多的举荐模型
（https://github.com/Oneflow-In…）。
3
结语
OneEmbedding 是应训练大规模举荐零碎模型需要而生的组件，灵便的分层存储、高度优化的数据流水线以及易于横向扩大的个性，能让用户轻松训练 TB 级别的举荐模型。
目前，OneFlow 框架提供了一些模型示例供你一键体验 OneEmbedding。后续，OneFlow 团队将推出笼罩业界支流模型的举荐零碎模型库 Flow-Recommender，它不仅反对分布式训练，还反对分布式推理。欢送感兴趣的敌人关注。
- OneEmbedding 地址：
https://github.com/Oneflow-In… - OneEmbedding 文档：
https://docs.oneflow.org/mast… - OneEmbedding API 文档：
https://oneflow.readthedocs.i…
欢送下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In…