HugeCTR源码阅读

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的初始化