算子优化：逼近 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-Inc/models/tree/main/RecommenderSystems）。
结语
OneEmbedding 是应训练大规模推荐系统模型需求而生的组件，灵活的分层存储、高度优化的数据流水线以及易于横向扩展的特性，能让用户轻松训练 TB 级别的推荐模型。
目前，OneFlow 框架提供了一些模型示例供你一键体验 OneEmbedding。后续，OneFlow 团队将推出覆盖业界主流模型的推荐系统模型库 Flow-Recommender ，它不仅支持分布式训练，还支持分布式推理。欢迎感兴趣的朋友关注。
- OneEmbedding 地址：https://github.com/Oneflow-Inc/models/tree/main/RecommenderSystems
- OneEmbedding 文档：https://docs.oneflow.org/master/cookies/one_embedding.html
- OneEmbedding API 文档：https://oneflow.readthedocs.io/en/master/one_embedding.html
