Leoda

nccl4py

2 min read

1. Intro

1.1 call chain

Python: nccl_comm.reduce(data, data, nccl.SUM)
    ↓
Python: communicator.py::reduce() 
    ↓ (root=None → all_reduce)
Cython: nccl.pyx::all_reduce()
    ↓
Cython: cynccl.pyx::ncclAllReduce()
    ↓
Cython: _internal/nccl_linux.pyx::_ncclAllReduce()
    ↓ (通过 dlsym 获取函数指针)
C 库: libnccl.so::ncclAllReduce()
    ↓
GPU 通信执行

从 Python 调用开始，经过 Python 包装 → Cython 绑定 → 动态库加载 → NCCL C 库，最终在 GPU 上执行通信操作。

1.2 compile and test

  • step1. VCCL/nccl4py路径下直接编译(主 node 执行就可以完成 nccl4py的编译)：
export CUDA_HOME=/usr/local/cuda
python setup.py build_ext --inplace
  • step2. 手动pip download了nccl4py需要的 python 依赖，比如 cuda.core 等来应对集群没有网的 case。主要包括一下四个： image.png 写一个 requirements.txt 一键 pip install 这些 whl 轮子。
# requirements.txt
packaging==24.2
mpi4py
cuda.core

安装（每个 node 上执行）：

cd nccl4py
pip install -r requirements.txt --no-index --find-links=./third/
  • step3. 完成安装后直接就可以使用 vccl alltoallv 简单测试03_alltoallv.py的功能：
export PYTHONPATH=/inspire/hdd/global_user/huxiaohe-p-huxiaohe/liuda/a2av/nccl4py/build:$PYTHONPATH
export LD_LIBRARY_PATH=/inspire/hdd/global_user/huxiaohe-p-huxiaohe/liuda/a2av/build/lib:$LD_LIBRARY_PATH
mpirun -np 4 \
        --allow-run-as-root \
        -x LD_LIBRARY_PATH=/workspace/liuda/iw/VCCL/build/lib:$LD_LIBRARY_PATH \
        python examples/01_basic/03_alltoallv.py

发包

2. VCCL AlltoallV 4py

目前C++接口为:

ncclResult_t ncclAlltoAllv(const void* sendbuff, const size_t* sendcounts,
    const size_t* sdispls, void* recvbuff, const size_t* recvcounts, const size_t* rdispls,
    const void* relaybuff, ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);
ncclResult_t pncclAlltoAllv(const void* sendbuff, const size_t* sendcounts,
    const size_t* sdispls, void* recvbuff, const size_t* recvcounts, const size_t* rdispls,
    const void* relaybuff, ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);

python接口为：

def alltoallv(
        self,
        sendbuf: NcclBufferSpec,
        recvbuf: NcclBufferSpec,
        sendcounts: Sequence[int],
        sdispls: Sequence[int],
        recvcounts: Sequence[int],
        rdispls: Sequence[int],
        relaybuf: NcclBufferSpec | None = None,
        *,
        stream: NcclStreamSpec | None = None,
    ) -> None:

这里count和displs都是nRanks2的长度，每个rank能找到自己发给目的rank的长度和起始地址

3. test example

vccl alltoallv的测试脚本路径为： VCCL/nccl4py/examples/01_basic/03_alltoallv.py

Graph View

Backlinks