Skip to content

Latest commit

 

History

28 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Mini-NCCL: 高性能轻量级分布式集合通信库

Language Platform License


Mini-NCCL 是一款专为受限算力环境打造的高性能、轻量级分布式集合通信库 。本项目深度复刻了 NVIDIA NCCL 的核心架构与设计哲学 ,基于 RDMA (RoCEv2/InfiniBand) 与 CUDA 底层技术,实现了极高带宽利用率的 Ring All-Reduce 算法 ,并内建了微秒级的分布式控制平面(Hera-Core) 。

为了在物理硬件边界内极限压榨通信性能,本项目深度融合了显存注册缓存(MR Cache)、双缓冲内存池以及 GPU 侧异步轮询等底层调优手段 ,成功在虚拟化(Soft-RoCE)与物理网卡双重环境中实现了极低延迟与高吞吐量 。整体架构全面涵盖了异步流水线 (Asynchronous Pipelining)、窗口化流控、零拷贝资源管理、无锁对象池、工业级异常安全体系以及 C++ 模板元编程等系统级实践 。


🎯 为什么要做 Mini-NCCL?

现状问题

  • 在真实 AI Infra 场景中NVIDIA 的 NCCL:

    • 面向 A100 / H100 + NVLink + IB 集群
    • 黑盒实现,难调试
    • 消费级硬件 + RoCE 场景表现不可控

项目背景与使用场景

存在大量“非顶配算力场景”:中小 AI 团队(4090 / 4060)、高校实验室、初创公司、量化交易低延迟集群以及平民玩家算力场景。

这些环境的特点:无 NVLink、无 IB 交换机、使用以太网 / Soft-RoCE、拓扑复杂、不规则

在缺乏高端显卡和昂贵 InfiniBand 物理网络的受限硬件环境下(如以太网、多显卡直连、Soft-RoCE 模拟),为了最大化榨取普通硬件的分布式计算潜力,降低多节点间的数据同步延迟,从零构建了这个集合通信库。它摒弃了庞大框架的冗杂依赖,通过极致的底层系统级优化,实现了物理硬件的带宽打满。

👉 Mini-NCCL 的意义

在“低成本算力环境”下,通过软件层极致优化逼近高性能通信能力

尽管 NVIDIA 官方的 NCCL 已经是业界的绝对标准,但 Mini-NCCL 的存在具有以下不可替代的价值与特定应用场景:

  1. AI 基础设施的“白盒化”参考实现:官方 NCCL 是一个极其庞大且高度优化的黑盒 。Mini-NCCL 剥离了复杂的硬件兼容性边缘逻辑,提炼了最核心的 Ring All-Reduce、RDMA 通信与 GPU 同步机制,是理解和二次开发底层大模型训练数据流转的工业级参考原型。
  2. 极轻量级的定制化集成场景:在某些定制化的边缘计算集群、国产算力芯片适配,或对二进制体积和第三方依赖极其敏感的场景下,不需要引入动辄数十万行的完整框架。Mini-NCCL 提供了零重型第三方依赖(仅依赖 Verbs 和 CUDA、无 Protobuf)的即插即用方案。
  3. 系统级极限性能工程实践:本项目展示了对系统底层资源的极致榨取能力,涵盖了异步流水线、无锁内存池、零拷贝通信、泛型元编程以及软硬件协同监控等高级系统级编程技术。

🚀 核心特性 (Key Features)

  • RDMA 传输层:基于 InfiniBand Verbs (ibverbs) API 构建的高速通信层,支持 Soft-RoCE (RXE) 及物理 RDMA 网卡。

  • 异步流水线:通过 计算-通信重叠 (Overlap) 技术,利用 CUDA Stream 并发执行,有效掩盖通信延迟。

  • 双缓冲机制:引入 Ping-Pong Buffer (双缓冲) 策略,最大化链路利用率。

  • 热路径零内存分配:设计了自定义 对象池 (Object Pool) 内存管理,消除了关键路径上的动态内存分配 (malloc/new) 开销,显著降低延迟抖动。

  • 泛型计算框架:利用 C++ 模板元编程 (Template Metaprogramming) 和分发器,支持任意数据类型 (FP32, FP64, Int32) 和规约操作 (Sum, Prod, Max, Min)。

  • 标准 API:提供 ABI 兼容的 C 语言标准接口 (ncclComm_t, ncclAllReduce),易于集成到 PyTorch 等框架。

  • 显存注册缓存 (MR Cache) :在 RDMATransport 内部维护 std::unordered_map 缓存,重复利用的通信缓冲区只需在首次调用时注册,后续直接命中缓存返回,消除系统调用导致的毫秒级延迟

  • GPU 侧轮询同步 (GPU-Side Polling) :开发了专门的 CUDA 核函数 (wait_kernel),使 GPU 能够直接在显存中轮询对端写入的 Flag 标志位,彻底摆脱了 CPU 轮询导致的频繁同步,使 GPU 能够真正独立地调度后续规约计算。

  • 窗口化流控 (Window-based Flow Control) : 实现了基于滑动窗口的请求调度 (Window Size = 64),在保证链路饱和的同时,防止请求队列溢出导致的资源耗尽。

  • 选择性信号 (Selective Signaling) : 通过稀疏的 CQ (Completion Queue) 信号机制(每 16 个包一次 Signal),大幅降低 CPU 轮询开销。

  • GPUDirect RDMA Ready: 底层传输层 (RDMATransport) 已实现 Device Pointer 的直接注册逻辑。在支持 nvidia-peermem 的物理集群上可自动激活 零拷贝 (Zero-Copy) 数据通路。

  • 自动拓扑感知: 自动探测节点内拓扑,优先使用 CUDA IPC (P2P) 进行节点内通信,仅在跨节点时使用 RDMA。

  • 柔性降级 (Graceful Degradation) : 在 WSL2 或 Docker 等受限环境中,自动屏蔽不支持的特性(如 P2P),回退到兼容模式,确保 "Run Anywhere"。

  • 异常传播体系: 全面摒弃 exit(1),构建了基于 std::exception 的错误传播链,确保 Library 错误不会导致宿主进程(如 PyTorch)意外崩溃。

  • NUMA 绑定: 自动探测网卡所属 NUMA 节点,并调用 sched_setaffinity 绑定 CPU,消除跨 Socket 访问瓶颈。


