HugeCTR: GPU加速的大规模推荐系统训练和推理框架

HugeCTR简介

HugeCTR是NVIDIA开发的一个开源GPU加速推荐系统框架,专门用于大规模深度学习推荐模型的高效训练和推理。它是NVIDIA Merlin推荐系统套件的核心组件之一,旨在解决推荐系统面临的计算和存储挑战。

HugeCTR的设计目标主要包括:

1. 快速 - 在包括MLPerf在内的各项推荐系统基准测试中表现出色。
2. 易用 - 提供丰富的文档、示例和notebook,方便数据科学家和机器学习从业者使用。
3. 领域专用 - 针对推荐系统的特点进行优化,支持超大规模嵌入。

HugeCTR的核心特性

HugeCTR支持多种推荐系统所需的关键功能:

1. 高层次抽象的Python接口
2. 模型并行训练
3. 优化的GPU工作流
4. 多节点训练
5. 混合精度训练
6. 嵌入训练缓存
7. GPU嵌入缓存
8. GPU/CPU内存共享机制
9. HugeCTR到ONNX的转换器
10. 分层参数服务器
11. Sparse Operation Kit

快速上手HugeCTR

要快速开始使用HugeCTR,可以按照以下步骤操作:

1. 使用NVIDIA提供的Docker镜像启动容器:

docker run --gpus=all --rm -it --cap-add SYS_NICE -v /your/host/dir:/your/container/dir -w /your/container/dir -it -u $(id -u):$(id -g) nvcr.io/nvidia/merlin/merlin-hugectr:24.06

2. 编写Python脚本生成合成数据集:

import hugectr
from hugectr.tools import DataGeneratorParams, DataGenerator

data_generator_params = DataGeneratorParams(
  format = hugectr.DataReaderType_t.Parquet,
  label_dim = 1,
  dense_dim = 13,
  num_slot = 26,
  i64_input_key = False,
  source = "./dcn_parquet/file_list.txt",
  eval_source = "./dcn_parquet/file_list_test.txt",
  slot_size_array = [39884, 39043, 17289, 7420, 20263, 3, 7120, 1543, 39884, 39043, 17289, 7420, 
                     20263, 3, 7120, 1543, 63, 63, 39884, 39043, 17289, 7420, 20263, 3, 7120,
                     1543 ],
  dist_type = hugectr.Distribution_t.PowerLaw,
  power_law_type = hugectr.PowerLaw_t.Short)
data_generator = DataGenerator(data_generator_params)
data_generator.generate()

3. 编写训练脚本:

import hugectr
from mpi4py import MPI

solver = hugectr.CreateSolver(max_eval_batches = 1280,
                              batchsize_eval = 1024,
                              batchsize = 1024,
                              lr = 0.001,
                              vvgpu = [[0]],
                              repeat_dataset = True)
reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Parquet,
                                 source = ["./dcn_parquet/file_list.txt"],
                                 eval_source = "./dcn_parquet/file_list_test.txt",
                                 slot_size_array = [39884, 39043, 17289, 7420, 20263, 3, 7120, 1543, 39884, 39043, 17289, 7420, 
                                                   20263, 3, 7120, 1543, 63, 63, 39884, 39043, 17289, 7420, 20263, 3, 7120, 1543 ])
optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam,
                                    update_type = hugectr.Update_t.Global)
model = hugectr.Model(solver, reader, optimizer)
model.add(hugectr.Input(label_dim = 1, label_name = "label",
                        dense_dim = 13, dense_name = "dense",
                        data_reader_sparse_param_array =
                        [hugectr.DataReaderSparseParam("data1", 1, True, 26)]))
model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash,
                           workspace_size_per_gpu_in_mb = 75,
                           embedding_vec_size = 16,
                           combiner = "sum",
                           sparse_embedding_name = "sparse_embedding1",
                           bottom_name = "data1",
                           optimizer = optimizer))
# 添加更多层...
model.compile()
model.summary()
model.fit(max_iter = 5120, display = 200, eval_interval = 1000, snapshot = 5000, snapshot_prefix = "dcn")

4. 运行训练脚本:

python dcn_parquet_train.py

通过以上步骤,我们就可以快速上手HugeCTR,训练一个简单的推荐模型。

HugeCTR的高级功能

除了基本的模型训练,HugeCTR还提供了许多高级功能来优化推荐系统的性能:

1. 模型并行训练

HugeCTR支持在多个GPU上并行训练模型,可以有效处理超大规模的嵌入表。它提供了多种策略来在GPU之间分配嵌入表:

• 模型并行嵌入表(Model-Parallel Embedding Table)
• 数据并行嵌入表(Data-Parallel Embedding Table)
• 本地化嵌入表(Localized Embedding Table)

2. 混合精度训练

HugeCTR支持FP32、FP16和TF32精度的混合使用,可以在保证精度的同时提高训练速度。

3. 嵌入训练缓存

对于超大规模的嵌入表,HugeCTR提供了嵌入训练缓存机制,可以将部分嵌入存储在CPU内存中,按需加载到GPU。

4. 分层参数服务器

HugeCTR的分层参数服务器架构可以有效管理大规模嵌入表,支持多GPU和多节点部署。

5. Sparse Operation Kit

Sparse Operation Kit (SOK)是HugeCTR的一个子项目,提供了GPU加速的稀疏操作,可以与TensorFlow等框架集成使用。

HugeCTR在工业界的应用

HugeCTR已经在多家科技公司得到了实际应用:

• 腾讯广告系统: 使用HugeCTR将CTR预估模型的训练速度提升了7倍以上。
• 美团机器学习平台: 利用HugeCTR优化了推荐模型的训练流程。

这些案例都证明了HugeCTR在处理大规模推荐系统时的优势。

总结

HugeCTR作为一个专门针对推荐系统优化的GPU加速框架,为构建高性能推荐系统提供了强大的工具。它不仅支持高效的模型训练,还提供了完整的推理解决方案。对于需要处理大规模数据和复杂模型的推荐系统开发者来说,HugeCTR无疑是一个值得尝试的选择。

随着推荐系统在各行各业的广泛应用,以及模型规模和数据量的不断增长,像HugeCTR这样的高性能框架将发挥越来越重要的作用。我们期待看到HugeCTR在未来能够支持更多新的模型结构和训练技术,为推荐系统的发展做出更大的贡献。

如果你对HugeCTR感兴趣,可以访问HugeCTR的GitHub仓库获取更多信息,或者加入HugeCTR的Slack频道与其他开发者交流讨论。让我们一起探索GPU加速推荐系统的无限可能吧!