浅谈NCCL Proxy线程

目录

（1）proxyService线程用到变量的初始化

（2）proxyService线程和proxyProgress线程的创建

（3）P2P建立connection，P2P和proxyService线程建立一个Socket连接的流程

（4）proxyProgressAsync根据不同的type做不同的动作

（5）proxyProgress线程的处理过程

（6）建立NET的transport中recv侧proxyService/proxyProgress线程的函数调用流程（Send侧类似）

（7）核函数和ProxyProgress线程变量的对应关系

（8）sendProxyProgress的流程

1、Post buffers to the GPU

2、Check whether we received data from the GPU and send it to the network

3、Check whether the network has completed some send operations.

（9）recvProxyProgress的流程

（10）proxy connection建立的时机

（11）几个问题（持续补充）

参考链接：

本文档针对NCCL 2.19版本。

在intra-node，GPU与GPU之间建立P2P transport的时候，和inter-node中，通过NET建立NET transport的时候，都需要proxy线程的参与，其实总共有两个proxy线程，一个叫做proxyService线程，是每个NODE中每个GPU对应的一个，主要维护连接建立，建立transport的setup和connect阶段。另一个叫做proxyProgress线程，也是每个NODE中每个GPU对应的一个，主要在inter-node通信过程中，处理kernel和IB之间的数据交互。

ncclProxyArgs保存了通信需要的参数，proxyProgress线程会根据这些args执行相应的通信流程。

（1）proxyService线程用到变量的初始化

（2）proxyService线程和proxyProgress线程的创建

NCCL为每个Rank创建一个proxyService线程，线程处理函数为：ncclProxyService，每个Rank上的线程名字分别为：NCCL Service 0，NCCL Service 1，NCCL Service 2，NCCL Service 3（4 GPU）

NCCL为每个Rank创建一个proxyProgress线程，线程处理函数为：ncclProxyProgress，线程名字一样，都为：NCCL Progress 4（4 GPU）

（3）P2P建立connection，P2P和proxyService线程建立一个Socket连接的流程

整个流程分为四部分：

a、发起连接的客户端：P2P

b、proxy处理函数

c、proxyService线程

d、proxyProgress线程

首先P2P建立transport的时候，本Rank发起到Proxy rank（P2P的时候，Proxy rank ID就是本rank ID）的proxyService线程的连接，p2pSendSetup调用ncclProxyConnect

p2pSendSetup
  tpProxyRank = comm->topParentRanks[info->rank];
  NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_P2P, 1, tpProxyRank, &send->proxyConn));

Proxy会设置此次消息的type为ncclProxyMsgInit，并且调用ncclProxyCallBlocking，先调用ncclProxyCallAsync进行异步请求，然后ncclPollProxyResponse接收异步消息。如果这个ncclTransport定义了proxyProgress，会调用proxyProgressInit创建一个proxyProgress线程（proxyProgress其实是在init阶段创建的），主要来处理跟核函数之间的数据

（4）proxyProgressAsync根据不同的type做不同的动作

proxyProgressAsync()函数是在proxyService线程中的，当proxyService收到数据之后，该函数会根据op→type进行不同的动作。

（5）proxyProgress线程的处理过程

如果当前ncclTransport定义了proxyProgress，比如：netTransport，会调用proxyProgressInit创建一个proxyProgress线程，线程的处理函数是：ncclProxyProgress

在while循环中，progressOps函数用来执行添加的progress动作，而ncclProxyGetPostedOps是用来添加progress动作。

NCCL考虑到频繁的调用ncclProxyGetPostedOps的话，将导致小包消息传递退化，为了减少调用ncclProxyGetPostedOps的次数，NCCL定义了一个环境变量：NCCL_PROGRESS_APPENDOP_FREQ，初始值为8，就是while每循环8次，去调用一次ncclProxyGetPostedOps，添加progress动作。

并且如果检测到本次ncclProxyGetPostedOps没有做任何progress动作的添加的话，会主动调用sched_yield释放线程的执行权，让操作系统执行调度。

