HugeCTR简介

基于参数服务器架构的大规模稀疏训练，可以说好几年没有新的变化和进步了，直到百度的aibox论文出现，以及后来nvidia开发的hugectr开源出来。总算是看到参数服务器架构又朝前走了一步。可以预见的是，这样的异构训练架构，相比之前的纯CPU的方式，一定会随着更多的高性能硬件、新的训练优化的出现，有更进一步的改进空间。希望大家看到后有什么问题可以留言讨论，一起进步。

相关的代码库链接：https://github.com/NVIDIA/HugeCTR/

hugectr是nvidia开发的GPU分布式训练框架，它主要针对的是推荐ctr场景，支持大规模稀疏参数的分布式训练与评估。

hugectr是一个基于参数服务器架构的训练框架，它的主要亮点在于，它有基于GPU显存的参数服务器（通俗一点说就是GPU显存里有个hashmap用来存参数），这样在做GPU训练的时候，参数可以直接从GPU拷贝或者利用GPU通信，大大加速了参数通信（pull和push），因为参数通信不再经过CPU了。

当然这也抛出了几个问题，不妨先思考一下：
（1）我们知道hashmap通常在大小接近一定阈值的时候，会做rehash的操作，而对参数服务器来说，随着不断训练，参数服务器存的参数也可能不断增多，尤其是训练刚开始的阶段，hashmap大小增长较快，如果频繁rehash或者创建新的key，那么就会存在频繁的GPU显存申请释放和拷贝，会很影响性能。
（2）GPU多了一步数据从CPU拷贝到GPU的过程，如果不做什么改进，计算的加速是否一定相比额外的数据拷贝要快呢？
（3）大规模稀疏参数，我们key的规模在十亿，value（即embedding向量）假如是32维float向量，假如我们使用的是adam优化器，那么算上优化器状态就是32 * 3 = 96维，总的参数规模大于300G，全放在显存里放得下吗？

这些问题在后面看的过程中会逐步解答，接下来让我们先看一下整体架构。

整体架构

训练流程：首先reader从dataset中读取batch_size（比如32）的原始数据，解析原始数据，得到输入的sparse key，dense向量，label等，并根据sparse key从参数服务器（下文简称ps，即parameter server）中拉取（pull）对应的embedding向量，然后输入到深度学习神经网络中做前向-反向计算，并把反向计算得到的参数梯度推送（push）到ps，由ps根据梯度更新参数。

训练的过程可以认为是数据并行+模型并行。数据并行主要体现在每个GPU卡的同时读取不同的数据做训练，而模型训练主要体现在sparse参数的存储是存储在多个节点（node）上，每个节点分配一部分参数。

我们知道，ctr场景中，sparse参数的规模通常很大，从千万到万亿级别的大小，dense参数（网络中的权重）通常很小，占用内存大小也就几MB到几十MB。因此对sparse参数的存取需要好好设计。hugectr中对sparse的存储方式有两种：local和distribute。

让我们先看看local模式：一个slot的参数，只会在一个gpu卡上，这样查完embedding之后，因为已经拿到了这个slot的所有embedding，可以做完pooling之后再做GPU多卡通信，可以降低通信量。（这里slot的意思是特征种类，也可以称作field）

举个例子，我们是单机训练有4张GPU卡，有8个slot：slot0到slot7。如果是local模式，那么就是GPU0存slot0和slot1，GPU1存slot2和slot3，GPU2存slot4和slot5，GPU3存slot6和slot7。

对于distribute模式，每个GPU上上都会存所有slot的一部分参数，至于如何将一个参数分配到哪个GPU上，可以通过哈希的方法。

下图是从多线程数据读取、数据从CPU拷贝到GPU、训练的过程，图里的worker，其实指的就是reader，多个reader同时解析dataset的数据，然后由collector模块将数据拷贝到GPU。图里的worker、collector、training这三个是通过流水线串起来，各个部分相互独立，同时在不同的线程中运行。

下图就是流水线的具体示例，每种颜色代表一级流水线，共三级流水线。当第一级解析完batch0后，扔给第二级用来拷贝给GPU，这时候第一级继续解析batch1。当batch0在训练的时候，同时在做的是batch1拷贝到GPU。

需要注意的是，上图各级流水线的时间默认是相等的，但是实际情况一般不会这么巧，那么一般是需要灵活调整各级流水线的线程数的，让各级流水线的速度匹配起来。举个例子比如，readfile开10个线程，copy开5个线程，训练开8个线程。另外，上面实际是纵向在看流水线，如果横向来看，是各个batch互相独立的在做readfile-copy-train。