📁 项目结构

.
├── src
│   ├── hera           # Hera-Core 控制平面源码 (Master/Worker/Socket)
│   ├── transport      # 通信层 (RDMA Verbs / CUDA IPC 封装)
│   ├── api.cpp        # NCCL 标准 API 实现
│   └── mini_nccl.cu   # 核心 CUDA Kernel 与 Ring 算法实现
├── include
│   └── mini_nccl_api.h # 对外公开头文件
├── tests
│   └── perf_test.cpp  # 性能测试工具 (支持自动组网)
└── demos              # 早期原型验证代码

🏗️ 项目架构设计 (Architectural Design)

Mini-NCCL 采用 Control Plane (控制面)Data Plane (数据面) 分离的设计:

  1. Data Plane (Mini-NCCL) : 负责极致吞吐的数据传输,基于 RDMA Verbs 和 CUDA IPC。
  2. Control Plane (Hera-Core) : 负责自动组网、拓扑管理与故障熔断,基于自定义二进制协议。
[User App]  -->  [ mini_nccl_api ]  -->  [ Context ]
                                            |
                        +-------------------+-------------------+
                        |                                       |
                  [ Algorithm ]                           [ Transport ]
             (Ring AllReduce Logic)                  (Verbs/IPC Wrapper)
                        |                                       |
              +---------+---------+                   +---------+---------+
              |   GPU Kernels     |                   |   RDMA Queue Pair |
              | (Reduce/Scale)    |                   |   CUDA IPC Handle |
              +-------------------+                   +-------------------+

1. 控制平面 (Control Plane - Hera System)

👉 本质:轻量级 MPI 替代

  • Hera-Master: 充当集群编排器。负责监听 Worker 注册,动态分配唯一的 Rank,并维护全局拓扑信息(如广播 Root IP),为 RDMA 建立连接提供引导(Bootstrap)服务。Rank 分配,拓扑发现,Bootstrap 信息分发
  • Hera-Worker: 嵌入在应用进程中的代理。负责与 Master 通信同步状态,并引导传输层完成 Queue Pair (QP) 的握手。注册节点,建立通信上下文,引导 RDMA 建链
  • Fail-Fast 机制:实时监控连接状态,一旦检测到节点崩溃,立即广播全局熔断指令 (Global Abort),防止分布式集群死锁 Hang 死 。

2. 算法层 (Algorithm Layer - Ring Engine)

  • Ring Topology: 实现逻辑环形拓扑。每个节点仅与邻居(Left/Right)通信,将数据分为 个分片(Slice)进行 Scatter-Reduce 和 All-Gather 两个阶段的操作。

  • Pipelining: 将大包拆分为 SLICE_SIZE(默认 128KB),利用双缓冲机制实现计算与传输的重叠。

  • 泛型计算分发器 (Dispatcher) :利用 C++ 模板元编程支持任意数据类型 (FP32, FP64, Int32) 和规约操作 (Sum, Prod, Max, Min) 的编译时分发,并提供 ABI 兼容的 C 语言标准接口 (ncclComm_t),极易集成至 PyTorch 等框架 。

3. 传输层 (Transport Layer - RDMA & UVA)

  • RDMA Transport: 封装 Verbs API。支持单边 RDMA_WRITE 操作,通过无锁队列管理异步请求。
  • UVA (Unified Virtual Addressing) : 结合 cudaHostAllocMapped 标志,实现 Host 与 Device 的虚拟地址统一,支撑零拷贝(Zero-copy)传输。
  • Ring Topology:实现逻辑环形拓扑。数据在 $2(N-1)$ 个分片 (Slice) 的调度下,依次经过 Scatter-Reduce 和 All-Gather 阶段,带宽利用率与节点数无关。

📊 性能基准 (Benchmark)

