0%

ARK: GPU-driven Code Execution for Distributed Deep Learning 论文主要包括两个工作:改善 GPU Driven 模式下的 DMA 调用方式、使用 GPU 用户态调度避免 CUDA Kernel 函数调度带来的延迟。笔者主要阅读了 ARK 论文中的第一个部分,即使用外部 DMA 驱动代替由 GPU thread 启动的 DMA 启动。

ARK 在背景介绍中谈到,集合通讯中存在的两种控制模式 CPU Driven 和 GPU Driven。CPU Driven 模式的缺点主要在于小数据包情况下,控制 overhead 不可忽略,通信带宽大幅度降低,而 GPU Driven 模式的缺点主要在于大数据包情况下,GPU 中 L2 Cache 利用率下降的问题。

事实上,在 NCCL 中,并没有使用 CPU Driven 这一模式,所以论文中使用 CPU Driven 作为 baseline 具有一定的误导性。NCCL 在 2.13.4-1 版本中已经实现了论文中的软件 GPU Driven 模式,即将 DMA 启动的部分从 GPU Thread 卸载到 CPU Thread 上面,从而避免了 CUDA IPC 模式下MMIO 带来的性能损耗。

所以可以得到,论文中的硬件 DMA 模型相比于现阶段 NCCL 的模型将会在小数据包的情况下具有一定的性能改善。这是因为现阶段 NCCL 中 GPU - CPU 通信需要通过共享内存模式进行控制位的同步,而小数据包情况下,即使是共享内存的方式,其通信延迟也是不可忽略的。使用一个 GPU 能够直接控制的 DMA 能够避免共享内存这一步的通信延迟,因而能够很大程度上改善小数据包情况下的延迟。在 NCCL 的测试数据下,使用 CudaMemcpy 来搬运数据,在单次通讯的数据小于 1M 时性能是低于默认方案的。

但是,论文中提出的方案有两点要求:Hardware DMA 与 GPU 在同一 PCIe switch 下,GPU 能够直接控制 Hardware DMA。这种要求的部署难度是比较大的,现阶段容易实现的 RNIC 方案,如果使用 GPU Thread 直接控制,控制延迟在几十 us 左右,而如果使用 CPU Thread 来控制 RNIC,虽然控制延迟降低到几 us,但是无法消除共享内存带来的延迟。另外,本论文的所有测试结果都是在 PCIe 上面测试得到的,在 NVLINK 下能否得到性能提升,还是一个疑问。

相关链接

NVIDIA 集合通信库(NCCL)是一个实现了多 GPU collective communication 的集合通信库。NCCL 在设计上是以 GPU Driven 为核心的,而非传统的 CPU Driven 代码,以获取更小的传输延迟和更大的传输带宽。在源码中,NCCL 有许多实现上的细节,本文将分享 NCCL 源码中有关网络传输以及 GDR 部分的相关内容。

GPU Driven 模式

在讲解网络以及 RDMA 部分的实现之前,首先需要了解 GPU Driven 这一核心的思路。由于 CPU 与 GPU 之间存在着通信的延迟,如果使用 CPU 进行数据链路的控制,在数据包较小时,CPU - GPU 中间的通信延迟就会显得不可忽略。以一个经典的 CPU 模型来展示:

GPU Driven 就是绕过 CPU 线程的控制,直接使用 GPU 线程去做数据链路的控制。在具体实现上,CPU Driven 其实就是传统的事件\阻塞模型,当 GPU 数据准备完毕后,通知 CPU 去做一些控制工作;而 GPU Driven 则是通过轮询某个标识位,让控制工作去等待数据准备完毕。这样就能够避免线程调度所带来的不可避免的延迟。听起来,GPU Driven 确实延迟会很低,但是不可避免地其 CPU 占用率会较高。但是当数据的传输工作较为密集时,CPU 占用率的提升就显得可以接收,这正是 NCCL 的应用场景。

作为集合通信库,NCCL 完成每一个集合通信操作,例如 allreduce,都是以小块的形式去发送数据,但是当一个集合通信操作被下发到 NCCL 中时,未来所要完成的所有操作都是确定的了。因此,每一个节点所需要接收和发送的数据量都是已经确定好的,这样就能够实现控制事件提前下发并等待数据准备完毕。

网络通信

代理线程

nccl 中所有网络 IO 都是由代理线程完成的,代理线程会以轮询的模式处理已经下发的网络请求。代理线程接收控制消息的方式与接收下发网络请求的方式所有不同。控制消息是通过 socket 接收的,是一个典型的事件通知模式。

