在这篇文章中，我们介绍了 HugeCTR，这是一个面向行业的推荐系统训练框架，针对具有模型并行嵌入和数据并行密集网络的大规模 CTR 模型进行了优化。
本文以GitHub 源码文档 https://github.com/NVIDIA-Merlin/HugeCTR/blob/master/docs/python_interface.md 的翻译为基础，并且结合源码进行分析。其中借鉴了HugeCTR源码阅读 这篇大作，特此感谢。
为了更好的说明，下面类定义之中，只保留其成员变量，成员函数会等到分析时候才会给出。
本系列其他代码为：
[源码解析] NVIDIA HugeCTR，GPU 版本参数服务器 --(1)
HugeCTR 训练的过程可以看作是数据并行+模型并行。
训练流程如下：
首先构建三级流水线，初始化模型网络。初始化参数和优化器状态。
Reader 会从数据集加载一个 batch 的数据，放入 Host 内存之中。
开始解析数据，得到 sparse 参数，dense 参数，label 等等。
嵌入层进行前向传播，即从参数服务器读取 embedding，进行处理。
对于网络层进行前向传播和后向传播，具体区分是多卡，单卡，多机，单机等。
嵌入层反向操作。
多卡之间交换 dense 参数的梯度。
嵌入层更新 sparse 参数。就是把反向计算得到的参数梯度推送到参数服务器，由参数服务器根据梯度更新参数。
我们从一个例子中可以看到，总体逻辑和单机很像，就是解析配置，使用 session 来读取数据，训练等等，其中 vvgpu 是 device map。
# train.pyimport sysimport hugectrfrom mpi4py import MPIdef train(json_config_file): solver_config = hugectr.solver_parser_helper(batchsize = 16384, batchsize_eval = 16384, vvgpu = [[0,1,2,3,4,5,6,7]], repeat_dataset = True) sess = hugectr.Session(solver_config, json_config_file) sess.start_data_reading() for i in range(10000): sess.train() if (i % 100 == 0): loss = sess.get_current_loss()if __name__ == "__main__": json_config_file = sys.argv[1] train(json_config_file)既然知道了 Session 是核心，我们就通过 Session 看看如何构建 HugeCTR。
我们首先看看Session的定义，只保留其成员变量，可以看到其主要是：
class Session { public: Session(const SolverParser& solver_config, const std::string& config_file); Session(const Session&) = delete; Session& operator=(const Session&) = delete; private: std::vector<std::shared_ptr<Network>> networks_; /**< networks (dense) used in training. */ std::vector<std::shared_ptr<IEmbedding>> embeddings_; /**< embedding */ std::shared_ptr<IDataReader> init_data_reader_; std::shared_ptr<IDataReader> train_data_reader_; /**< data reader to reading data from data set to embedding. */ std::shared_ptr<IDataReader> evaluate_data_reader_; /**< data reader for evaluation. */ std::shared_ptr<ResourceManager> resource_manager_; /**< GPU resources include handles and streams etc.*/ std::shared_ptr<Parser> parser_; std::shared_ptr<ExchangeWgrad> exchange_wgrad_; metrics::Metrics metrics_; SolverParser solver_config_; struct HolisticCudaGraph { std::vector<bool> initialized; std::vector<cudaGraphExec_t> instance; std::vector<cudaEvent_t> fork_event; } train_graph_; // TODO: these two variables for export_predictions. // There may be a better place for them. bool use_mixed_precision_; size_t batchsize_eval_;};构造函数大致分为以下步骤：
Session::Session(const SolverParser& solver_config, const std::string& config_file) : resource_manager_(ResourceManagerExt::create(solver_config.vvgpu, solver_config.seed, solver_config.device_layout)), solver_config_(solver_config) { // 检查设备 for (auto dev : resource_manager_->get_local_gpu_device_id_list()) { if (solver_config.use_mixed_precision) { check_device(dev, 7, 0); // to support mixed precision training earliest supported device is CC=70 } else { check_device(dev, 6, 0); // earliest supported device is CC=60 } } // 生成 Parser，用来解析配置 parser_.reset(new Parser(config_file, solver_config.batchsize, solver_config.batchsize_eval, solver_config.num_epochs < 1, solver_config.i64_input_key, solver_config.use_mixed_precision, solver_config.enable_tf32_compute, solver_config.scaler, solver_config.use_algorithm_search, solver_config.use_cuda_graph)); // 建立流水线 parser_->create_pipeline(init_data_reader_, train_data_reader_, evaluate_data_reader_, embeddings_, networks_, resource_manager_, exchange_wgrad_);#ifndef DATA_READING_TEST#pragma omp parallel num_threads(networks_.size()) { // 多线程并行初始化模型 size_t id = omp_get_thread_num(); networks_[id]->initialize(); if (solver_config.use_algorithm_search) { networks_[id]->search_algorithm(); } CK_CUDA_THROW_(cudaStreamSynchronize(resource_manager_->get_local_gpu(id)->get_stream())); }#endif // 加载dense feature需要的参数 init_or_load_params_for_dense_(solver_config.model_file); // 加载sparse feature需要的参数 init_or_load_params_for_sparse_(solver_config.embedding_files); // 加载信息 load_opt_states_for_sparse_(solver_config.sparse_opt_states_files); load_opt_states_for_dense_(solver_config.dense_opt_states_file); int num_total_gpus = resource_manager_->get_global_gpu_count(); for (const auto& metric : solver_config.metrics_spec) { metrics_.emplace_back( std::move(metrics::Metric::Create(metric.first, solver_config.use_mixed_precision, solver_config.batchsize_eval / num_total_gpus, solver_config.max_eval_batches, resource_manager_))); } if (solver_config_.use_holistic_cuda_graph) { train_graph_.initialized.resize(networks_.size(), false); train_graph_.instance.resize(networks_.size()); for (size_t i = 0; i < resource_manager_->get_local_gpu_count(); i++) { auto& gpu_resource = resource_manager_->get_local_gpu(i); CudaCPUDeviceContext context(gpu_resource->get_device_id()); cudaEvent_t event; CK_CUDA_THROW_(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); train_graph_.fork_event.push_back(event); } } if (embeddings_.size() == 1) { auto lr_scheds = embeddings_[0]->get_learning_rate_schedulers(); for (size_t i = 0; i < lr_scheds.size(); i++) { networks_[i]->set_learning_rate_scheduler(lr_scheds[i]); } }}这里有几个相关类需要注意一下。
我们首先看看 ResourceManager。
首先是两个接口，ResourceManagerBase 是最顶层接口，ResourceManager 进行了扩展。
/** * @brief Top-level ResourceManager interface * * The top level resource manager interface shared by various components */class ResourceManagerBase { public: virtual void set_local_gpu(std::shared_ptr<GPUResource> gpu_resource, size_t local_gpu_id) = 0; virtual const std::shared_ptr<GPUResource>& get_local_gpu(size_t local_gpu_id) const = 0; virtual size_t get_local_gpu_count() const = 0; virtual size_t get_global_gpu_count() const = 0;};/** * @brief Second-level ResourceManager interface * * The second level resource manager interface shared by training and inference */class ResourceManager : public ResourceManagerBase { // 省略了函数定义}然后是核心实现：ResourceManagerCore，这里记录了各种资源。
/** * @brief GPU resources manager which holds the minimal, essential set of resources * * A core GPU Resource manager */class ResourceManagerCore : public ResourceManager { private: int num_process_; int process_id_; DeviceMap device_map_; std::shared_ptr<CPUResource> cpu_resource_; std::vector<std::shared_ptr<GPUResource>> gpu_resources_; /**< GPU resource vector */ std::vector<std::vector<bool>> p2p_matrix_; std::vector<std::shared_ptr<rmm::mr::device_memory_resource>> base_cuda_mr_; std::vector<std::shared_ptr<rmm::mr::device_memory_resource>> memory_resource_;}ResourceManagerExt 是在 ResourceManagerCore 基础之上进行再次封装，其核心就是 core_，这是一个 ResourceManagerCore 类型。我们用 ResourceManagerExt 来分析。
/** * @brief GPU resources manager which holds all the resources required by training * * An extended GPU Resource manager */class ResourceManagerExt : public ResourceManager { std::shared_ptr<ResourceManager> core_;#ifdef ENABLE_MPI std::unique_ptr<IbComm> ib_comm_ = NULL;#endif std::shared_ptr<AllReduceInPlaceComm> ar_comm_ = NULL;};其创建代码如下，可以看到其利用 MPI 做了一些通信上的配置：
std::shared_ptr<ResourceManager> ResourceManagerExt::create( const std::vector<std::vector<int>>& visible_devices, unsigned long long seed, DeviceMap::Layout layout) { int size = 1, rank = 0;#ifdef ENABLE_MPI HCTR_MPI_THROW(MPI_Comm_size(MPI_COMM_WORLD, &size)); HCTR_MPI_THROW(MPI_Comm_rank(MPI_COMM_WORLD, &rank));#endif DeviceMap device_map(visible_devices, rank, layout); std::random_device rd; if (seed == 0) { seed = rd(); }#ifdef ENABLE_MPI HCTR_MPI_THROW(MPI_Bcast(&seed, 1, MPI_UNSIGNED_LONG_LONG, 0, MPI_COMM_WORLD));#endif std::shared_ptr<ResourceManager> core( new ResourceManagerCore(size, rank, std::move(device_map), seed)); return std::shared_ptr<ResourceManager>(new ResourceManagerExt(core));}ResourceManagerExt::ResourceManagerExt(std::shared_ptr<ResourceManager> core) : core_(core) {#ifdef ENABLE_MPI int num_process = get_num_process(); if (num_process > 1) { int process_id = get_process_id(); ib_comm_ = std::make_unique<IbComm>(); ib_comm_->init(num_process, get_local_gpu_count(), process_id, get_local_gpu_device_id_list()); }#endif}void ResourceManagerExt::set_ar_comm(AllReduceAlgo algo, bool use_mixed_precision) { int num_process = get_num_process();#ifdef ENABLE_MPI ar_comm_ = AllReduceInPlaceComm::create(num_process, algo, use_mixed_precision, get_local_gpus(), ib_comm_.get());#else ar_comm_ = AllReduceInPlaceComm::create(num_process, algo, use_mixed_precision, get_local_gpus());#endif}具体资源上的配置还是调用了<ResourceManager> core_ 来完成。
// from ResourceManagerBasevoid set_local_gpu(std::shared_ptr<GPUResource> gpu_resource, size_t local_gpu_id) override { core_->set_local_gpu(gpu_resource, local_gpu_id);}const std::shared_ptr<GPUResource>& get_local_gpu(size_t local_gpu_id) const override { return core_->get_local_gpu(local_gpu_id);}size_t get_local_gpu_count() const override { return core_->get_local_gpu_count(); }size_t get_global_gpu_count() const override { return core_->get_global_gpu_count(); }// from ResourceManagerint get_num_process() const override { return core_->get_num_process(); }int get_process_id() const override { return core_->get_process_id(); }前面提到了Parser，我们接下来就看看。Parser 负责解析配置文件，建立流水线。其类似的支撑文件还有 SolverParser，Solver，InferenceParser 等等。可以说，Parser 是自动化运作的关键，是支撑系统的灵魂。
/** * @brief The parser of configure file (in json format). * * The builder of each layer / optimizer in HugeCTR. * Please see User Guide to learn how to write a configure file. * @verbatim * Some Restrictions: * 1. Embedding should be the first element of layers. * 2. layers should be listed from bottom to top. * @endverbatim */class Parser { private: nlohmann::json config_; /**< configure file. */ size_t batch_size_; /**< batch size. */ size_t batch_size_eval_; /**< batch size. */ const bool repeat_dataset_; const bool i64_input_key_{false}; const bool use_mixed_precision_{false}; const bool enable_tf32_compute_{false}; const float scaler_{1.f}; const bool use_algorithm_search_; const bool use_cuda_graph_; bool grouped_all_reduce_ = false;}我们接下来的这些分析，其实都是调用了 Parser 或者其相关类。
我们首先看看配置文件，看看其中是如何组织一个模型网络，这里以 test/scripts/deepfm_8gpu.json 为例。
这里需要说明一下 json 字段的作用：
bottom_names: 本层的输入张量名字。top_names: 本层的输出张量名字。所以，模型就是通过 bottom 和 top 从下往上组织起来的。
输入层如下，dense 是 slice 层的输入，Sparse 是sparse_embedding1 的输入，其中包含了 26 个 slots。
{ "name": "data", "type": "Data", "source": "./file_list.txt", "eval_source": "./file_list_test.txt", "check": "Sum", "label": { "top": "label", "label_dim": 1 }, "dense": { "top": "dense", "dense_dim": 13 }, "sparse": [ { "top": "data1", "slot_num": 26, "is_fixed_length": false, "nnz_per_slot": 2 } ]},此时模型图如下：
我们看看其定义：
embedding_vec_size 是向量维度。
combiner ：查找得到向量之后，如何做pooling，是做sum还是avg。
workspace_size_per_gpu_in_mb ：每个GPU之上的内存大小。
{ "name": "sparse_embedding1", "type": "DistributedSlotSparseEmbeddingHash", "bottom": "data1", "top": "sparse_embedding1", "sparse_embedding_hparam": { "embedding_vec_size": 11, "combiner": "sum", "workspace_size_per_gpu_in_mb": 10 }},此时模型如下：
这里我们把其它层也包括进来，就是目前输入数据和嵌入层的再上一层，我们省略了很多层，这里只是给大家一个大致的逻辑。
Reshape 层把一个 3D 输入转换为 2D 形状。此层是嵌入层的消费者。
{ "name": "reshape1", "type": "Reshape", "bottom": "sparse_embedding1", "top": "reshape1", "leading_dim": 11},Slice 层把一个bottom分解成多个top。
{ "name": "slice2", "type": "Slice", "bottom": "dense", "ranges": [ [ 0, 13 ], [ 0, 13 ] ], "top": [ "slice21", "slice22" ]},这就是我们最终的损失层，label直接会输出到这里。
{ "name": "loss", "type": "BinaryCrossEntropyLoss", "bottom": [ "add", "label" ], "top": "loss"}目前逻辑如下，是从下往上组织的模型，我们省略了其他部分：
我们对每个层进行精简，省略内部标签，把配置文件中所有层都整理出来，看看一个DeepFM在HugeCTR之中的整体架构。
我们接着看如何建立流水线。Create_pipeline 函数是用来构建流水线的，其就是转移给了create_pipeline_internal 方法。
void Parser::create_pipeline(std::shared_ptr<IDataReader>& init_data_reader, std::shared_ptr<IDataReader>& train_data_reader, std::shared_ptr<IDataReader>& evaluate_data_reader, std::vector<std::shared_ptr<IEmbedding>>& embeddings, std::vector<std::shared_ptr<Network>>& networks, const std::shared_ptr<ResourceManager>& resource_manager, std::shared_ptr<ExchangeWgrad>& exchange_wgrad) { if (i64_input_key_) { create_pipeline_internal<long long>(init_data_reader, train_data_reader, evaluate_data_reader, embeddings, networks, resource_manager, exchange_wgrad); } else { create_pipeline_internal<unsigned int>(init_data_reader, train_data_reader, evaluate_data_reader, embeddings, networks, resource_manager, exchange_wgrad); }}create_pipeline_internal 主要包含了四步：
template <typename TypeKey>void Parser::create_pipeline_internal(std::shared_ptr<IDataReader>& init_data_reader, std::shared_ptr<IDataReader>& train_data_reader, std::shared_ptr<IDataReader>& evaluate_data_reader, std::vector<std::shared_ptr<IEmbedding>>& embeddings, std::vector<std::shared_ptr<Network>>& networks, const std::shared_ptr<ResourceManager>& resource_manager, std::shared_ptr<ExchangeWgrad>& exchange_wgrad) { try { // 建立allreduce通信相关 create_allreduce_comm(resource_manager, exchange_wgrad); std::map<std::string, SparseInput<TypeKey>> sparse_input_map; std::vector<TensorEntry> train_tensor_entries_list[resource_manager->get_local_gpu_count()]; std::vector<TensorEntry> evaluate_tensor_entries_list[resource_manager->get_local_gpu_count()]; { if (!networks.empty()) { CK_THROW_(Error_t::WrongInput, "vector network is not empty"); } // 校验网络 auto j_layers_array = get_json(config_, "layers"); auto j_optimizer = get_json(config_, "optimizer"); check_graph(tensor_active_, j_layers_array); // Create Data Reader // 建立 Data Reader { // TODO: In using AsyncReader, if the overlap is disabled, // scheduling the data reader should be off. // THe scheduling needs to be generalized. auto j_solver = get_json(config_, "solver"); auto enable_overlap = get_value_from_json_soft<bool>(j_solver, "enable_overlap", false); const nlohmann::json& j = j_layers_array[0]; create_datareader<TypeKey>()(j, sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list, init_data_reader, train_data_reader, evaluate_data_reader, batch_size_, batch_size_eval_, use_mixed_precision_, repeat_dataset_, enable_overlap, resource_manager); } // Create Data Reader // Create Embedding { for (unsigned int i = 1; i < j_layers_array.size(); i++) { // 网路配置的每层是从底到上，因此只要遇到非嵌入层，就不检查其后的层了 // if not embedding then break const nlohmann::json& j = j_layers_array[i]; auto embedding_name = get_value_from_json<std::string>(j, "type"); Embedding_t embedding_type; if (!find_item_in_map(embedding_type, embedding_name, EMBEDDING_TYPE_MAP)) { Layer_t layer_type; if (!find_item_in_map(layer_type, embedding_name, LAYER_TYPE_MAP) && !find_item_in_map(layer_type, embedding_name, LAYER_TYPE_MAP_MP)) { CK_THROW_(Error_t::WrongInput, "No such layer: " + embedding_name); } break; } // 建立嵌入层 if (use_mixed_precision_) { create_embedding<TypeKey, __half>()( sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list, embeddings, embedding_type, config_, resource_manager, batch_size_, batch_size_eval_, exchange_wgrad, use_mixed_precision_, scaler_, j, use_cuda_graph_, grouped_all_reduce_); } else { create_embedding<TypeKey, float>()( sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list, embeddings, embedding_type, config_, resource_manager, batch_size_, batch_size_eval_, exchange_wgrad, use_mixed_precision_, scaler_, j, use_cuda_graph_, grouped_all_reduce_); } } // for () } // Create Embedding // 建立网络层 // create network int total_gpu_count = resource_manager->get_global_gpu_count(); if (0 != batch_size_ % total_gpu_count) { CK_THROW_(Error_t::WrongInput, "0 != batch_size\%total_gpu_count"); } // create network，在每张GPU卡之中构建一个network副本 for (size_t i = 0; i < resource_manager->get_local_gpu_count(); i++) { networks.emplace_back(Network::create_network( j_layers_array, j_optimizer, train_tensor_entries_list[i], evaluate_tensor_entries_list[i], total_gpu_count, exchange_wgrad, resource_manager->get_local_cpu(), resource_manager->get_local_gpu(i), use_mixed_precision_, enable_tf32_compute_, scaler_, use_algorithm_search_, use_cuda_graph_, false, grouped_all_reduce_)); } } exchange_wgrad->allocate(); // 建立梯度交换类 } catch (const std::runtime_error& rt_err) { std::cerr << rt_err.what() << std::endl; throw; }}create_allreduce_comm 的功能是设置通信算法，比如建立 AllReduceInPlaceComm，构建 GroupedExchangeWgrad。
void Parser::create_allreduce_comm(const std::shared_ptr<ResourceManager>& resource_manager, std::shared_ptr<ExchangeWgrad>& exchange_wgrad) { auto ar_algo = AllReduceAlgo::NCCL; bool grouped_all_reduce = false; // 获取通信算法配置 if (has_key_(config_, "all_reduce")) { auto j_all_reduce = get_json(config_, "all_reduce"); std::string ar_algo_name = "Oneshot"; if (has_key_(j_all_reduce, "algo")) { ar_algo_name = get_value_from_json<std::string>(j_all_reduce, "algo"); } if (has_key_(j_all_reduce, "grouped")) { grouped_all_reduce = get_value_from_json<bool>(j_all_reduce, "grouped"); } if (!find_item_in_map(ar_algo, ar_algo_name, ALLREDUCE_ALGO_MAP)) { CK_THROW_(Error_t::WrongInput, "All reduce algo unknown: " + ar_algo_name); } } // 设置通信算法，比如建立 AllReduceInPlaceComm resource_manager->set_ar_comm(ar_algo, use_mixed_precision_); // 构建 GroupedExchangeWgrad grouped_all_reduce_ = grouped_all_reduce; if (grouped_all_reduce_) { if (use_mixed_precision_) { exchange_wgrad = std::make_shared<GroupedExchangeWgrad<__half>>(resource_manager); } else { exchange_wgrad = std::make_shared<GroupedExchangeWgrad<float>>(resource_manager); } } else { if (use_mixed_precision_) { exchange_wgrad = std::make_shared<NetworkExchangeWgrad<__half>>(resource_manager); } else { exchange_wgrad = std::make_shared<NetworkExchangeWgrad<float>>(resource_manager); } }}其中 GroupedExchangeWgrad 是用来交换梯度的。
template <typename TypeFP>class GroupedExchangeWgrad : public ExchangeWgrad { public: const BuffPtrs<TypeFP>& get_network_wgrad_buffs() const { return network_wgrad_buffs_; } const BuffPtrs<TypeFP>& get_embed_wgrad_buffs() const { return embed_wgrad_buffs_; } void allocate() final; void update_embed_wgrad_size(size_t size) final; void allreduce(size_t device_id, cudaStream_t stream); GroupedExchangeWgrad(const std::shared_ptr<ResourceManager>& resource_manager); ~GroupedExchangeWgrad() = default; private: BuffPtrs<TypeFP> network_wgrad_buffs_; BuffPtrs<TypeFP> embed_wgrad_buffs_; std::vector<std::shared_ptr<GeneralBuffer2<CudaAllocator>>> bufs_; std::shared_ptr<ResourceManager> resource_manager_; AllReduceInPlaceComm::Handle ar_handle_; size_t network_wgrad_size_ = 0; size_t embed_wgrad_size_ = 0; size_t num_gpus_ = 0;};比如通过allreduce进行交换：
template <typename T>void GroupedExchangeWgrad<T>::allreduce(size_t device_id, cudaStream_t stream) { auto ar_comm = resource_manager_->get_ar_comm(); ar_comm->all_reduce(ar_handle_, stream, device_id);}DataReader 是流水线的主体，它实际包含了流水线的前两级：data reader worker 与 data collector。
create_datareader 的调用如下，
create_datareader<TypeKey>()(j, sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list, init_data_reader, train_data_reader, evaluate_data_reader, batch_size_, batch_size_eval_, use_mixed_precision_, repeat_dataset_, enable_overlap, resource_manager);回忆一下，在下面代码之中会调用到create_datareader创建了几个 reader。
parser_->create_pipeline(init_data_reader_, train_data_reader_, evaluate_data_reader_, embeddings_, networks_, resource_manager_, exchange_wgrad_);其实就是 Session 之中的几个成员变量，比如：
std::shared_ptr<IDataReader> init_data_reader_;std::shared_ptr<IDataReader> train_data_reader_; /**< data reader to reading data from data set to embedding. */std::shared_ptr<IDataReader> evaluate_data_reader_; /**< data reader for evaluation. */分别用于训练，评估。
因为代码太长，我们只保留部分关键代码。我们先看create_datareader里面做了什么：这里有两个 reader，一个train_data_reader和一个evaluate_data_reader，也就是一个用于训练，一个用于评估。然后会为他们建立workgroup。
对于Reader，HugeCTR 提供了三种实现：
template <typename TypeKey>void create_datareader<TypeKey>::operator()( switch (format) { case DataReaderType_t::Norm: { bool start_right_now = repeat_dataset; train_data_reader->create_drwg_norm(source_data, check_type, start_right_now); evaluate_data_reader->create_drwg_norm(eval_source, check_type, start_right_now); break; } case DataReaderType_t::Raw: { const auto num_samples = get_value_from_json<long long>(j, "num_samples"); const auto eval_num_samples = get_value_from_json<long long>(j, "eval_num_samples"); std::vector<long long> slot_offset = f(); bool float_label_dense = get_value_from_json_soft<bool>(j, "float_label_dense", false); train_data_reader->create_drwg_raw(source_data, num_samples, float_label_dense, true, false); evaluate_data_reader->create_drwg_raw(eval_source, eval_num_samples, float_label_dense, false, false); break; } case DataReaderType_t::Parquet: { // @Future: Should be slot_offset here and data_reader ctor should // be TypeKey not long long std::vector<long long> slot_offset = f(); train_data_reader->create_drwg_parquet(source_data, slot_offset, true); evaluate_data_reader->create_drwg_parquet(eval_source, slot_offset, true); break; } }}我们以 norm 为例进行解析，首先提一下，其内部建立了 WorkerGroup。
void create_drwg_norm(std::string file_name, Check_t check_type, bool start_reading_from_beginning = true) override { source_type_ = SourceType_t::FileList; worker_group_.reset(new DataReaderWorkerGroupNorm<TypeKey>( thread_buffers_, resource_manager_, file_name, repeat_, check_type, params_, start_reading_from_beginning)); file_name_ = file_name;}在 DataReaderWorkerGroupNorm 之中，建立了DataReaderWorker，其中 file_list_ 是需要读取的数据文件。
template <typename TypeKey>class DataReaderWorkerGroupNorm : public DataReaderWorkerGroup { std::string file_list_; /**< file list of data set */ std::shared_ptr<Source> create_source(size_t worker_id, size_t num_worker, const std::string &file_name, bool repeat) override { return std::make_shared<FileSource>(worker_id, num_worker, file_name, repeat); } public: // Ctor DataReaderWorkerGroupNorm(const std::vector<std::shared_ptr<ThreadBuffer>> &output_buffers, const std::shared_ptr<ResourceManager> &resource_manager_, std::string file_list, bool repeat, Check_t check_type, const std::vector<DataReaderSparseParam> ¶ms, bool start_reading_from_beginning = true) : DataReaderWorkerGroup(start_reading_from_beginning, DataReaderType_t::Norm) { int num_threads = output_buffers.size(); size_t local_gpu_count = resource_manager_->get_local_gpu_count(); // create data reader workers int max_feature_num_per_sample = 0; for (auto ¶m : params) { max_feature_num_per_sample += param.max_feature_num; } set_resource_manager(resource_manager_); for (int i = 0; i < num_threads; i++) { std::shared_ptr<IDataReaderWorker> data_reader(new DataReaderWorker<TypeKey>( i, num_threads, resource_manager_->get_local_gpu(i % local_gpu_count), &data_reader_loop_flag_, output_buffers[i], file_list, max_feature_num_per_sample, repeat, check_type, params)); data_readers_.push_back(data_reader); } create_data_reader_threads(); }};然后创建了多个线程 data_reader_threads_ 分别运行这些 woker。
/** * Create threads to run data reader workers>>>>>>> v3.1_preview */ void create_data_reader_threads() { size_t local_gpu_count = resource_manager_->get_local_gpu_count(); for (size_t i = 0; i < data_readers_.size(); ++i) { auto local_gpu = resource_manager_->get_local_gpu(i % local_gpu_count); data_reader_threads_.emplace_back(data_reader_thread_func_, data_readers_[i], &data_reader_loop_flag_, local_gpu->get_device_id()); } }我们总结一下。DataReader 包含了流水线的前两级，目前分析之涉及到了第一级。在 Reader之中，有一个 worker group，里面包含了若干worker，也有若干对应线程来运行这些 worker， Data Reader worker 就是流水线第一级。第二级 collecotr 我们会暂时跳过去在下一章进行介绍。
我们直接调过来看流水线第三级，如下代码建立了嵌入。
create_embedding<TypeKey, float>()( sparse_input_map, train_tensor_entries_list, evaluate_tensor_entries_list, embeddings, embedding_type, config_, resource_manager, batch_size_, batch_size_eval_, exchange_wgrad, use_mixed_precision_, scaler_, j, use_cuda_graph_, grouped_all_reduce_);这里建立了一些embedding，比如DistributedSlotSparseEmbeddingHash。
如前文所述，HugeCTR 包含了若干 Hash，比如：
LocalizedSlotEmbeddingHash：同一个槽（特征域）中的特征会存储在一个GPU中，这就是为什么它被称为“本地化槽”，根据槽的索引号，不同的槽可能存储在不同的GPU中。
DistributedSlotEmbeddingHash：所有特征都存储于不同特征域/槽上，不管槽索引号是多少，这些特征都根据特征的索引号分布到不同的GPU上。这意味着同一插槽中的特征可能存储在不同的 GPU 中，这就是将其称为“分布式插槽”的原因。
以下代码省略了很多，有兴趣的读者可以深入源码进行阅读。
template <typename TypeKey, typename TypeFP>void create_embedding<TypeKey, TypeFP>::operator()( std::map<std::string, SparseInput<TypeKey>>& sparse_input_map, std::vector<TensorEntry>* train_tensor_entries_list, std::vector<TensorEntry>* evaluate_tensor_entries_list, std::vector<std::shared_ptr<IEmbedding>>& embeddings, Embedding_t embedding_type, const nlohmann::json& config, const std::shared_ptr<ResourceManager>& resource_manager, size_t batch_size, size_t batch_size_eval, std::shared_ptr<ExchangeWgrad>& exchange_wgrad, bool use_mixed_precision, float scaler, const nlohmann::json& j_layers, bool use_cuda_graph, bool grouped_all_reduce) { #ifdef ENABLE_MPI int num_procs = 1, pid = 0; // 建立 MPI相关 MPI_Comm_rank(MPI_COMM_WORLD, &pid); MPI_Comm_size(MPI_COMM_WORLD, &num_procs);#endif // 从配置文件之中读取 auto j_optimizer = get_json(config, "optimizer"); auto embedding_name = get_value_from_json<std::string>(j_layers, "type"); auto bottom_name = get_value_from_json<std::string>(j_layers, "bottom"); auto top_name = get_value_from_json<std::string>(j_layers, "top"); auto j_hparam = get_json(j_layers, "sparse_embedding_hparam"); size_t workspace_size_per_gpu_in_mb = get_value_from_json_soft<size_t>(j_hparam, "workspace_size_per_gpu_in_mb", 0); auto embedding_vec_size = get_value_from_json<size_t>(j_hparam, "embedding_vec_size"); size_t max_vocabulary_size_per_gpu = (workspace_size_per_gpu_in_mb * 1024 * 1024) / (sizeof(float) * embedding_vec_size); auto combiner_str = get_value_from_json<std::string>(j_hparam, "combiner"); int combiner; // 设定combiner方法 if (combiner_str == "sum") { combiner = 0; } else if (combiner_str == "mean") { combiner = 1; } else { CK_THROW_(Error_t::WrongInput, "No such combiner type: " + combiner_str); } // 设定slot配置 std::vector<size_t> slot_size_array; if (has_key_(j_hparam, "slot_size_array")) { auto slots = get_json(j_hparam, "slot_size_array"); assert(slots.is_array()); for (auto slot : slots) { slot_size_array.emplace_back(slot.get<size_t>()); } } SparseInput<TypeKey> sparse_input; // 设定优化器配置 OptParams embedding_opt_params; if (has_key_(j_layers, "optimizer")) { embedding_opt_params = get_optimizer_param(get_json(j_layers, "optimizer")); } else { embedding_opt_params = get_optimizer_param(j_optimizer); } embedding_opt_params.scaler = scaler; // 建立不同的hash switch (embedding_type) { case Embedding_t::DistributedSlotSparseEmbeddingHash: { const SparseEmbeddingHashParams embedding_params = {batch_size, batch_size_eval, max_vocabulary_size_per_gpu, {}, embedding_vec_size, sparse_input.max_feature_num_per_sample, sparse_input.slot_num, combiner, // combiner: 0-sum, 1-mean embedding_opt_params}; embeddings.emplace_back(new DistributedSlotSparseEmbeddingHash<TypeKey, TypeFP>( sparse_input.train_sparse_tensors, sparse_input.evaluate_sparse_tensors, embedding_params, resource_manager)); break; } case Embedding_t::LocalizedSlotSparseEmbeddingHash: { const SparseEmbeddingHashParams embedding_params = {batch_size, batch_size_eval, max_vocabulary_size_per_gpu, slot_size_array, embedding_vec_size, sparse_input.max_feature_num_per_sample, sparse_input.slot_num, combiner, // combiner: 0-sum, 1-mean embedding_opt_params}; embeddings.emplace_back(new LocalizedSlotSparseEmbeddingHash<TypeKey, TypeFP>( sparse_input.train_sparse_tensors, sparse_input.evaluate_sparse_tensors, embedding_params, resource_manager)); break; } case Embedding_t::LocalizedSlotSparseEmbeddingOneHot: { const SparseEmbeddingHashParams embedding_params = {...}; embeddings.emplace_back(new LocalizedSlotSparseEmbeddingOneHot<TypeKey, TypeFP>( sparse_input.train_sparse_tensors, sparse_input.evaluate_sparse_tensors, embedding_params, resource_manager)); break; } case Embedding_t::HybridSparseEmbedding: { const HybridSparseEmbeddingParams<TypeFP> embedding_params = {...}; embeddings.emplace_back(new HybridSparseEmbedding<TypeKey, TypeFP>( sparse_input.train_sparse_tensors, sparse_input.evaluate_sparse_tensors, embedding_params, embed_wgrad_buff, get_gpu_learning_rate_schedulers(config, resource_manager), graph_mode, resource_manager)); break; } } // switch }接下来是建立网络环节，这部分过后，hugeCTR系统就正式建立起来，可以进行训练了，大体逻辑是：
Network* Network::create_network(const nlohmann::json& j_array, const nlohmann::json& j_optimizer, std::vector<TensorEntry>& train_tensor_entries, std::vector<TensorEntry>& evaluate_tensor_entries, int num_networks_in_global, std::shared_ptr<ExchangeWgrad>& exchange_wgrad, const std::shared_ptr<CPUResource>& cpu_resource, const std::shared_ptr<GPUResource>& gpu_resource, bool use_mixed_precision, bool enable_tf32_compute, float scaler, bool use_algorithm_search, bool use_cuda_graph, bool inference_flag, bool grouped_all_reduce) { Network* network = new Network(cpu_resource, gpu_resource, use_mixed_precision, use_cuda_graph); auto& train_layers = network->train_layers_; auto* bottom_layers = &network->bottom_layers_; auto* top_layers = &network->top_layers_; auto& evaluate_layers = network->evaluate_layers_; auto& train_loss_tensor = network->train_loss_tensor_; auto& evaluate_loss_tensor = network->evaluate_loss_tensor_; auto& train_loss = network->train_loss_; auto& evaluate_loss = network->evaluate_loss_; auto& enable_cuda_graph = network->enable_cuda_graph_; auto& raw_metrics = network->raw_metrics_; // 会进行GPU内存分配，这里大量使用了 create_block，其中就是 BufferBlockImpl std::shared_ptr<GeneralBuffer2<CudaAllocator>> blobs_buff = GeneralBuffer2<CudaAllocator>::create(); std::shared_ptr<BufferBlock2<float>> train_weight_buff = blobs_buff->create_block<float>(); std::shared_ptr<BufferBlock2<__half>> train_weight_buff_half = blobs_buff->create_block<__half>(); std::shared_ptr<BufferBlock2<float>> wgrad_buff = nullptr; std::shared_ptr<BufferBlock2<__half>> wgrad_buff_half = nullptr; if (!inference_flag) { if (use_mixed_precision) { auto id = gpu_resource->get_local_id(); wgrad_buff_half = (grouped_all_reduce) ? std::dynamic_pointer_cast<GroupedExchangeWgrad<__half>>(exchange_wgrad) ->get_network_wgrad_buffs()[id] : std::dynamic_pointer_cast<NetworkExchangeWgrad<__half>>(exchange_wgrad) ->get_network_wgrad_buffs()[id]; wgrad_buff = blobs_buff->create_block<float>(); // placeholder } else { auto id = gpu_resource->get_local_id(); wgrad_buff = (grouped_all_reduce) ? std::dynamic_pointer_cast<GroupedExchangeWgrad<float>>(exchange_wgrad) ->get_network_wgrad_buffs()[id] : std::dynamic_pointer_cast<NetworkExchangeWgrad<float>>(exchange_wgrad) ->get_network_wgrad_buffs()[id]; wgrad_buff_half = blobs_buff->create_block<__half>(); // placeholder } } else { wgrad_buff = blobs_buff->create_block<float>(); wgrad_buff_half = blobs_buff->create_block<__half>(); } std::shared_ptr<BufferBlock2<float>> evaluate_weight_buff = blobs_buff->create_block<float>(); std::shared_ptr<BufferBlock2<__half>> evaluate_weight_buff_half = blobs_buff->create_block<__half>(); std::shared_ptr<BufferBlock2<float>> wgrad_buff_placeholder = blobs_buff->create_block<float>(); std::shared_ptr<BufferBlock2<__half>> wgrad_buff_half_placeholder = blobs_buff->create_block<__half>(); std::shared_ptr<BufferBlock2<float>> opt_buff = blobs_buff->create_block<float>(); std::shared_ptr<BufferBlock2<__half>> opt_buff_half = blobs_buff->create_block<__half>(); // 建立训练网络层 if (!inference_flag) { // create train layers create_layers(j_array, train_tensor_entries, blobs_buff, train_weight_buff, train_weight_buff_half, wgrad_buff, wgrad_buff_half, train_loss_tensor, gpu_resource, use_mixed_precision, enable_tf32_compute, num_networks_in_global, scaler, enable_cuda_graph, inference_flag, train_layers, train_loss, nullptr, top_layers, bottom_layers); } // 建立评估网络层 // create evaluate layers create_layers(j_array, evaluate_tensor_entries, blobs_buff, evaluate_weight_buff, evaluate_weight_buff_half, wgrad_buff_placeholder, wgrad_buff_half_placeholder, evaluate_loss_tensor, gpu_resource, use_mixed_precision, enable_tf32_compute, num_networks_in_global, scaler, enable_cuda_graph, inference_flag, evaluate_layers, evaluate_loss, &raw_metrics); // 建立优化器 // create optimizer if (!inference_flag) { if (use_mixed_precision) { auto opt_param = get_optimizer_param(j_optimizer); network->optimizer_ = std::move(Optimizer::Create(opt_param, train_weight_buff->as_tensor(), wgrad_buff_half->as_tensor(), scaler, opt_buff_half, gpu_resource)); } else { auto opt_param = get_optimizer_param(j_optimizer); network->optimizer_ = std::move(Optimizer::Create(opt_param, train_weight_buff->as_tensor(), wgrad_buff->as_tensor(), scaler, opt_buff, gpu_resource)); } } else { try { TensorEntry pred_tensor_entry = evaluate_tensor_entries.back(); if (use_mixed_precision) { network->pred_tensor_half_ = Tensor2<__half>::stretch_from(pred_tensor_entry.bag); } else { network->pred_tensor_ = Tensor2<float>::stretch_from(pred_tensor_entry.bag); } } catch (const std::runtime_error& rt_err) { std::cerr << rt_err.what() << std::endl; throw; } } // 初始化网络其他信息 network->train_weight_tensor_ = train_weight_buff->as_tensor(); network->train_weight_tensor_half_ = train_weight_buff_half->as_tensor(); network->wgrad_tensor_ = wgrad_buff->as_tensor(); network->wgrad_tensor_half_ = wgrad_buff_half->as_tensor(); network->evaluate_weight_tensor_ = evaluate_weight_buff->as_tensor(); network->evaluate_weight_tensor_half_ = evaluate_weight_buff_half->as_tensor(); network->opt_tensor_ = opt_buff->as_tensor(); network->opt_tensor_half_ = opt_buff_half->as_tensor(); CudaDeviceContext context(gpu_resource->get_device_id()); blobs_buff->allocate(); return network;}create_layers 有两个版本，分别是HugeCTR/src/parsers/create_network.cpp 和 HugeCTR/src/cpu/create_network_cpu.cpp，我们使用 create_network.cpp 的代码来看看。
其实就是遍历从配置读取的json数组，然后建立每一层，因为层类型太多，所以我们只给出了两个例子。
void create_layers(const nlohmann::json& j_array, std::vector<TensorEntry>& tensor_entries, const std::shared_ptr<GeneralBuffer2<CudaAllocator>>& blobs_buff, const std::shared_ptr<BufferBlock2<float>>& weight_buff, const std::shared_ptr<BufferBlock2<__half>>& weight_buff_half, const std::shared_ptr<BufferBlock2<float>>& wgrad_buff, const std::shared_ptr<BufferBlock2<__half>>& wgrad_buff_half, Tensor2<float>& loss_tensor, const std::shared_ptr<GPUResource>& gpu_resource, bool use_mixed_precision, bool enable_tf32_compute, int num_networks_in_global, float scaler, bool& enable_cuda_graph, bool inference_flag, std::vector<std::unique_ptr<Layer>>& layers, std::unique_ptr<ILoss>& loss, metrics::RawMetricMap* raw_metrics, std::vector<Layer*>* top_layers = nullptr, std::vector<Layer*>* bottom_layers = nullptr) { for (unsigned int i = 1; i < j_array.size(); i++) { // 遍历json数组 const nlohmann::json& j = j_array[i]; const auto layer_type_name = get_value_from_json<std::string>(j, "type"); Layer_t layer_type; std::vector<TensorEntry> output_tensor_entries; // 这里获得本层的输入和输出 auto input_output_info = get_input_tensor_and_output_name(j, tensor_entries); switch (layer_type) { // 建立对应的每一层 case Layer_t::ReduceMean: { int axis = get_json(j, "axis").get<int>(); // 本层输入 Tensor2<float> in_tensor = Tensor2<float>::stretch_from(input_output_info.inputs[0]); Tensor2<float> out_tensor; emplaceback_layer( new ReduceMeanLayer<float>(in_tensor, out_tensor, blobs_buff, axis, gpu_resource)); // 本层输出 output_tensor_entries.push_back({input_output_info.output_names[0], out_tensor.shrink()}); break; } case Layer_t::Softmax: { // 本层输入 Tensor2<float> in_tensor = Tensor2<float>::stretch_from(input_output_info.inputs[0]); Tensor2<float> out_tensor; blobs_buff->reserve(in_tensor.get_dimensions(), &out_tensor); // 本层输出 output_tensor_entries.push_back({input_output_info.output_names[0], out_tensor.shrink()}); emplaceback_layer(new SoftmaxLayer<float>(in_tensor, out_tensor, blobs_buff, gpu_resource)); break; } } // end of switch } // for layers}HugeCTR 属于一个具体而微的深度学习系统，它实现的具体层类型如下：
enum class Layer_t { BatchNorm, BinaryCrossEntropyLoss, Reshape, Concat, CrossEntropyLoss, Dropout, ELU, InnerProduct, FusedInnerProduct, Interaction, MultiCrossEntropyLoss, ReLU, ReLUHalf, GRU, MatrixMultiply, Scale, FusedReshapeConcat, FusedReshapeConcatGeneral, Softmax, PReLU_Dice, ReduceMean, Sub, Gather, Sigmoid, Slice, WeightMultiply, FmOrder2, Add, ReduceSum, MultiCross, Cast, DotProduct, ElementwiseMultiply};我们使用 SigmoidLayer 作为例子，大家来看看。
/** * Sigmoid activation function as a derived class of Layer */template <typename T>class SigmoidLayer : public Layer { /* * stores the references to the input tensors of this layer. */ Tensors2<T> in_tensors_; /* * stores the references to the output tensors of this layer. */ Tensors2<T> out_tensors_; public: /** * Ctor of SigmoidLayer. * @param in_tensor the input tensor * @param out_tensor the output tensor which has the same dim with in_tensor * @param device_id the id of GPU where this layer belongs */ SigmoidLayer(const Tensor2<T>& in_tensor, const Tensor2<T>& out_tensor, const std::shared_ptr<GPUResource>& gpu_resource); /** * A method of implementing the forward pass of Sigmoid * @param stream CUDA stream where the foward propagation is executed */ void fprop(bool is_train) override; /** * A method of implementing the backward pass of Sigmoid * @param stream CUDA stream where the backward propagation is executed */ void bprop() override;};其前向传播如下：
template <typename T>void SigmoidLayer<T>::fprop(bool is_train) { CudaDeviceContext context(get_device_id()); int len = in_tensors_[0].get_num_elements(); auto fop = [] __device__(T in) { return T(1) / (T(1) + exponential(-in)); }; MLCommon::LinAlg::unaryOp(out_tensors_[0].get_ptr(), in_tensors_[0].get_ptr(), len, fop, get_gpu().get_stream());#ifndef NDEBUG cudaDeviceSynchronize(); CK_CUDA_THROW_(cudaGetLastError());#endif}其后向传播如下：
template <typename T>void SigmoidLayer<T>::bprop() { CudaDeviceContext context(get_device_id()); int len = in_tensors_[0].get_num_elements(); auto bop = [] __device__(T d_out, T d_in) { T y = T(1) / (T(1) + exponential(-d_in)); return d_out * y * (T(1) - y); }; MLCommon::LinAlg::binaryOp(in_tensors_[0].get_ptr(), out_tensors_[0].get_ptr(), in_tensors_[0].get_ptr(), len, bop, get_gpu().get_stream());#ifndef NDEBUG cudaDeviceSynchronize(); CK_CUDA_THROW_(cudaGetLastError());#endif}至此，HugeCTR 已经初始化完毕，接下来可以开始训练了，我们摘录官方HugeCTR_Webinar 之中的图来给大家做一个梳理，这里CSR是嵌入层依赖的数据格式，我们下文会分析。
我们尚且有一个疑问，那就是层与层之间如何串联起来？
在create_layers之中有如下代码：
// 获取本层的输入和输出auto input_output_info = get_input_tensor_and_output_name(j, tensor_entries);get_input_tensor_and_output_name 的代码如下，可以看到，每一层都会记录自己的输入和输出，结合内部解析模块，这些层就建立其了逻辑关系。
static InputOutputInfo get_input_tensor_and_output_name( const nlohmann::json& json, const std::vector<TensorEntry>& tensor_entries) { auto bottom = get_json(json, "bottom"); auto top = get_json(json, "top"); // 从jason获取输入，输出名字 std::vector<std::string> bottom_names = get_layer_names(bottom); std::vector<std::string> top_names = get_layer_names(top); std::vector<TensorBag2> bottom_bags; // 把输出组成一个向量列表 for (auto& bottom_name : bottom_names) { for (auto& top_name : top_names) { if (bottom_name == top_name) { CK_THROW_(Error_t::WrongInput, "bottom and top include a same layer name"); } } TensorBag2 bag; if (!get_tensor_from_entries(tensor_entries, bottom_name, &bag)) { CK_THROW_(Error_t::WrongInput, "No such bottom: " + bottom_name); } bottom_bags.push_back(bag); } return {bottom_bags, top_names}; // 返回}最终，建立流水线逻辑关系如下：
具体训练代码逻辑如下：
bool Session::train() { try { // 确保 train_data_reader_ 已经启动 if (train_data_reader_->is_started() == false) { CK_THROW_(Error_t::IllegalCall, "Start the data reader first before calling Session::train()"); }#ifndef DATA_READING_TEST // 需要 reader 先读取一个 batchsize 的数据。 long long current_batchsize = train_data_reader_->read_a_batch_to_device_delay_release(); if (!current_batchsize) { return false; // 读不到就退出，没有数据了 } #pragma omp parallel num_threads(networks_.size()) //其后语句将被networks_.size()个线程并行执行 { size_t id = omp_get_thread_num(); CudaCPUDeviceContext ctx(resource_manager_->get_local_gpu(id)->get_device_id()); cudaStreamSynchronize(resource_manager_->get_local_gpu(id)->get_stream()); } // reader 可以开始解析数据 train_data_reader_->ready_to_collect();#ifdef ENABLE_PROFILING global_profiler.iter_check();#endif // If true we're gonna use overlaping, if false we use default if (solver_config_.use_overlapped_pipeline) { train_overlapped(); } else { for (const auto& one_embedding : embeddings_) { one_embedding->forward(true); // 嵌入层进行前向传播，即从参数服务器读取embedding，进行处理 } // Network forward / backward if (networks_.size() > 1) { // 因为之前是把模型分别拷贝到GPU之上，所以size大于1，就说明多卡 // 单机多卡或多机多卡 // execute dense forward and backward with multi-cpu threads #pragma omp parallel num_threads(networks_.size()) { // dense网络的前向反向 size_t id = omp_get_thread_num(); long long current_batchsize_per_device = train_data_reader_->get_current_batchsize_per_device(id); networks_[id]->train(current_batchsize_per_device); // 前向操作 const auto& local_gpu = resource_manager_->get_local_gpu(id); local_gpu->set_compute_event_sync(local_gpu->get_stream()); local_gpu->wait_on_compute_event(local_gpu->get_comp_overlap_stream()); } } else if (resource_manager_->get_global_gpu_count() > 1) { // 多机单卡 long long current_batchsize_per_device = train_data_reader_->get_current_batchsize_per_device(0); networks_[0]->train(current_batchsize_per_device); // 前向操作 const auto& local_gpu = resource_manager_->get_local_gpu(0); local_gpu->set_compute_event_sync(local_gpu->get_stream()); local_gpu->wait_on_compute_event(local_gpu->get_comp_overlap_stream()); } else { // 单机单卡 long long current_batchsize_per_device = train_data_reader_->get_current_batchsize_per_device(0); networks_[0]->train(current_batchsize_per_device); // 前向操作 const auto& local_gpu = resource_manager_->get_local_gpu(0); local_gpu->set_compute_event_sync(local_gpu->get_stream()); local_gpu->wait_on_compute_event(local_gpu->get_comp_overlap_stream()); networks_[0]->update_params(); } // Embedding backward for (const auto& one_embedding : embeddings_) { one_embedding->backward(); // 嵌入层反向操作 } // Exchange wgrad and update params if (networks_.size() > 1) { #pragma omp parallel num_threads(networks_.size()) { size_t id = omp_get_thread_num(); exchange_wgrad(id); // 多卡之间交换dense参数的梯度 networks_[id]->update_params(); } } else if (resource_manager_->get_global_gpu_count() > 1) { exchange_wgrad(0); networks_[0]->update_params(); } for (const auto& one_embedding : embeddings_) { one_embedding->update_params(); // 嵌入层更新sparse参数 } // Join streams 各个流进行同步 if (networks_.size() > 1) { #pragma omp parallel num_threads(networks_.size()) { size_t id = omp_get_thread_num(); const auto& local_gpu = resource_manager_->get_local_gpu(id); local_gpu->set_compute2_event_sync(local_gpu->get_comp_overlap_stream()); local_gpu->wait_on_compute2_event(local_gpu->get_stream()); } } else { const auto& local_gpu = resource_manager_->get_local_gpu(0); local_gpu->set_compute2_event_sync(local_gpu->get_comp_overlap_stream()); local_gpu->wait_on_compute2_event(local_gpu->get_stream()); } return true; }#else data_reader_->read_a_batch_to_device();#endif } catch (const internal_runtime_error& err) { std::cerr << err.what() << std::endl; throw err; } catch (const std::exception& err) { std::cerr << err.what() << std::endl; throw err; } return true;}训练流程如下：
至此，我们大体知道了 HugeCTR如何初始化和训练，下一篇我们介绍如何读取数据。
https://developer.nvidia.com/blog/introducing-merlin-hugectr-training-framework-dedicated-to-recommender-systems/
https://developer.nvidia.com/blog/announcing-nvidia-merlin-application-framework-for-deep-recommender-systems/
https://developer.nvidia.com/blog/accelerating-recommender-systems-training-with-nvidia-merlin-open-beta/
HugeCTR源码阅读
embedding层如何反向传播
https://web.eecs.umich.edu/~justincj/teaching/eecs442/notes/linear-backprop.html
HugeCTR_Webinar