有没有发现，这里的流水线解答了开头的第2个问题。

代码阅读

我看代码喜欢自顶向下的看，这样不仅能开始就能掌握运行的整个流程，后面看细节的时候也可以有针对性得看。另外，我没有弄太多子标题，顺序往下看就好。

我们首先看readme中的例子，这是一个调用了python api的例子。hugectr与常见的深度学习框架一样，分为python端和c++端，python封装user api，c++实现底层训练逻辑。

   # train.py
   import sys
   import hugectr
   from mpi4py import MPI
   def 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()
         print("[HUGECTR][INFO] iter: {}; loss: {}".format(i, loss))
   if __name__ == "__main__":
     json_config_file = sys.argv[1]
     train(json_config_file)

另外，观察这里的api会发现，长得跟tensorflow单机的api是不是很像，确实分布式框架的一个目标就是用起来像写单机程序一样顺手，也就是所谓的“易用性”。
这里的solver_config就是把各种训练配置传入hugectr，session就是封装了分布式训练逻辑，start_data_reading就是字面意思，启动上文的readfile的异步线程，也就是第一级流水线。接着就是train，然后打印oss。

前面说过，自顶向下看代码的好处是可以对细节有针对性，那咱们先看重点，也就是sess.train。

python端与c++端的连接，可以使用pybind库。连接的“桥梁”定义在pybind/session_wrapper.hpp这个文件里：

void SessionPybind(pybind11::module &m) {
  pybind11::class_<HugeCTR::Session, std::shared_ptr<HugeCTR::Session>>(m, "Session")
      .def(pybind11::init<const SolverParser &, const std::string &, bool, const std::string>(),
           pybind11::arg("solver_config"), pybind11::arg("config_file"),
           pybind11::arg("use_model_oversubscriber") = false,
           pybind11::arg("temp_embedding_dir") = std::string())
      .def("train", &HugeCTR::Session::train)
      .def("eval", &HugeCTR::Session::eval)
      .def("start_data_reading", &HugeCTR::Session::start_data_reading)
      ....

在python调用sess.train，对应了c++的HugeCTR::Session::train，让我们来看一下这个函数，我加了一些注释：

bool Session::train() {
    // 判断reader是否启动，未启动就开始训练则报错
    if (train_data_reader_->is_started() == false) { 
      CK_THROW_(xxxx);
    }
    // 等待reader读取至少一个batchsize的数据
    long long current_batchsize = 0;
    while ((current_batchsize = train_data_reader_->read_a_batch_to_device_delay_release()) &&
           (current_batchsize < train_data_reader_->get_full_batchsize())) {
      // 告诉reader可以开始解析数据了，通过设置flag：READY_TO_WRITE
      train_data_reader_->ready_to_collect();
    }
    // 读不到数据了，即没有数据可以训练了，直接返回
    if (!current_batchsize) {
      return false;
    }
    // reader解析完一个batch的数据后，flag会被设置为READY_TO_READ
    // 上面通过read_a_batch_to_device_delay_release把数据已经从reader中取出来，
    // 并且正在异步的拷贝到GPU，
    // 调用ready_to_collect，首先sync上面的异步拷贝，然后让reader继续解析下一个batch
    train_data_reader_->ready_to_collect();
    // 从ps查embedding，做sum或者avg
    for (auto& one_embedding : embeddings_) {
      one_embedding->forward(true);
    }
    // 这里的逻辑看着有点乱，也就是多卡数据并行训练，
    // 一个网络有gpu卡数个副本，也就是networks大小大于1的原因。 
    if (networks_.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);
        // 多卡之间交换dense参数的梯度
        networks_[id]->exchange_wgrad();
        // 更新dense参数
        networks_[id]->update_params();
      }
    } 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);
      networks_[0]->exchange_wgrad();
      networks_[0]->update_params();
    } else {
      // 单机单卡
      long long current_batchsize_per_device =
          train_data_reader_->get_current_batchsize_per_device(0);
      networks_[0]->train(current_batchsize_per_device);
      networks_[0]->update_params();
    }
    // embedding的反向
    for (auto& one_embedding : embeddings_) {
      one_embedding->backward();
      // 更新sparse参数
      one_embedding->update_params();
    }
    return true;
}

看到这里，基本上训练中的大体流程是清楚了。接下来，我们继续深入往下看reader、embedding、参数存储和通信等部分。首先有必要看一下初始化。

