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

HugeCTR

HugeCTR简介

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

HugeCTR的设计目标主要包括:

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

HugeCTR的核心特性

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

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

HugeCTR架构

快速上手HugeCTR

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

使用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

编写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()

编写训练脚本:

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")