夜雨聆风学习资料网

ARTICLE · 1073512

TensorRT-LLM 0.5.0 源码之四十七

TensorRT-LLM 0.5.0 源码之四十七

NcclCommunicator

struct ncclComm;typedef struct ncclComm* ncclComm_t;
class NcclCommunicator{public:    template <typename T>void send(T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const;    template <typename T>void send(IBuffer const& buf, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const{        send(bufferCast<T>(buf), buf.getSize(), peer, stream, logger);    }    template <typename T>void receive(T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const;    template <typename T>void receive(IBuffer& buf, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const{        receive(bufferCast<T>(buf), buf.getSize(), peer, stream, logger);    }static std::shared_ptr<NcclCommunicator> createPipelineComm(        WorldConfig const& worldConfig, nvinfer1::ILogger& logger);private:    ncclComm_t mComm;};
#if ENABLE_MULTI_DEVICE#include <nccl.h>#endif // ENABLE_MULTI_DEVICE
namespace{#if ENABLE_MULTI_DEVICE//! \brief For converting a C++ data type to a Nccl data type.template <typename T>struct NcclDataType{};template <>struct NcclDataType<half>{    static constexpr auto value = ncclDataType_t::ncclHalf;};template <>struct NcclDataType<float>{    static constexpr auto value = ncclDataType_t::ncclFloat;};template <>struct NcclDataType<std::uint8_t>{    static constexpr auto value = ncclDataType_t::ncclUint8;};template <>struct NcclDataType<std::int32_t>{    static constexpr auto value = ncclDataType_t::ncclInt32;};#endif // ENABLE_MULTI_DEVICE} // namespace

send

template <typename T>void NcclCommunicator::send(    T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const{#if ENABLE_MULTI_DEVICE    auto datatype = NcclDataType<std::remove_cv_t<T>>::value;    TLLM_NCCL_CHECK(ncclSend(sendbuff, count, datatype, peer, mComm, stream.get()), logger);#else    TLLM_THROW("Multi device support is disabled.");#endif // ENABLE_MULTI_DEVICE}template void NcclCommunicator::send(std::uint8_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const;template void NcclCommunicator::send(std::int32_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const;template void NcclCommunicator::send(std::uint8_t const*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const;template void NcclCommunicator::send(std::int32_t const*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const;

receive

template <typename T>void NcclCommunicator::receive(    T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const{#if ENABLE_MULTI_DEVICE    auto datatype = NcclDataType<std::remove_cv_t<T>>::value;    TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, datatype, peer, mComm, stream.get()), logger);#else    TLLM_THROW("Multi device support is disabled.");#endif // ENABLE_MULTI_DEVICE}template void NcclCommunicator::receive(std::uint8_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const;template void NcclCommunicator::receive(std::int32_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const;

createPipelineComm

std::shared_ptr<NcclCommunicator> NcclCommunicator::createPipelineComm(    WorldConfig const& worldConfig, nvinfer1::ILogger& logger){#if ENABLE_MULTI_DEVICE    int const myRank = worldConfig.getRank();    int const worldSize = worldConfig.getSize();    ncclUniqueId id;    if (myRank == 0)    {        ncclGetUniqueId(&id);        for (auto peer = 1; peer < worldSize; ++peer)        {            TLLM_MPI_CHECK(MPI_Send(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD), logger);        }    }    else    {        auto constexpr peer = 0;        MPI_Status status;        TLLM_MPI_CHECK(MPI_Recv(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD, &status), logger);    }    auto pipelineComm = std::make_shared<NcclCommunicator>();    TLLM_NCCL_CHECK(ncclCommInitRank(&pipelineComm->mComm, worldSize, id, myRank), logger);    return pipelineComm;#else    // Python runtime requires instantiation of a communicator even though it may never be used to enable    // pipeline parallel code-path. To enable this, have an empty communicator with uninitialized state.    return nullptr;#endif // ENABLE_MULTI_DEVICE}

这是一个用于创建管道并行通信器的函数,主要使用NCCL(NVIDIA集合通信库)和MPI(消息传递接口)来实现多GPU/多节点间的通信初始化。下面我将详细解释代码的每个部分。

函数整体功能

该函数createPipelineComm的核心功能是为管道并行训练模式创建一个NCCL通信器。它通过MPI进行进程间协调,确保所有参与训练的进程能够加入到同一个NCCL通信域中,为后续的集合通信操作(如All-Reduce、Broadcast等)奠定基础。

代码逐段解释

  1. 1. 条件编译与基本变量获取
    #if ENABLE_MULTI_DEVICEint const myRank = worldConfig.getRank();int const worldSize = worldConfig.getSize();
    • • #if ENABLE_MULTI_DEVICE 表明该函数仅在启用多设备(分布式)模式时编译和执行。
    • • myRank 获取当前进程的排名(Rank),用于唯一标识一个进程。
    • • worldSize 获取参与通信的进程总数。
  2. 2. 生成并广播NCCL唯一ID
    ncclUniqueId id;if (myRank == 0){    ncclGetUniqueId(&id);    for (auto peer = 1; peer < worldSize; ++peer)    {        TLLM_MPI_CHECK(MPI_Send(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD), logger);    }}else{    auto constexpr peer = 0;    MPI_Status status;    TLLM_MPI_CHECK(MPI_Recv(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD, &status), logger);}
    • • ncclGetUniqueId(&id):由Rank 0进程调用,生成一个唯一的标识符(ID)。这个ID用于确保所有进程初始化的NCCL通信器属于同一个通信域。
    • • 广播过程:Rank 0进程作为根节点(root),通过MPI的MPI_Send函数将生成的ID发送给所有其他进程(Rank 1 到 Rank worldSize-1)。其他进程则通过MPI_Recv函数从Rank 0接收这个ID。这个过程确保了所有进程都拿到了相同的ncclUniqueId。
  3. 3. 初始化NCCL通信器
    auto pipelineComm = std::make_shared<NcclCommunicator>();TLLM_NCCL_CHECK(ncclCommInitRank(&pipelineComm->mComm, worldSize, id, myRank), logger);
    • • 创建通信器对象:首先创建一个NcclCommunicator对象的智能指针。
    • • ncclCommInitRank:这是关键函数,它使用之前广播的id、当前进程的myRank和总的进程数worldSize来初始化一个NCCL通信器(pipelineComm->mComm)。每个进程(Rank)都会调用此函数,从而加入到这个共同的通信域中。
  4. 4. 返回通信器
    return pipelineComm;
    • • 函数最终返回初始化好的NCCL通信器的智能指针。
  5. 5. 单机模式处理
    #else// Python runtime requires instantiation of a communicator even though it may never be used to enable// pipeline parallel code-path. To enable this, have an empty communicator with uninitialized state.return nullptr;#endif // ENABLE_MULTI_DEVICE
    • • 当没有启用多设备支持(ENABLE_MULTI_DEVICE为假)时,例如在单机单卡环境下,函数返回nullptr。注释说明这是为了兼容某些可能实例化通信器但实际并不使用它的运行环境(如Python运行时)。

关键协作:MPI 与 NCCL

这段代码体现了MPI和NCCL在分布式训练中的典型协作模式:

  • • MPI 负责控制面通信:在这里,MPI负责进程间的协调工作,即可靠地将NCCL的唯一ID从根节点(Rank 0)广播到所有工作节点,确保大家获得一致的初始化信息。
  • • NCCL 负责数据面通信:当所有进程获得统一的ID后,NCCL库负责建立高效的、基于GPU的通信路径(如通过NVLink、InfiniBand等),并进行实际的GPU间数据交换(集合通信)。

总结

总而言之,createPipelineComm函数完成了一次分布式的NCCL通信器初始化。它通过MPI协调,让所有进程基于一个共同的ID创建NCCL通信域,为后续在管道并行等模型并行策略中进行高效的GPU间通信做好了准备。

参考文献

  • • https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/runtime/ncclCommunicator.h
点个「赞」+「在看」❤️
让我们知道这份文字有温暖到你,也是我们持续创作的最大动力!
推荐
TensorRT-LLM 0.5.0 源码之四十六
SGLang-Omni 通俗解析(六):新 TTS 车间入驻指南与血泪避坑手册
DALI Audio Resample
把"炼丹炉"直播给你看:小米 MiMo-V2.6-Pro/Flash 强化学习训练全景解读
全双工语音模型究竟为语音AI带来了哪些变革
SGLang-Omni 通俗解析(五):厂内物流与通信——对讲机只管喊话,物流只管搬货
Moshi:面向实时对话的语音‑文本基础模型
线程别名
Claude Fable 5:自改进智能体——14步循环工程指南
MOSS Transcribe Diarize 技术报告
SGLang-Omni 通俗解析(四):流水线"四大金刚"与反馈循环黑科技
京东大溶洞
SGLang-Omni 通俗解析(三):一个请求的奇幻漂流——架构全景
SGLang-Omni 通俗解析(二):进程命名潜规则与5个高级机制
Qwen3-Omni 技术报告
SGLang-Omni 通俗解析(一):用"开工厂"的思路理解多模态配置
权限、沙箱与自主智能体
技能:规定智能体应该如何思考
循环工程:设计编码代理系统,而不是每次都手动提示
VoxCPM2 技术报告
Agent 数据生产与训练
什么是循环工程?AI 编码智能体的新范式
从空文件夹到生成图表:Claude Code 实战教程
Claude Code 入门:研究者配置指南
GLM-5.2:面向长时序任务打造
循环工程(Loop Engineering)
Agent SFT 标准数据格式 + Loss Mask 完整实现
4D Parallelism
TileLang与OpenAI Triton的核心区别
Agent SFT 数据
Claude Code 的上限,就是你的上限
如何用Claude Code提升软件工程工作效率、改善生活
LLM推理优化的核心技术:深入理解KV缓存与分页注意力机制
Qwen3-TTS 技术报告
PagedAttention
如何让AI听懂你的“话外音”?GOAT-SLM模型实现更懂情感的语言交互
FlashAttention与PagedAttention详解:拯救GPU显存,让大模型飞起来的核心技术

相关学习资料