HugeCTR::Session的初始化代码如下：

parser.create_pipeline(
    train_data_reader_, evaluate_data_reader_,
    embeddings_, networks_, resource_manager_);
#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()));
}
init_or_load_params_for_dense_(solver_config.model_file);
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);

也就是分为以下步骤：
（1）创建三级流水线，即create_pipeline
（2）初始化network
（3）初始化参数，以及对应的优化器状态

这里面比较重要的部分是创建三级流水线，我们看一下create_pipeline的实现（函数传参先忽略掉）：

// create reader
create_datareader<TypeKey>()(...);
// create embedding
for (unsigned int i = 1; i < j_layers_array.size(); i++) {
  // 网路配置的每层是从bottom到top的，因此只要遇到非embedding的layer，
  // 后面的layer就不用检查了
  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)) {
    break;
  }
  create_embedding<TypeKey, float>()(...);
}
// create network，每张GPU卡创建一个network副本
for (size_t i = 0; i < resource_manager->get_local_gpu_count(); i++) {
  network.emplace_back(Network::create_network(...));
}

可以看到create_pipeline主要包含了三步：create_datareader、create_embedding、create_network

我们接下来先看create_datareader里面做了什么：创建了一个train_data_reader和一个evaluate_data_reader，也就是一个用于训练，一个用于评估。然后还各自创建了WorkerGroup。

DataReader<TypeKey>* data_reader_tk = new DataReader<TypeKey>(...);
train_data_reader.reset(data_reader_tk);
DataReader<TypeKey>* data_reader_eval_tk = new DataReader<TypeKey>(...);
evaluate_data_reader.reset(data_reader_eval_tk);

train_data_reader->create_drwg_norm(source_data, check_type, repeat_dataset_);
evaluate_data_reader->create_drwg_norm(eval_source, check_type, repeat_dataset_);

  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>(
        csr_heap_, file_name, repeat_, check_type, params_, start_reading_from_beginning));
    file_name_ = file_name;
  }

// DataReaderWorkerGroupNorm的构造函数主要是如下功能，创建DataReaderWorker
for (int i = 0; i < NumThreads; i++) {
  std::shared_ptr<IDataReaderWorker> data_reader(new DataReaderWorker<TypeKey>(
          i, NumThreads, csr_heap, file_list, max_feature_num_per_sample, repeat, check_type, params));
  data_readers_.push_back(data_reader);
}
// 然后创建了多个线程。每个线程对应一个reader，执行如下逻辑
while (*p_loop_flag) {
  data_reader->read_a_batch();
}

好了，看到这里，出现了两个类DataReader和DataReaderWorkerGroupNorm，这两个类有必要看一下细节，从而弄清楚数据读取。

先看DataReader里的重要的函数和变量：