WSL2 (Ubuntu 24.04) + Soft-RoCE (RXE) 严格受限的模拟环境下,对单节点数据量进行了双进程测试,结果如下:

数据量 (Data Size) 耗时 (Avg Time) 算法带宽 (AlgBW) 校验状态 (Check) 状态
1 MB ~4244 us 0.25 GB/s PASSED Latency Bound
16 MB ~42599 us 0.39 GB/s PASSED Bandwidth Saturation
64 MB ~165358 us 0.41 GB/s PASSED Stable Peak
128 MB ~345252 us 0.39 GB/s PASSED Stable

性能分析:0.41 GB/s(约 3.3 Gbps)是 WSL2 虚拟网络栈的物理极限 。由于缺乏物理网卡卸载,数据路径为:GPU -> PCIe -> CPU Mem -> Linux Kernel (RXE) -> Hyper-V -> Windows Network,涉及至少 4 次 CPU 内存拷贝 。尽管受限于宿主机网络栈,但极度平稳的带宽曲线证明了软件异步流水线与内存池控制极其高效,软件层面没有引入任何额外抖动 。在物理机环境配合 GPUDirect,该架构完全具备跑满 100Gbps+ 物理网卡的扩展能力 。

AVX2 SIMD 硬件级极速校验:为确保 All-Reduce 分布式累加结果的 100% 准确性,项目中引入了 256 位向量寄存器 (__m256) 校验逻辑。对比传统循环,校验速度提升 8倍,处理 128MB 以上数据包耗时缩短至微秒级,几乎零额外开销 。


🛠️ 编译与运行指南

1. 依赖环境

  • 操作系统:Linux (Ubuntu 22.04+ 测试通过)

  • 编译器:GCC 7+ & NVCC (CUDA Toolkit 11+)

  • 依赖库:libibverbs (RDMA-Core), NVTX

  • 网络环境: 配置好的 Soft-RoCE (RXE) 接口 (例如 rxe0)

  • CMake >= 3.10

2. 编译指南

支持 AVX2 硬件加速编译 :

mkdir build && cd build
cmake ..
make -j$(nproc)

3. 运行测试

使用内置的 Benchmark 工具进行双进程模拟测试:单机上模拟 2 个 Rank 的分布式环境

# Terminal 1
./build/perf_test 0 2 127.0.0.1
# Terminal 2
./build/perf_test 1 2 127.0.0.1

预期输出: Result: [PASS] All values are 3.0!

4. 分布式性能测试运行

推荐使用 Hera 自动化组网模式进行测试。

步骤 A:启动 Hera-Master 控制平面

在一个终端指定等待的计算节点数量(如 2 个):

./hera_master 2  # Master 将监听 9999 端口等待 Worker 加入

步骤 B:启动计算节点 (Workers) 在另外两个终端分别执行测试程序,使用 -1 激活自动组网模式 :

# 建议使用 prlimit 提升内存锁定权限以支持 RDMA 注册
sudo prlimit --pid $$ --memlock unlimited:unlimited

# 终端 1 (Rank 0)
./perf_test -1 2 127.0.0.1
# 终端 2 (Rank 1)
./perf_test -1 2 127.0.0.1

预期:Worker 自动连接 Master -> Master 广播拓扑 -> 自动分配 Rank -> 触发 GPU 预热并输出 AVX2 校验结果 [PASS]

说明:Hera 模式下,perf_test 的第三个参数会被当作 Hera Master 的 IP。

5. 安装 (Optional)

将库安装到系统路径,供第三方集成:

sudo make install

⚙️ 核心配置参数 (Environment Variables)

支持通过环境变量实现对系统调优的高级控制:

变量名 默认值 描述说明
MINI_NCCL_NET_DEVICE 自动匹配 指定强制使用的 IB/RoCE 设备名 (如mlx5_0,rxe0)
MINI_NCCL_SLICE_SIZE 131072 RDMA 通信数据流水分片大小 (Bytes)
MINI_NCCL_WINDOW_SIZE 64 滑动窗口大小,控制 RDMA 在途请求的流控上限
MINI_NCCL_SIGNAL_BATCH 16 批量生成 RDMA 完成信号的间隔数,优化 CQ 轮询压力

示例:

export MINI_NCCL_NET_DEVICE=mlx5_0
export MINI_NCCL_SLICE_SIZE=$((256*1024))
export MINI_NCCL_WINDOW_SIZE=32
export MINI_NCCL_SIGNAL_BATCH=8

📄 代码结构

  • include/mini_nccl_api.h: 公开的 C API 头文件 (ABI 兼容)。
  • src/api.cpp: API 实现与类型分发逻辑。
  • src/mini_nccl.cu: 泛型 Ring All-Reduce 算法与 CUDA 核函数。
  • src/transport/RDMATransport.h: 底层 Verbs 封装与 内存池 (Memory Pool) 实现。
  • src/main.cpp: 端到端集成测试代码。

📜 License

MIT License.

Copyright (c) 2026 宫旭东 (Gong Xudong)

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

About

No description, website, or topics provided.

Resources

Stars

8 stars

Watchers

2 watching

Forks

Releases

Packages

Contributors

Languages