上一章中,我们确实能够使用8张显卡32张网卡,然而传输速度仅为 287.089 Gbps,只达到了总带宽 3200 Gbps 的 9.0%。

如果我们仔细观察上一章程序的运行过程,我们会看到程序一开始的传输速度在 100 Gbps 左右,然后迅速升高至 260 Gbps 左右,然后逐渐缓慢上升至 287 Gbps。这一现象十分可疑。这说明了程序一开始的时候受到了一个比较大的延迟,随后便能够一直保持相同的传输速度。这一章让我们来解决这个慢速启动的问题。我们把本章的程序命名为 10_warmup.cpp

我们很容易能够猜到,这个现象可能是因为第一个 WRITE 操作花费了特别长的时间导致的。而之所以第一个操作花费了特别长的时间,是因为除了第一张网卡以外,其他的网卡都还没有与客户端建立连接。

预热

要解决这个问题也很简单,在正式进入速度测试之前,我们可以引入一个预热阶段。在预热阶段,我们让每一个网卡都向对方发送一个 WRITE 操作,这样就确保了两端的连接都已经建立。在预热阶段结束后,我们再开始正式的速度测试。

我们在服务器端的状态机里面加入一些关于预热的状态:

struct RandomFillRequestState {
  enum class State {
    kWaitRequest,
    kPostWarmup,  // Added
    kWaitWarmup,  // Added
    kWrite,
    kDone,
  };

  struct WriteState {
    bool warmup_posted = false;  // Added
    size_t i_repeat = 0;
    size_t i_buf = 0;
    size_t i_page = 0;
  };

  // ...
  size_t posted_warmups = 0;
  size_t cnt_warmups = 0;
  // ...
};

当服务器端收到 RANDOM_FILL 请求时,我们进入 kPostWarmup 状态:

struct RandomFillRequestState {
  // ...

  void HandleRequest(Network &net, RdmaOp &op) {
    // ...
    // Generate random data and copy to local GPU memory
    // ...

    // Prepare for warmup
    write_states.resize(connect_msg->num_gpus);
    state = State::kPostWarmup;
  }
};

然后我们增加一个 PostWarmup(gpu_idx) 函数,为这张显卡对应的所有网卡提交一个 WRITE 操作。当所有显卡都提交完了预热操作之后,我们让状态机进入 kWaitWarmup 状态:

struct RandomFillRequestState {
  // ...

  void PostWarmup(size_t gpu_idx) {
    // Warmup the connection.
    // Write 1 page via each network
    auto &s = write_states[gpu_idx];
    if (s.warmup_posted) {
      return;
    }

    auto page_size = request_msg->page_size;
    auto &group = (*net_groups)[gpu_idx];
    for (size_t k = 0; k < group.nets.size(); ++k) {
      auto net_idx = group.GetNext();
      const auto &mr =
          connect_msg->mr((gpu_idx * nets_per_gpu + net_idx) * buf_per_gpu);
      auto write = RdmaWriteOp{ ... };
      group.nets[net_idx]->PostWrite(std::move(write),
                                     [this](Network &net, RdmaOp &op) {
                                       HandleWarmupCompletion(net, op);
                                     });
    }
    s.warmup_posted = true;
    if (++posted_warmups == connect_msg->num_gpus) {
      state = State::kWaitWarmup;
    }
  }
};

然后在预热操作的回调函数中,我们检查是否所有的预热操作都已经完成。如果是,我们就进入 kWrite 状态:

struct RandomFillRequestState {
  // ...

  void HandleWarmupCompletion(Network &net, RdmaOp &op) {
    if (++cnt_warmups < connect_msg->num_nets) {
      return;
    }
    printf("Warmup completed.\n");

    // Prepare RDMA WRITE the data to remote GPU.
    printf("Started RDMA WRITE to the remote GPU memory.\n");
    total_write_ops = connect_msg->num_gpus * buf_per_gpu *
                      request_msg->num_pages * total_repeat;
    write_op_size = request_msg->page_size;
    write_states.resize(connect_msg->num_gpus);
    write_start_at = std::chrono::high_resolution_clock::now();
    state = State::kWrite;
  }
};

最后就是在服务器端的主循环中,我们判断如果状态机处于 kPostWarmup 状态,我们就调用 PostWarmup(gpu_idx) 函数:

int ServerMain(int argc, char **argv) {
    // ...
    while (s.state != RandomFillRequestState::State::kDone) {
      for (size_t gpu_idx = 0; gpu_idx < net_groups.size(); ++gpu_idx) {
        for (auto *net : net_groups[gpu_idx].nets) {
          net->PollCompletion();
        }
        switch (s.state) {
        case RandomFillRequestState::State::kWaitRequest:
          break;
        case RandomFillRequestState::State::kPostWarmup:  // Added
          s.PostWarmup(gpu_idx);
        case RandomFillRequestState::State::kWaitWarmup:  // Added
          break;
        case RandomFillRequestState::State::kWrite:
          s.ContinuePostWrite(gpu_idx);
          break;
        case RandomFillRequestState::State::kDone:
          break;
        }
      }
    }
    // ...
}

运行效果

在上面的视频中我们可以看到,在加入了预热之后,程序一开始的传输速度就能够达到 290 Gbps 左右,并且之后一直稳定在这个速度附近。最终的传输速度为 293.461 Gbps,达到了总带宽 3200 Gbps 的 9.2%。我们会继续在下一章中优化这个程序。

本章代码:https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/10_warmup.cpp