template <typename TypeKey>
class DataReader : public IDataReader {
  std::shared_ptr<HeapEx<CSRChunk<TypeKey>>> csr_heap_;
  std::shared_ptr<DataCollector<TypeKey>> data_collector_;
  std::shared_ptr<DataReaderWorkerGroup> worker_group_;
  //还有各种tensor：label_tensors_、dense_tensors_、row_offsets_tensors_、value_tensors_等
  DataReader(...) {
    // 初始化heap，这个类后面介绍
    csr_heap_.reset(new HeapEx<CSRChunk<TypeKey>>(...));
    // 为每个GPU初始化一个buffer
    std::vector<std::shared_ptr<GeneralBuffer2<CudaAllocator>>> buffs;
    for (size_t i = 0; i < local_gpu_count; i++) {
      buffs.push_back(GeneralBuffer2<CudaAllocator>::create());
    }
    // create label and dense tensor
    size_t batch_size_per_device = batchsize_ / total_gpu_count;
    for (size_t i = 0; i < local_gpu_count; i++) {
      // Tensor2并不持有内存或者显存，是在buffs里
      {
        Tensor2<float> tensor;
        buffs[i]->reserve({batch_size_per_device, label_dim_}, &tensor);
        label_tensors_.push_back(tensor);
      }
      {
        Tensor2<float> tensor;
        buffs[i]->reserve({batch_size_per_device, dense_dim_}, &tensor);
        dense_tensors_.push_back(tensor.shrink());
      }
    }
    ...
    // 这里又出现一个DataCollector类
    data_collector_.reset(new DataCollector<TypeKey>(...);
    data_collector_->start();
    // buffs在每个GPU上分配显存
    for (size_t i = 0; i < local_gpu_count; i++) {
      CudaDeviceContext context(resource_manager_->get_local_gpu(i)->get_device_id());
      buffs[i]->allocate();
    }
  }

详细介绍一下上面出现的类：
（1）CSR：一种用来压缩稀疏矩阵的存储格式，例子如下

假如有这样一组数据
 *   4,5,1,2
 *   3,5,1
 *   3,2
用CSR可以表示为
  row offset: 0,4,7,9
  value: 4,5,1,2,3,5,1,3,2

这里的CSR由于是用来存储slot里的sparse key，其实少了column index，因为一个slot里的sparse key直接顺序存就好了，CSR可以参考这篇文章。总是令人想起百度paddle的lodtensor。

（2）HeapEx：为每个数据解析的线程维护了三个队列，ready queue 、 wait queue 和 credit queue，队列中的元素就是CSR存储。当credit queue非空，意味着有空闲的CSR可以用来存储解析好的数据。

• 解析数据（多线程）：当前线程只会操作自己的queue。从credit queue取出一个空闲的CSR后，会同时把它塞到wait queue里，数据解析完存到该CSR后，会把它从wait queue移除，塞到ready queue里（之所以多了一步放在wait queue里的步骤，主要是为了记录该CSR，从而塞到ready queue里）。
• 读数据（单线程）：从上述多个队列里取数据。HeapEx里有一个变量count，记录了上次是从哪个index的数据解析线程的ready queue中取得的数据，下次取数据还是从这个index的ready queue开始遍历，也就是“”轮询”，归还数据还给这个index。归还是归还到credit queue里，并将其从ready queue移除。当归还成功后count++。（为啥不是取数据的时候，就把数据从ready queue移除呢？因为取数据是单线程的，所以这样做是可以的）

（3）DataReaderWorker：就是上面的解析数据线程。
（4）DataCollector：就是上面的读数据线程。把数据从CSR拷贝到GPU。它会启动一个线程不断的执行如下函数：

template <typename TypeKey>
void DataCollector<TypeKey>::collect_() {
  std::unique_lock<std::mutex> lock(stat_mtx_);
  CSRChunk<TypeKey>* chunk_tmp = csr_heap_->checkout_data_chunk();
  while (stat_ != READY_TO_WRITE && stat_ != STOP) {
    usleep(2);
  }
  ...
  cudaMemcpyAsync 做拷贝
  ...
  cudaStreamSynchronize
  ...
  csr_heap_->return_free_chunk();
  stat_ = READY_TO_READ;

这里checkout_data_chunk从一个队列里取出一个CSR，拷贝完成后再调用return_free_chunk归还。

为了清晰一点表示流程，我画了个简易的图，绿色的线是数据生产消费，红色是数据归还。data reader worker获取到空闲的CSR就解析数据填进去，并放到ready queue中，而data collector发现有可用的CSR，就拷贝到GPU中，然后归还CSR。可以说，data reader worker和data collector是相互独立的，两者通过存放数据的队列联系起来。

还有一些细节值得注意：
（1）hugectr还实现了两外两种data reader worker，一种是DataReaderWorkerRaw，它读取的数据是通过mmap直接映射到内存。还有一种是DataReaderWorkerGroup，读取parquet格式的文件。
（2）data reader worker解析数据放到CSR的细节：

// 返回的是get_batchsize训练配置的batch size
for (i = 0; i < csr_chunk->get_batchsize(); i++) {
  // dense输入
  {
    // label_dense_buffers的大小是当前节点的gpu卡数
    // buffer_id就是这条样本落在哪张卡上
    int buffer_id = i / (csr_chunk->get_batchsize() / label_dense_buffers.size());
    // local_id是这条样本在当前卡上的偏移
    int local_id = i % (csr_chunk->get_batchsize() / label_dense_buffers.size());
    // 拷贝，从解析数据的buffer拷贝到CSR里
    float* ptr = label_dense_buffers[buffer_id].get_ptr();
    for (int j = 0; j < label_dense_dim; j++) {
        ptr[local_id * label_dense_dim + j] = label_dense[j];  // row major for label buffer
    }
  }
  // sparse输入
  for (auto& param : params_) {
    for (int k = 0; k < param.slot_num; k++) {
      // 省略了读取数据
      ...
      // 下面这段代码，对理解embedding的两种存储方式很有帮助
      if (param.type == DataReaderSparse_t::Distributed) {
        // 所有节点的每张卡上都会存所有slot
        for (int dev_id = 0; dev_id < csr_chunk->get_num_devices(); dev_id++) {
          csr_chunk->get_csr_buffer(param_id, dev_id).new_row();
        }
        // 这里就是判断slot的key应该存在哪张卡上
        for (int j = 0; j < nnz; j++) {
          int dev_id = feature_ids_[j] % csr_chunk->get_num_devices();
          dev_id = std::abs(dev_id);
          T local_id = feature_ids_[j];
          csr_chunk->get_csr_buffer(param_id, dev_id).push_back(local_id);
        }
      } else if (param.type == DataReaderSparse_t::Localized) {
        // 一个slot只会存在一张卡上
        int dev_id = k % csr_chunk->get_num_devices();
        csr_chunk->get_csr_buffer(param_id, dev_id).new_row();
        for (int j = 0; j < nnz; j++) {
          T local_id = feature_ids_[j];
          csr_chunk->get_csr_buffer(param_id, dev_id).push_back(local_id);
        }
      }      
    }
  }
}

（3）data collector拷贝数据的细节：

  // total_device_count 所有节点的gpu之和
  for (int ix = 0; ix < total_device_count; ix++) {
    int i =
        ((id_ == 0 && !reverse_) || (id_ == 1 && reverse_)) ? ix : (total_device_count - 1 - ix);
    int pid = resource_manager_->get_process_id_from_gpu_global_id(i);
    int label_copy_num = (label_dense_buffers[0]).get_num_elements();
    if (pid == resource_manager_->get_process_id()) {
      ...
      for (int j = 0; j < num_params; j++) {
        // 这里的i * num_params + j取的就是全局的偏移
        unsigned int nnz = csr_cpu_buffers[i * num_params + j]
                               .get_row_offset_tensor()
                               .get_ptr()[csr_cpu_buffers[i * num_params + j].get_num_rows()];
        // cudaMemcpyAsync 异步拷贝
    }

我们发现：所有的节点都会解析所有的数据！在拷贝的时候，才会只拷贝出属于本节点的数据。这种实现，对于大数据量来说，性能是很值得怀疑的，带宽可能不够，并且会解析大量无用的数据。

一方面数据拷贝的时候会只保留属于本gpu的sparse key，另一方面本gpu中也只会存储属于本节点的sparse key，那么也就是说sparse pull和push的时候就不需要做节点间的通信了。当然后续还是需要做节点间通信，把数据拼成完整的batch做前向反向计算，得到梯度后做节点间的梯度平均。这块会在create embedding处详细展开。

实际上面的data reader worker 与 data collector就是本文开头的图里的第一、二级流水线了。接下来我们再看一下创建第三级流水线中的create_embedding部分。

创建embedding初始化，我们应该最关注的是如何存参数。在此之前，我们看看网络配置，看看是咋组织embedding的， 这里就以deepfm为例，首先看一下输入层

      "dense": {
        "top": "dense",
        "dense_dim": 13
      },
      "sparse": [
        {
          "top": "data1",
          "type": "DistributedSlot",
          "max_feature_num_per_sample": 30,
          "slot_num": 26
	    }
      ]

dense输入是deepfm的fm输入。sparse是deepfm的deep输入，包含了26个slot。再看embedding层的定义，max_vocabulary_size_per_gpu是一个gpu卡上的最大sparse key的个数，embedding_vec_size就是embedding向量维度，combiner表示查完embedding做pooling是sum还是avg。

    {
      "name": "sparse_embedding1",
      "type": "DistributedSlotSparseEmbeddingHash",
      "bottom": "data1",
      "top": "sparse_embedding1",
      "sparse_embedding_hparam": {
        "max_vocabulary_size_per_gpu": 1447751,
        "embedding_vec_size": 11,
        "combiner": 0
      }
    }

hugectr有两种embedding：DistributedSlotSparseEmbeddingHash和 LocalizedSlotSparseEmbeddingHash。我们一个一个看。

大致扫一眼hashmap创建的代码，两者是一样的：

// 注意到hash table的value type是个size_t,这个是记录了embedding在存储中的偏移量
using NvHashTable = HashTable<TypeHashKey, size_t>;
// 这个是hashmap的定义，发现外面套了个vector，需要弄清楚vector每个元素是啥
std::vector<std::shared_ptr<NvHashTable>> hash_tables_;
// 原来vector大小是本地gpu数，也就是说每个gpu卡对应一个hash table   
hash_tables_.resize(Base::get_resource_manager().get_local_gpu_count());
// 事先就固定了hash table容纳元素的最大数量为max_vocabulary_size_per_gpu_
#pragma omp parallel num_threads(Base::get_resource_manager().get_local_gpu_count())
{
  size_t id = omp_get_thread_num();
  CudaDeviceContext context(Base::get_local_gpu(id).get_device_id());
  // construct HashTable object: used to store hash table <key, value_index>
  hash_tables_[id].reset(new NvHashTable(max_vocabulary_size_per_gpu_));
  Base::get_buffer(id)->allocate();
}

这里可以解答开头的第一个问题了，答案很简单粗暴，就是事先固定了哈希表的大小。
这样其实对于大规模稀疏来说，支持的参数规模不会很大，就像开头第三个问题算过的亿级别就得300G。因此hugectr支持的参数规模会因为这样的设计，而受到比较大的限制。
那是不是第三个问题无解了呢，其实不然，百度的aibox不存在的这样的限制，因为它采用了多级ps（ssd+mem+gpu），并且gpu ps的大小会随着当前训练的增量数据的key的多少而动态创建不同大小的hashmap。这里aibox暂时先不展开了，我后面会写一篇文章专门介绍它的论文。：）

继续看一下HashTable这个类

template <typename KeyType, typename ValType>
class HashTable {
  // 查找key，如果key不在table里就insert
  void get_insert(const KeyType* d_keys, ValType* d_vals, size_t len, cudaStream_t stream);
  // 查找key
  void get(const KeyType* d_keys, ValType* d_vals, size_t len, cudaStream_t stream) const;
  // 把table里的kv都dump出来
  void dump(KeyType* d_key, ValType* d_val, size_t* d_dump_counter, cudaStream_t stream) const;
  HashTableContainer<KeyType, ValType>* container_

我们再看一下HashTableContainer，它是继承了concurrent_unordered_map这个类

template <typename KeyType, typename ValType>
class HashTableContainer
    : public concurrent_unordered_map<KeyType, ValType, std::numeric_limits<KeyType>::max()> {
 public:
  HashTableContainer(size_t capacity)
      : concurrent_unordered_map<KeyType, ValType, std::numeric_limits<KeyType>::max()>(
            capacity, std::numeric_limits<ValType>::max()) {}
};

concurrent_unordered_map是固定大小的、在显存中的map，它支持并发insert，但是不支持并发insert和get。因为hugectr训练是同步的训练，pull的时候只会有get，push的时候只会有insert，并且不会同时做pull和push，因此这样的concurrent_unordered_map满足要求。
首先看一下它的get函数：

// __forceinline__ 强制指定为内联函数
// __host__ __device__ 这个函数会同时为主机端和设备端编译
__forceinline__ __host__ __device__ const_iterator find(const key_type& k) const {
    // 对key做哈希
    size_type key_hash = m_hf(k);
    // 映射到table的一个index
    size_type hash_tbl_idx = key_hash % m_hashtbl_size;
    value_type* begin_ptr = 0;
    size_type counter = 0;
    while (0 == begin_ptr) {
      value_type* tmp_ptr = m_hashtbl_values + hash_tbl_idx;
      const key_type tmp_val = tmp_ptr->first;
      // 找到了这个key
      if (m_equal(k, tmp_val)) {
        begin_ptr = tmp_ptr;
        break;
      }
      // 这个位置是空的，或者找完了这个table也没找到
      if (m_equal(unused_key, tmp_val) || counter > m_hashtbl_size) {
        begin_ptr = m_hashtbl_values + m_hashtbl_size;
        break;
      }
      hash_tbl_idx = (hash_tbl_idx + 1) % m_hashtbl_size;
      ++counter;
    }
    return const_iterator(m_hashtbl_values, m_hashtbl_values + m_hashtbl_size, begin_ptr);
  }

可以看出，get一个key的时候，如果insert了这个key，可能还是get不到，或者get到的是错误的值（insert正在修改这个值的时候，get了这个值），或者旧的值。

再看一下insert，它的主要过程如下

const key_type insert_key = k;
bool insert_success = false;
size_type counter = 0;
while (false == insert_success) {
  // 哈希表满了
  if (counter++ >= hashtbl_size) {
    return end();
  }
  key_type& existing_key = current_hash_bucket->first;
  volatile mapped_type& existing_value = current_hash_bucket->second;
  // existing_key == unused_key时，insert_key会被赋值给existing_key，因为这个位置是空的。
  // existing_key == insert_key时，这个位置已经有这个key了，
  // 如果这时候existing_value == m_unused_element，就说明其他线程正在insert且还没来得及修改existing_value
  const key_type old_key = atomicCAS(&existing_key, unused_key, insert_key);
  if (keys_equal(unused_key, old_key)) {
    existing_value = (mapped_type)(atomicAdd(value_counter, 1));
    break;
  } else if (keys_equal(insert_key, old_key)) {
    while (existing_value == m_unused_element) { }
    break;
  }
  // 这个位置被其他key占了，继续往后遍历
  current_index = (current_index + 1) % hashtbl_size;
  current_hash_bucket = &(hashtbl_values[current_index]);
}
return iterator(m_hashtbl_values, m_hashtbl_values + hashtbl_size, current

atomicCAS函数参考这篇文章。

两种embedding：
（1）DistributedSlotSparseEmbeddingHash

void forward(bool is_train) override {
  // Read data from input_buffers_ -> look up -> write to output_tensors
  CudaDeviceContext context;
  for (size_t i = 0; i < Base::get_resource_manager().get_local_gpu_count(); i++) {
    context.set_device(Base::get_local_gpu(i).get_device_id());
    functors_.forward_per_gpu(..., Base::get_local_gpu(i).get_stream());
  }
  // do reduce scatter
  size_t recv_count = Base::get_batch_size_per_gpu(is_train) * 
                      Base::get_slot_num() *
                      Base::get_embedding_vec_size();
  functors_.reduce_scatter(recv_count, embedding_feature_tensors_,
                           Base::get_output_tensors(is_train), Base::get_resource_manager());
  // scale for combiner=mean after reduction
  if (Base::get_combiner() == 1) {
    size_t send_count = Base::get_batch_size(is_train) * Base::get_slot_num() + 1;
    functors_.all_reduce(send_count, Base::get_row_offsets_tensors(is_train),
                         row_offset_allreduce_tensors_, Base::get_resource_manager());
    // do average
    functors_.forward_scale(Base::get_batch_size(is_train), Base::get_slot_num(),
                            Base::get_embedding_vec_size(), row_offset_allreduce_tensors_,
                            Base::get_output_tensors(is_train), Base::get_resource_manager());
  }
  return;
}

我们先看forward，首先从当前gpu的hashmap做lookup，也就是functors_.forward_per_gpu，此时不需要做节点间通信，因为当前gpu对应的数据的key都在当前gpu。
接着做了reduce scatter，这个通信可以参考这篇官方文档

做完forward的时候，每个gpu的数据是batch size条，但是每条数据里的每个slot是一部分key。做完reduce scatter后，数据是完整的了，并且每个gpu上分到一部分完整的数据。

我们假设一共有2个gpu，batch size为2，一共3个slot，那么上面的过程如下：

如果是要做mean pooling，还需要做再做一次all reduce，拿到每个sample每个slot里的key的总个数（把csr里的offset求个allreduce，就可以得到全局offset了），然后把embedding的值除以这个个数，也就是求了平均。

再看一看backward

void backward() override {
  // Read dgrad from output_tensors -> compute wgrad
  // do all-gather to collect the top_grad
  size_t send_count =
      Base::get_batch_size_per_gpu(true) * Base::get_slot_num() * Base::get_embedding_vec_size();
  functors_.all_gather(send_count, Base::get_output_tensors(true), embedding_feature_tensors_,
                       Base::get_resource_manager());
  // do backward
  functors_.backward(...);
  return;
}

首先做all gather，每个gpu拿到当前batch所有样本的梯度，然后更新本地每个gpu上的参数。

void update_params() override {
  #pragma omp parallel num_threads(Base::get_resource_manager().get_local_gpu_count())
  {
    size_t id = omp_get_thread_num();
    CudaDeviceContext context(Base::get_local_gpu(id).get_device_id());
    // accumulate times for adam optimizer
    Base::get_opt_params(id).hyperparams.adam.times++;
    // do update params operation
    functors_.update_params(...);
  }
  return;
}

dump分为dump参数与dump优化器状态，两者代码比较类似，如下是dump参数：

// dump hash table from GPUs
for (size_t id = 0; id < local_gpu_count; id++) {
  // dump key
  hash_tables[id]->dump
  // 拷贝到内存
  cudaMemcpyAsync(...,cudaMemcpyDeviceToHost,...)
  // dump value
  functors_.get_hash_value
  // 拷贝到内存
  cudaMemcpyAsync(...,cudaMemcpyDeviceToHost,...)
}
functors_.sync_all_gpus(...)
for (size_t id = 0; id < local_gpu_count; id++) {
  // 每个gpu上参数总大小
  size_t size_in_B = count[id] * (sizeof(TypeHashKey) + sizeof(float) * embedding_vec_size);
  // memcpy到file_buf 
  ...
  // rank0节点负责写文件
  if (Base::get_resource_manager().is_master_process()) {
    weight_stream.write(file_buf.get(), size_in_B);
  } else {
    // 其他节点把数据发给rank0节点
    MPI_Send(file_buf.get(), size_in_B, ...);
  }
}
// rank0节点收到数据
if (Base::get_resource_manager().is_master_process()) {
  for (int r = 1; r < Base::get_resource_manager().get_num_process(); r++) {
    for (size_t id = 0; id < local_gpu_count; id++) {
      ...
      MPI_Recv(...);
      weight_stream.write(file_buf.get(), size_in_B);
    }
  }
}
// 释放gpu显存
cudaFree

注意到dump的时候需要把参数和优化器状态都通过MPI_SEND发给一个节点，参数规模比较大时，0号节点就会成为瓶颈。不如每个节点dump自己的参数，还可以按分片组织参数。

load时候，每个节点都会加载所有模型文件，然后判断每个key是否属于自己：

TypeHashKey key = key_ptr[...];
size_t gid = key % global_gpu_count; // global GPU ID
int dst_rank = get_process_id_from_gpu_global_id(gid);  // node id
if (my_rank == dst_rank) {
    memcpy(...)
} else {
  continue;
}

（2）LocalizedSlotSparseEmbeddingHash

void forward(bool is_train) override {
    CudaDeviceContext context;
    for (size_t i = 0; i < Base::get_resource_manager().get_local_gpu_count(); i++) {
      context.set_device(Base::get_local_gpu(i).get_device_id());  // set device
      functors_.forward_per_gpu(...);
    }
    functors_.all2all_forward(...);
    // reorder：重新组织收到的数据buffer
    functors_.forward_reorder(...);
    // 保存每个sparse参数对应对的slot id
    functors_.store_slot_id(...);
    return;
}

all2all_forward的过程如下：

reorder的过程可以这样理解：

之所以要保存参数对应的slot id，是因为每个gpu上存不同的slot，加载的时候需要知道加载哪个slot的参数。

create_network主要是创建神经网络的各个层。
执行完前向反向后，首先多卡之间会平均梯度，然后再更新dense参数。

void Network::exchange_wgrad() {
  CudaDeviceContext context(get_device_id());
  ncclAllReduce((const void*)wgrad_tensor_.get_ptr(),
                (void*)wgrad_tensor_.get_ptr(), wgrad_tensor_.get_num_elements(),
                ncclFloat, ncclSum, gpu_resource_->get_nccl(),
                pu_resource_->get_stream()));
}

说到这里，基本上整个训练流程应该是清楚了。这里再说一下混合精度训练，hugectr在官方介绍中的Highlighted features提到了混合精度训练。在hugectr中，embedding层的参数存储还是用的全精度float32。如果配置了使用mixed_precision，那么embedding层优化器状态用的是半精度fp16，其输出也为半精度fp16，并且其他dense支持fp16计算的层也会使用fp16计算。可以看这篇文章。

上面大致了解了hugectr之后，可以再运行一下hugectr官方代码库的例子。

看完代码后，这里小小的总结一下hugectr的优缺点：

优点：

• 大batch训练 / 复杂模型训练，sparse参数分布在多个gpu上，worker与ps在同一个进程。支持混合精度训练。
• embedding表存在gpu中，sparse参数通信比cpu的方式更快。embedding在本节点lookup后，先做sum/embedding，减少通信量。dense参数通信也是通过gpu。
• 三级流水线，隐藏数据时间：解析数据（data reader worker）、拷贝数据到gpu（collector）、训练（前向反向+更新参数）。

缺点：

• 每个节点下载解析全量数据，数据量大时候带宽会成为瓶颈。
• dump模型的时候只有0号节点节点执行，参数量大的时候0号节点会成为瓶颈。
• lookup之后需要节点间交换样本：DistributedSlotSparseEmbeddingHash与LocalizedSlotSparseEmbeddingHash都需要通信交换样本，两种模式相比之间没有明显区别优势。
• 事先需要分配好参数服务器的显存，并且大小固定，不适合参数量大的场景。

参考文献

【1】https://developer.nvidia.com/blog/introducing-merlin-hugectr-training-framework-dedicated-to-recommender-systems/