网络请求的下发则是通过队列的模式来完成接收的,proxy 会为每一个 rank 创建一个 op 池,网络请求会被下发到对应的 op pool 中。当有可以处理的网络请求时,代理线程将会优先处理这些请求。当已收割的网络请求处理完毕时,代理线程会尝试从 op pool 中收割全部下发的请求,如果 op pool 中也没有请求,则会进入阻塞状态。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
void* ncclProxyProgress(void *proxyState_) {

...

int lastIdle = 0;
int proxyOpAppendCounter = 0;
struct ncclProxyArgs profArgs; // Only used for profiling purposes
while ((state->stop == false || (state->stop == true && state->active)) && *proxyState->abortFlag == 0) {
int idle = 1;
// 处理已经完成收割并且转换为 args 的 op
ncclResult_t ret = progressOps(proxyState, state, state->active, &idle);

// 因为 pool 是一个临界区,所以这里通过调控 pool 的 access 频率来降低临界竞争
if (idle || (++proxyOpAppendCounter == ncclParamProgressAppendOpFreq())) {
int added = 0;
proxyOpAppendCounter = 0;

// 尝试收割全部下发的 op,并且转换为 args。如果没有 op 下发,则会进入阻塞
if (state->stop == false)
ret = ncclProxyGetPostedOps(proxyState, &added);

if (added == 0) {
sched_yield(); // No request progressed. Let others run.
}
}
lastIdle = idle;
}
return NULL;
}

网络收发中的代理线程在总体上是轮询 + 事件通知的模式。当没有下发的任务时,线程会进入阻塞。这里为了避免竞争,NCCL 也做出了优化。

网络事件的下发

NCCL 中的集合操作会经过一个很长的调用链去转换为网络收发操作,这里先忽略其中的转换。在使用以下参数运行all_reduce_perf时,我们观察网络不同的事件下发:

1
./all_reduce_perf -b 4M -e 4M -f 2 -g 1 -t 1 -n 20

这个 allreduce 操作中,每个节点所需要发送的数据量为 4M,数据发送的类型为 float32,因此一共需要发送 1M 个数据点。通过NCCL 中预先埋下的日志分析,我们可以看到,allreduce 最终被分解为两次网络发送和两次网络接收请求。其中,单次网络发送数据的数据量为 2M。allreduce 操作虽然需要经过很多次网络传输,但是 NCCL 只是将其分解为了 reduce-scatter 和 all-gather 这两个操作,并将这两个操作需要的总发送数据量下发到网络发送请求中。下发的网络事件不会立刻进行数据传输,而是由代理线程轮询并等待数据准备完毕。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
static ncclResult_t sendProxyProgress(struct ncclProxyState* proxyState, struct ncclProxyArgs* args) {
// --发送完毕的-- done --正在发送的-- transimitted --还没有发送的-- posted ----

// 检查是否有可以下发的请求
if (sub->posted < sub->nsteps && sub->posted < sub->done + maxDepth) {
// Post buffers to the GPU
...
}
// 如果有下发的网络发送请求,则检查当前是否可以发送数据
if (sub->transmitted < sub->posted && sub->transmitted < sub->done + NCCL_STEPS) {

// 检查标识位是否允许发送
if (sizesFifo[buffSlot] != -1 && ((*recvTail > (sub->base+sub->transmitted)) || p == NCCL_PROTO_LL)){
// STEP1 : 检查数据是否准备完毕

// STEP2 : 发送数据

}
}
if (sub->done < sub->transmitted) {
// 检查是否有数据发送完毕,这一步主要是由于 RDMA 通信的特点
}
}

这种网络事件的设计模式正是 NCCL 中 GPU Driven 的一个具体体现。当 GPU Kernel 函数完成数据的准备后,就可以通过修改队列尾的方式激活已经提前下发好的网络请求了。

RDMA 通信

与 rsocket 等通用的网络库不同,NCCL 中的 RDMA 通信也针对应用场景采用了不同的设计。使用 RDMA 通信,一个原则性要求是发送方在 post_send 之前,必须要确保接收方已经使用 post_receive_wr 下发了一个接收操作,这样才能够保证消息的正确传递。正常情况下,RDMA 网络库都会在初始化时下发一部分接收操作,并且采用计时的方式进行流量控制。而 NCCL 中则采取了不同的模式——接收方通知发送方进行发送,这是同样也是因为当一个操作被下发之前,所有的网络 IO 次数以及数量都是确定好的。

缓冲区的分配

控制缓冲区

控制缓冲区是用来调控收发请求的,一次操作对应一个槽。这部分是注册在 cudaMallocHost 中的锁页内存中,因为 GPU 和 CPU 都需要去访问这一块缓冲区。

收发缓冲区

这部分缓冲区会根据是否支持 gdr 选择分配在内存或者是显存上。如果分配到显存上面,会使用 gdr ,收发直接到显存上。gdr 优先使用 dmabuf 方式,否则使用 peermem 方式。

RDMA 通信优化

  1. 按组进行接收 comp channel 通知
  2. inline 标识位减少延迟
  3. 减少接收 comp channel 通知次数