MXNet源码分析 | KVStore进程内通信

本文主要基于MXNet1.6.0版本进行分析。

MXNet的KVStore模块下有几个比较重要的类。KVStore是一个抽象类，提供了一些通用的API，例如Init、Push和Pull等。因为KVSotre支持int和string两种类型的key，所以这些API以不同类型的key作为参数，提供了两种重载。

KVStoreLocal继承自KVStore，负责进程内通信。它主要维护了以下变量：负责不同设备间通信的comm_，机器上的页锁定内存（不进行页交换，一直在物理内存中），本机的key-val buffer，字符串key到整型key的映射以及整型key到字符串key的映射。在其构造函数中，它会根据传入的设备类型（CPU、GPU等）创建不同的通信对象。

KVStoreDist和KVStoreDistServer分别实现了参数服务器架构中的worker和server节点。KVStoreDist继承自KVStoreLocal，内部维护了一个指向ps::Worker<char>对象的指针。这个ps::Worker<char>对象实现在PS-Lite里面，主要负责执行push和pull操作。这里我们只需要知道这个ps::KVWorker<char>对象能够执行push和pull操作就可以了，至于内部如何实现的，后续会有文章进行进一步的分析。

KVStoreDistServer负责接收、处理各个worker的请求，并对这些请求进行响应。它内部维护了一个负责处理请求的ps::KVServer<char>对象（同样实现在PS-Lite中），以及用来更新模型参数的Updater方法。