proxyProgress线程中，当没有opts要处理的时候，线程就进入休眠，集合通信和点到点通信API可以唤醒proxyProgress线程

（6）建立NET的transport中recv侧proxyService/proxyProgress线程的函数调用流程（Send侧类似）

节点间首先与recvpeer建立NET的transport，会与proxyService线程和proxyProgress线程交互，会触发与IB建立connection以及listen的动作（Send侧就是connect的动作）

（7）核函数和ProxyProgress线程变量的对应关系

prims_simple.h中的几个关键变量：

1）connEltsFifo

loadRecvConn

    connEltsFifo = (T*)conn→buffs[NCCL_PROTO_SIMPLE];

loadSendConn

    connEltsFifo = (T*)conn→buffs[NCCL_PROTO_SIMPLE];

connEltsFifo保存了GPU上申请的共享buff，对于Send和recv的情况分别保存对应的buff，填入数据之后，sendProxyProgress就可以把这块buff数据通过IB verbs发送出去

2）connStepPtr

uint64_t *connStepPtr;

uint64_t connStepCache; // Cache last seen value of (*connStepPtr)      connStepCache就是connStepPtr这个指针指向的具体值，方便比较判断

在load阶段，connStepPtr保存了connection的head和tail指针

在发送完数据之后，会保存为当前的step值

3）connSizesFifoPtr （保存的就是一个GPU本次要往Proxy发送数据的长度）

loadSendConn

    connSizesFifoPtr = conn→sizesFifo;         // Sizes fifo from GPU to proxy

connSizesFifoPtr保存了conn->sizesFifo，是GPU向Proxy发送数据的长度

在waitPeer里，可以发送数据了之后去更新，将数据长度写入到sendConnFifoPtr中，即sizesFifo，这样sendProxyProgress函数就能知道有数据要发送了，也知道这次写的数据的长度。

__atomic_store_n(connSizesFifoPtr+step%NCCL_STEPS, nelts*sizeof(T), __ATOMIC_SEQ_CST);

4）connOffsFifoPtr（先不care）

这个变量是Proxy给GPU发送数据的时候用的，也就是GPU接收Proxy上数据用的fifo

loadRecvConn

    connOffsFifoPtr = conn→offsFifo;           // Buffer fifo from proxy to GPU

loadSendConn

    connOffsFifoPtr = conn->offsFifo;

当满足resources->shared的时候在recvProxyProgress里更新。

（8）sendProxyProgress的流程

在第6节中介绍了调用sendProxyProgress的时机，proxy其他部分都是为sendProxyProgress服务的。

sendProxyProgress的作用就是用来在节点间通信的时候，在发送方把GPU的Send buff上的数据通过调用IB verbs传递出去，同时更新head指针。

proxy线程定义了这个数据结构来维护数据的处理，包括数据加载，发送，确认等。

struct ncclProxyArgs {
  struct ncclProxySubArgs subs[NCCL_PROXY_MAX_SUBS];
  proxyProgressFunc_t progress;
  int nsubs;
  int done;
  uint64_t opCount;
  int sliceSteps;
  int chunkSteps;
  int chunkSize;
  uint8_t /*ncclDataType_t*/ dtype;
  uint8_t /*ncclDevRedOp_t*/ redOp;
  uint8_t /*ncclPattern_t*/ pattern;
  uint8_t protocol;
  int state;
  char* sharedBuff[NCCL_STEPS];
  int sharedSize[NCCL_STEPS];

  int idle;
  uint64_t hdp_flushed;

  // Element linking
  struct ncclProxyArgs* next;
  struct ncclProxyArgs* nextPeer;
  struct ncclProxyArgs** proxyAppendPtr;
};

主要的关键字段解析：