classDiagram Comm <|.. CommCPU Comm <|.. CommDevice CommDevice <|-- CommDeviceTree Comm <.. KVStoreLocal KVStore <|.. KVStoreLocal KVStoreLocal <|-- KVStoreDist KVStoreLocal <|-- KVStoreNCCL KVStoreDist ..> KVStoreDistServer class Comm { <<abstract>> +Init() void +Reduce() const NDArray& +Broadcast() void +BroadcastRowSparse() void #Context pinned_ctx_ } class CommCPU { -ReduceSumCPU() void -ReduceSumCPUExSerial() void -unordered_map<int, BufferEntry> merge_buf_ -size_t bigarray_bound_ } class CommDevice { +InitBuffersAndComm() void +ReduceCompressed() const NDArray& +InitMergeBuffer() void -EnableP2P() void -unordered_map<int, BufferEntry> merge_buf_ } class CommDeviceTree { +ReduceInner() const NDArray& +BroadcastInner() void -int depth_ -int gpuarray_bound_ } class KVStore { <<abstract>> +Create() static KVStore* +Push() void +Pull() void +Barrier() void +RunServer() void #Updater updater_ #StrUpdater str_updater_ } class KVStoreLocal { +Init() void -InitImpl() void -PushImpl() void -PullImpl() void -PushPullImpl() void #GroupKVPairsPush() void #GroupKVPairsPull() void #GroupKVPairs() void #LookupKeys() void #Comm* comm_ #unordered_map<int, NDArray> local_ #unordered_map<string, int> str_key_dict_ #unordered_map<int, string> reverse_str_key_dict_ } class KVStoreDist { -EncodeDefaultKey() void -EncodeCompressedKey() void -EncodeRowSparseKey() void -unordered_map<int, PSKV> ps_kv_ -KVWorker<char>* ps_worker_ -KVStoreDistServer* server_ -size_t bigarray_bound_ -unordered_map<int, NDArray> comm_buf_ -unordered_map<int, NDArray> compr_buf_ -unordered_map<int, NDArray> residual_ } class KVStoreDistServer { -ApplyUpdates() void -DefaultStorageResponse() void -DataHandleDefault() void -unordered_map<int, NDArray> store_ -unordered_map<int, NDArray> update_buf_ -unordered_map<int, NDArray> decom_buf_ -KVServer<char>* ps_server_ }

KVStore创建

MXNet支持6种不同的kvstore：local，device，dist_sync，dist_async，dist_sync_device，dist_async_device。其中local和device会创建支持单机训练的KVStore，而以dist开头的命令会创建支持分布式训练的KVStore。通过向KVStore::Create传入不同的参数，即可创建不同类型的KVStore。先判断有无dist，如果有就创建KVStoreDist，如果有dist又有async，就给server发送异步训练的消息。如果命令中包含device，就在GPU上聚合。

KVStore* KVStore::Create(const char *type_name) {
  std::string tname = type_name;
  std::transform(tname.begin(), tname.end(), tname.begin(), ::tolower);
  KVStore* kv = nullptr;
  bool use_device_comm = false;
  auto has = [tname](const std::string& pattern) {
    return tname.find(pattern) != std::string::npos;
  };
  if (has("device")) {
    use_device_comm = true;
  }

  if (has("dist")) {
    kv = new kvstore::KVStoreDist(use_device_comm);
    if (!has("_async") && kv->IsWorkerNode() && kv->get_rank() == 0) {
      // configure the server to be the sync mode
      kv->SendCommandToServers(static_cast<int>(kvstore::CommandType::kSyncMode), "");
    }
  } else {
    if (has("nccl")) {
      kv = new kvstore::KVStoreNCCL();
    } else {
      kv =  new kvstore::KVStoreLocal(use_device_comm);
    }
  }
  kv->type_ = tname;
  return kv;
}

进程内通信

对于单机训练以及分布式训练中每个进程内部的设备间数据交换需求，MXNet基于Reduce和Broadcast语义实现了进程内通信机制。KVStoreLocal负责进行进程内通信，它实现了抽象类KVStore中的大部分接口。在KVStoreLocal的构造函数中，它会根据下面的代码来创建负责进程内通信的对象。KVStoreLocal支持3种类型的通信，分别是CommCPU、CommDevice以及CommDeviceTree。顾名思义，CommCPU这个类把Reduce和Broadcast操作放在CPU上，而CommDevice和CommDeviceTree这两个类则实现了GPU上的Reduce和Broadcast，其中CommDeviceTree则是基于树形拓扑实现的。

explicit KVStoreLocal::KVStoreLocal(bool use_device_comm) : KVStore() {
  if (use_device_comm) {
    bool tree = dmlc::GetEnv("MXNET_KVSTORE_USETREE", 0)
    if (tree) {
      comm_ = new CommDeviceTree();
    } else {
      comm_ = new CommDevice();
    }
  } else {
    comm_ = new CommCPU();
  }
}

进程内Push实现

KVStoreLocal类的PushImpl函数实现了单进程的Push操作，该函数会先调用GroupKVPairsPush把具有相同key的NDArray汇总到同一个vector中。然后，针对每个key，调用相关通信对象实现的Reduce函数对数据进行规约。最后，如果设置了用来更新参数的updater_，那么就会调用相关的更新算法去更新权重；否则，会把规约后的结果拷贝或者移动到保存本地变量的buffer中。

virtual void 
KVStoreLocal::PushImpl(const std::vector<int>& keys, 
                       const std::vector<NDArray>& values, 
                       int priority) {
  std::vector<int> uniq_keys;
  std::vector<std::vector<NDArray>> grouped_val;
  GroupKVPairsPush(keys, values, &uniq_keys, &grouped_val, false);
  for (size_t i = 0; i <uniq_keys.size(); ++i) {
    int key = uniq_keys[i];
    const NDArray& merged = comm_->Reduce(key, grouped_val[i], priority);
    NDArray& local = local_[key];
    if (updater_ != nullptr) {
      if (kye_type_ == kStringKye && str_updater_ != nullptr) {
        str_updater_(str_key, merged, &local);
      } else {
        updater_(key, merged, &local);
      }
    } else {
      if (merged.storage_type() != local.storage_type()) {
        local = merged.Copy(local.ctx());
      } else {
        local = merged;
      }
    }
  }
}

在上面这个函数中，最终执行规约操作的是由本地通信对象实现的Reduce方法。前面我们提到，KVStoreLocal可以创建3中不同类型的本地通信对象——CommCPU、CommDevice以及CommDeviceTree。这里我们先来看下CommCPU这个类是如何实现Reduce操作的。

CommCPU上的通信

CommCPU::Reduce将输入vector<NDArray>& src的每个元素求和并返回。当src只有一个元素时直接返回src[0]；否则就按照下面的代码，把数据规约到变量reduce[0]中。可以看到，下面的代买调用PushAsync方法把一个lambda表达式压入到依赖引擎中，这个lambda表达式首先捕获了reduce变量，然后在函数体中调用ReduceSumCPU方法在CPU上执行Reduce操作，最终操作的结果会存放到reduce[0]中。

std::vector<Engine::VarHandle> const_vars(src.size() - 1);
std::vector<NDArray> reduce(src.size());
CopyFromTo(src[0], &buf_merged, priority);
reduce[0] = buf_merged;

if (buf.copy_buf.empty()) {
  buf.copy_buf.resize(src.size()-1);
   for (size_t j = 0; j < src.size() - 1; ++j) {
     // allocate copy buffer
     buf.copy_buf[j] = NDArray(src[0].shape(), pinned_ctx_, false, src[0].dtype());
   }
}
CHECK(stype == buf.copy_buf[0].storage_type()) 
  << "Storage type mismatch detected. " << stype << "(src) vs. "
  << buf.copy_buf[0].storage_type() << "(buf.copy_buf)";
 for (size_t i = 1; i < src.size(); ++i) {
   CopyFromTo(src[i], &(buf.copy_buf[i-1]), priority);
   reduce[i] = buf.copy_buf[i-1];
   const_vars[i-1] = reduce[i].var();
}

Engine::Get()->PushAsync(
  [reduce, this](RunContext rctx, Engine::CallbackOnComplete on_complete) {
  ReduceSumCPU(reduce);
  on_complete();
  }, Context::CPU(), const_vars, {reduce[0].var()},
  FnProperty::kCPUPrioritized, priority, "KVStoreReduce");

在下面这个函数中，通过MSHADOW_TYPE_SWITCH这个宏把in_data中的每个NDArray转换成对应数据类型的指针，然后再调用ReduceSumCPUImpl方法去做Reduce。

inline void ReduceSumCPU(const std::vector<NDArray> &in_data) {
    MSHADOW_TYPE_SWITCH(in_data[0].dtype(), DType, {
      std::vector<DType*> dptr(in_data.size());
      for (size_t i = 0; i < in_data.size(); ++i) {
        TBlob data = in_data[i].data();
        CHECK(data.CheckContiguous());
        dptr[i] = data.FlatTo2D<cpu, DType>().dptr_;
      }
      size_t total = in_data[0].shape().Size();
      ReduceSumCPUImpl(dptr, total);
    });
  }

下面的代码展示了模板函数ReduceSumCPUImpl，它的第一个参数可以接收由任意元素类型的指针组成的vector。这个函数主要负责将Reduce操作并行化，如果需要Reduce的张量的大小没有超过bigarray_bound_，或者只有一个线程执行Reduce操作，那么这个函数会直接调用一个接收3个参数的ReduceCPUSum的重载版本；否则，它会使用OpenMP针对较大的张量进行并行化的Reduce。

template<typename DType>
inline void ReduceSumCPUImpl(std::vector<DType*> dptr, size_t total) {
  const size_t step = std::min(bigarray_bound_, static_cast<size_t>(4 << 10));
  long ntask = (total + step - 1) / step; // NOLINT(*)
  if (total < bigarray_bound_ || nthread_reduction_ <= 1) {
    ReduceSumCPU(dptr, 0, total);
  } else {
    #pragma omp parallel for schedule(static) num_threads(nthread_reduction_)
    for (long j = 0; j < ntask; ++j) { // NOLINT(*)
      size_t k = static_cast<size_t>(j);
      size_t begin = std::min(k * step, total);
      size_t end = std::min((k + 1) * step, total);
      if (j == ntask - 1) CHECK_EQ(end, total);
      ReduceSumCPU(dptr, begin, static_cast<index_t>(end - begin));
    }
  }
}

这里的ReduceSumCPU也是一个模板函数，内部实现了Reduce的操作逻辑，具体代码如下所示。

template<typename DType>
inline static void ReduceSumCPU(const std::vector<DType*> &dptr, size_t offset, index_t size) {
  using namespace mshadow;  // NOLINT(*)
  Tensor<cpu, 1, DType> in_0(dptr[0] + offset, Shape1(size));
  for (size_t i = 1; i < dptr.size(); i+=4) {
    switch (dptr.size() - i) {
      case 1: {
        Tensor<cpu, 1, DType> in_1(dptr[i] + offset, Shape1(size));
        in_0 += in_1;
        break;
      }
      case 2: {
        Tensor<cpu, 1, DType> in_1(dptr[i] + offset, Shape1(size));
        Tensor<cpu, 1, DType> in_2(dptr[i+1] + offset, Shape1(size));
        in_0 += in_1 + in_2;
        break;
      }
      case 3: {
        Tensor<cpu, 1, DType> in_1(dptr[i] + offset, Shape1(size));
        Tensor<cpu, 1, DType> in_2(dptr[i+1] + offset, Shape1(size));
        Tensor<cpu, 1, DType> in_3(dptr[i+2] + offset, Shape1(size));
        in_0 += in_1 + in_2 + in_3;
        break;
      }
      default: {
        Tensor<cpu, 1, DType> in_1(dptr[i] + offset, Shape1(size));
        Tensor<cpu, 1, DType> in_2(dptr[i+1] + offset, Shape1(size));
        Tensor<cpu, 1, DType> in_3(dptr[i+2] + offset, Shape1(size));
        Tensor<cpu, 1, DType> in_4(dptr[i+3] + offset, Shape1(size));
        in_0 += in_1 + in_2 + in_3 + in_4;
        break;
      }
    }
  }
}

进程内Pull实现

进程内Pull操作实现在KVStoreLocal::PullImpl函数中，类似于Push操作，在Pull之前会把具有相同key的NDArray汇总到一个vector里，然后再调用通信对象实现的Broadcast函数把数据广播到每个线程。

virtual void KVStoreLocal::PullImpl(const std::vector<int>& keys,
                                    const std::vector<NDArray*>& values,
                                    int priority,
                                    bool ignore_sparse) {
  std::vector<int> uniq_keys;
  std::vector<std::vector<NDArray*>> grouped_keys;
  GroupKVPairsPull(keys, values, &unique_keys, &grouped_vals, ignore_sparse);
  
  for (size_t i = 0; i < uniq_keys.size(); ++i) {
    int key = uniq_keys[i];
    const NDArray& local = local_[key];
    comm_->Broadcast(key, local, grouped_vals[i], priority);
  }
}

CPU实现

Broadcast这一部分的实现逻辑还是比较清晰的，如果原始数据存放在内存中，那么就直接进行拷贝；否则，会先把数据从GPU显存拷贝到页锁定内存（pinned memory），然后再进行数据的拷贝。

void Broadcast(int key, const NDArray& src,
               const std::vector<NDArray*> dst, int priority) override {
  int mask = src.ctx().dev_mask();
  if (mask == Context::kCPU) {
    for (auto d : dst) CopyFromTo(src, d, priority);
  } else {
    // First copy data to pinned_ctx, then broadcast.
    // Note that kv.init initializes the data on pinned_ctx.
    // This branch indicates push() with ndarrays on gpus were called,
    // and the source is copied to gpu ctx.
    // Also indicates that buffers are already initialized during push().
    auto& buf = merge_buf_[key].merged_buf(src.storage_type());
    CopyFromTo(src, &buf, priority);
    for (auto d : dst) CopyFromTo(buf, d, priority);
  }
}

GPU实现

如果merge_buf_没有被初始化，那么就先把src拷贝到一个随机的设备上，然后再把数据从随机设备拷贝到其他所有设备；如果merge_buf_已经被初始化，那么src中的数据会首先拷贝到merge_buf_，然后再拷贝到目标设备上。

void Broadcast(int key, const NDArray& src,
               const std::vector<NDArray*> dst, int priority) override {
  if (!inited_) {
    // copy to a random device first
    int dev_id = key % dst.size();
    CopyFromTo(src, dst[dev_id], priority);
    for (size_t i = 0; i < dst.size(); ++i) {
      if (i != static_cast<size_t>(dev_id)) {
        CopyFromTo(*dst[dev_id], dst[i], priority);
      }
    }
  } else {
    auto& buf_merged = merge_buf_[key].merged_buf(src.storage_type());
    CopyFromTo(src, &buf_merged, priority);
    for (auto d : dst) {
      CopyFromTo(buf_merged, d, priority);
    }
  }
}