(gdb) p *op
$169 = {
  subs = {{
      connection = 0x14764c000f30,                                                 //在发送方保存了Send connection，接收方保存的是Recv connection
      channelId = 0,                                                               //保存了第几个channel，因为所有的channel的数据都由这个Rank上这一个proxyOrogress线程来处理，每个channel对应的共享显存是单独的，不会冲突
      nsteps = 896,                                                                //这个是总的字节数划分成了多少个step，在Allreduce，Simple协议中，一个step是512KB
      nbytes = 1048576,                                                            //一个buff大小的数量，int buffSize = stepSize*args->sliceSteps=512K*2=1MB，没啥用。。。
      peer = 0,
      groupSize = 0,
      base = 0,                                                                    //相当于这次处理数据的起始，就是一个偏移量，来描述已经处理过了多少数据，下面的处理在这个基础上往后处理就好了 sub->base = ROUNDUP(resources->step, args->chunkSteps);
      posted = 12,                                                                 //在sendProxyProgress中，用来描述第一步中从GPU中post到了多少数据到proxy中，只是一个逻辑的概念，单位是args->sliceSteps，Allreduce中是2，也就是每次post2个step，也就是1KB数据。
                                                                                   //在recvProxyProgress中，用来描述第一步中从proxy中post到了多少数据到GPU中，在proxyState->ncclNet->irecv之后更新，只是一个逻辑的概念，单位是args->sliceSteps，Allreduce中是2，也就是每次post2个step，也就是1KB数据。
      received = 0,                                                                //只在recvProxyProgress中用，用来描述接收第二步中，真正接收的数据step数量，单位是args->sliceSteps，Allreduce中是2，也就是每次receive个step，也就是1KB数据。proxyState->ncclNet->test之后更新
      flushed = 0,                                                                 //在coll_net中用的，net中没用到
      transmitted = 8,                                                             //在sendProxyProgress中，用来描述第二步中从proxy线程通过IB verbs发送了多少数据出去，单位也是args->sliceSteps，Allreduce中是2，也就是说每次传递2个step，也就是1KB数据。proxyState->ncclNet->isend之后更新。
                                                                                   //在recvProxyProgress中，用来描述第二步中从proxy中post到了多少数据到GPU中，在proxyState->ncclNet->irecv之后更新，只是一个逻辑的概念，单位是args->sliceSteps，Allreduce中是2，也就是每次post2个step，也就是1KB数据。
      done = 4,                                                                    //在sendProxyProgress中用来描述第三步中有多少数据通过IB verbs发送成功了，会调用proxyState->ncclNet->test来确认，单位是args->sliceSteps，Allreduce中是2，也就是每次确认2个step，1KB数据。proxyState->ncclNet->test之后更新。
                                                                                   //在recvProxyProgress中，用来描述第四步中有多少数据通过IB verbs接收成功了。。。
      end = 0,
      requests = {0x14764c025020, 0x0, 0x14764c025068, 0x0, 0x14764c0250b0, 0x0, 0x14764c0250f8, 0x0},
      profilingEvents = {0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0}
    }, {
      connection = 0x0,
      channelId = 0,
      nsteps = 0,
      nbytes = 0,
      peer = 0,
      groupSize = 0,
      base = 0,
      posted = 0,
      received = 0,
      flushed = 0,
      transmitted = 0,
      done = 0,
      end = 0,
      requests = {0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
      profilingEvents = {0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0}
    } <repeats 31 times>},
  progress = 0x147b577f1200 <sendProxyProgress(ncclProxyState*, ncclProxyArgs*)>,
  nsubs = 1,
  done = 0,
  opCount = 0,
  sliceSteps = 2,                                                                                         //Allreduce中一个slice是两个step
  chunkSteps = 4,                                                                                         //Allreduce中一个chunk是4个step
  chunkSize = 2097152,                                                                                    //2097152 = 2*1024*1024B=2MB
  dtype = 7 '\a',
  redOp = 0 '\000',
  pattern = 1 '\001',
  protocol = 2 '\002',                                                                                    //代表NCCL_PROTO_SIMPLE协议
  state = 2,                                                                                              //ncclProxyArgs的状态，enum ncclProxyOpState { ncclProxyOpNone, ncclProxyOpReady, ncclProxyOpProgress };
  sharedBuff = {0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
  sharedSize = {0, 0, 0, 0, 0, 0, 0, 0},
  idle = 0,
  hdp_flushed = 0,
  next = 0x1476400028f0,
  nextPeer = 0x147640007c18,
  proxyAppendPtr = 0x14764c000f58
}

下面说一下sendProxyProgress中的三个步骤：

sendProxyProgress总共分为三部分

1、Post buffers to the GPU

判断条件：

if (sub->posted < sub->nsteps && sub->posted < sub->done + maxDepth)

int maxDepth = std::min(NCCL_STEPS, NCCL_SHARED_STEPS/args→nsubs);           //maxDepth的值是8，跟NCCL_STEPS一样大

就是先从GPU的buff中准备好要通过IB发出去的数据，Allreduce里是每次发512KB数据，这块数据是要发送的

每次更新sub->posted

2、Check whether we received data from the GPU and send it to the network

判断条件：

if (sub->transmitted < sub->posted && sub->transmitted < sub->done + NCCL_STEPS)

到这里已经准备好了post的数据，要求是一个buff大小（8个step），才通过IB verbs按照每次512KB的顺序发出去

每次更新sub->transmitted

3、Check whether the network has completed some send operations.

判断条件：

if (sub->done < sub->transmitted)

通过IB发出去之后，还要去确认IB是否发完，因为proxyProgress线程发送数据是异步的，需要通过test来判断一下是否发送完成。

反正要通过调用test来确认一下

每次更新sub→done

三个步骤都完成之后，需要更新一下head指针，因为head指针是proxy来更新的，通知kernel可以产生数据了。kernel是生产者，proxyProgress线程是消费者。

*sendHead = sub->base + sub->done;

举例说一下sendProxyProgress的一个流程。

（1）进行第一步，把GPU的Send buff数据post到proxyProgress线程中，这个时候是每次post 2 个step，sub→posted每次+2，每个channel轮流处理，要凑够8个step，才能继续下一步，这一步就是先攒够8个step，够了之后到下一步

（2）进行第二步，把proxyProgress中的数据通过IB verbs（proxyState->ncclNet→isend）transmit出去，也是每次transmit 2个step，sub→transmitted每次+2，每个channel轮流处理，要凑够8个step，才能继续下一步，这一步就是先攒够8个step，够了之后到下一步

（3）进行第三步，需要对上一步中transmit出去的数据确认真正完成了（proxyState→ncclNet→test），这一步是先test 2个step，sub→done +2，更新head指针，就跳到第一步了，因为不能保证上一步中8个step都transmit完了，所以先做一些第一步中的事情。

（4）再进行第一步，每次post 2 个step，就可以转到第三步test 2个step了，等着第一步post又攒够了8个step，正好第一个8个step也test完了，就转到第二步进行transmit了，还是一样，一直transmit 8个step，再转到第三步test 2个step，又转到第一步，依次循环

注意：

当posted = 16，transmitted = 8，done = 8的时候，虽然符合可以继续transmit的条件，但是这个时候proxyState->ncclNet->isend返回值中sub→requests[buffSlot] == NULL，说明不能继续发送了，不会更新sub→transmitted。

因为对端这个时候还没有把buff中的数据收走，需要等待对端收走数据，本端才可以继续发送，并且更新sub->transmitted

（9）recvProxyProgress的流程

recvProxyProgress的作用就是用来在节点间通信的时候，在接收方把通过IB verbs接收到数据，保存到GPU的Send buff上，同时更新tail指针。

下面说一下recvProxyProgress中的三个步骤：

recvProxyProgress总共分为三部分

1、把proxyProgress的数据通过IB verbs post到GPU的buff中（post是接收过程中的一个步骤）

判断条件：

if (sub->posted < sub->nsteps)

proxyState->ncclNet->irecv，并且更新sub->posted

2、把上一步irecv收到的数据，用test来测试一下是否真正接收结束，并且flush

判断条件：

if (subGroup->posted > subGroup->received)

proxyState->ncclNet->test，并且更新sub->received，然后proxyState->ncclNet→iflush

为什么在接收端需要flush：

NCCL GDR Flush Operation · Issue #683 · NVIDIA/nccl · GitHub

3、test一下flush是否结束，然后更新sub->transmitted和recvTail指针

判断条件：

if (subGroup->received > subGroup->transmitted)

proxyState->ncclNet->test，并且更新sub->transmitted和recvTail指针，告诉内核函数，数据已经产生，你可以去消费了

4、更新sub→done，表示完成了从IB上收数据到GPU中

判断条件：

if (sub->transmitted > sub->done)

更新sub->done，需要判断sendHead是否已经处理完了从IB上收到的数据，真正处理完了才更新sub->done

举例说一下recvProxyProgress的一个流程。

（1）进行第一步，把IB上收到的数据，post到GPU的recv buff数据中，这个时候是每次post 2 个step，sub→posted每次+2，每个channel轮流处理，要凑够8个step，才能继续下一步，这一步就是先攒够8个step，够了之后到下一步

（2）进行第二步，把GPU的recv buff数据通过IB verbs（proxyState->ncclNet→test）test一下，也是每次test2个step，sub→received每次+2，然后proxyState->ncclNet->iflush，每个channel轮流处理，要凑够8个step，才能继续下一步，这一步就是先攒够8个step，够了之后到下一步

（3）进行第三步，需要对上一步中flush的数据确认真正完成了（proxyState→ncclNet→iflush），这一步是先test 2个step，sub→transmitted+2，更新recvTail指针

（4）进行第四步，通过sendHead指针去判断当前GPU的recv buff的数据，有么有被发送到下一个GPU的recv buff上。每次done2 个step，又转到第一步，依次循环

（10）proxy connection建立的时机

ncclProxyConnect是用来在rank和proxyService建立connection

对于节点内：

P2P的transport的情况下，本rank会沿着Ring中channel指定的ring的顺序，或者TREE中channel指定的chain的顺序，与自己的send peer rank和recv peer rank之间建立proxyService的connection

shm的transport的情况下，本rank只会在ncclParamShmUseCudaMemcpy的情况下，跟peer的proxyService建立connection

net的transport的情况下，本rank会本rank的proxyService建立connection，用来申请共享内存用的（ncclProxyMsgSharedInit）

对于节点间：

net的transport的情况下，本rank会和send peer rank，recv peer rank的proxyService之间建立connection（描述有问题）

（应该这样说）对于节点间需要建立TRANSPORT_NET的情况下，本rank先找到正确的网卡（ncclTopoGetNetDev）：如果不支持PXN的话，那么proxyRank就是本rank，否则就是离这个网卡最近的rank

自己rank和管理网卡的那个rank"建立proxyservice的connection"，然后自己rank和管理网卡的rank再建立NET transport。

CollNet的transport的情况下，本rank会和send peer rank，recv peer rank的proxyService之间建立connection

（11）几个问题（持续补充）

1、为什么在inter-node接收端需要GPU fulsh一下？

https://zhuanlan.zhihu.com/p/685361884

2、一个proxyProgress线程处理有所channel，如何保证实时性？

RDMA操作不需要太多的CPU，所以一个CPU线程可以处理所有channel

3、inter-node中，发送端的Proxy和接收端的Proxy分别是什么作用？怎么和kernel配合的？

在发送端的Proxy，要接收来自本rank上kernel产生的数据，kernel是生产者，发送端Proxy是消费者；

在接收端的Proxy，要发数据给本rank上的kernel，发送端Proxy是生产者，kernel是消费者。

生产者更新tail指针，通知消费者可以接收数据了；消费者更新head指针，通知生产者可以发送数据了。

4、send侧支持PXN的时候，send的rank跟哪个NET之间建立net transport？

参考链接：

Some questions about proxy and kernel in nccl · Issue #966 · NVIDIA/nccl · GitHub

NCCL GDR Flush Operation · Issue #683 · NVIDIA/nccl · GitHub

NVIDIA NCCL 源码学习（十）- 多机间ncclSend和ncclRecv的过程_nccl编程-CSDN博客

NCCL源码解析: proxy 线程_nccl中的proxy指什么-CSDN博客