1. 分布式训练概述：从单卡到多卡，为什么需要NVSHMEM和NCCL？

大家好，我是你们的老朋友。今天咱们聊聊分布式训练。说实话，我刚入行那会儿，觉得单卡训练挺香的——一张V100怼上去，小模型跑得飞起。但后来模型越来越大，单卡根本塞不下，我才意识到：分布式训练不是选择题，而是必答题。

1.1 单卡训练的瓶颈：为什么我们不得不走向多卡？

先说说单卡训练的痛点。你想想看，一张显卡的显存是有限的。比如NVIDIA A100是80GB，H100是80GB。但现在的模型呢？GPT-3有1750亿参数，光参数就要700GB。单卡根本装不下。

我遇到过最尴尬的一次：客户要训练一个百亿参数的推荐模型，我一看单卡显存，连模型参数都放不下，更别说优化器状态和梯度了。嗯，这时候你只有两条路：要么换更大的卡（但成本爆炸），要么上多卡分布式训练。

单卡还有另一个问题：算力天花板。一张卡的计算能力是有限的，训练一个千亿模型可能要跑几个月。这谁受得了？

1.2 多卡训练的挑战：通信成为新瓶颈

好，既然决定上多卡，那问题来了：多卡之间怎么协同？

最简单的想法：每张卡存一部分模型参数，各自算各自的梯度，然后汇总更新。但这里有个关键问题——通信。

我刚开始做分布式训练时，天真地以为多卡就是简单地把数据分一分。结果一跑起来，发现大部分时间都花在卡间通信上了。GPU在算，但更多时候在等数据传过来。说白了，通信成了新的瓶颈。

为什么会这样？因为GPU计算速度极快（每秒万亿次浮点运算），但卡间通信带宽有限。比如PCIe 4.0 x16的理论带宽是32GB/s，而GPU内部显存带宽是2TB/s以上。差了两个数量级！

所以，分布式训练的核心问题变成了：如何高效地在多卡之间传输数据和梯度？

1.3 NCCL：GPU通信的“高速公路”

NCCL（NVIDIA Collective Communications Library）就是来解决这个问题的。它是NVIDIA官方推出的GPU间通信库，专门优化了多卡之间的集体通信操作。

我个人习惯把NCCL比作“GPU通信的高速公路”。它提供了几个核心操作：

• AllReduce：所有卡的数据规约后广播给所有卡。这是数据并行训练中最常用的操作。
• AllGather：收集所有卡的数据，然后广播给所有卡。
• ReduceScatter：先规约，再分散到各卡。
• Broadcast：一张卡的数据广播给所有卡。

NCCL厉害在哪？它利用了NVLink、InfiniBand等高速互联技术，实现了接近硬件极限的通信带宽。我记得有一次做测试，用NCCL做AllReduce，8张A100之间的通信带宽能达到600GB/s，几乎是线性的。

1.4 NVSHMEM：共享内存的“魔法”

NCCL解决了集体通信的问题，但有些场景需要更灵活的通信模式。比如，在模型并行或流水线并行中，不同卡之间可能需要点对点通信，或者需要访问其他卡上的数据。

这时候NVSHMEM就派上用场了。NVSHMEM是基于OpenSHMEM标准的GPU间通信库，它提供了一种全局共享内存的抽象。

说白了，NVSHMEM让你感觉所有GPU的显存是连在一起的，你可以像访问本地内存一样访问其他卡上的数据。这听起来很魔幻，但确实能做到。

我举个例子：在模型并行中，模型的不同层分布在不同的GPU上。前向传播时，前一层的输出需要传给下一层。如果用NCCL，你需要显式地做一次点对点通信。但用NVSHMEM，你可以直接通过shmem_put或shmem_get操作，把数据写到目标GPU的显存里，就像写本地内存一样。

// NVSHMEM 示例：将数据从 GPU 0 发送到 GPU 1
int my_pe = shmem_my_pe();  // 当前 GPU 编号
int target_pe = 1;          // 目标 GPU 编号

float *src_buf = (float*)shmem_malloc(N * sizeof(float));
float *dst_buf = (float*)shmem_malloc(N * sizeof(float));

// 在 GPU 0 上填充数据
if (my_pe == 0) {
    for (int i = 0; i < N; i++) src_buf[i] = i * 1.0f;
    // 将数据从 GPU 0 发送到 GPU 1
    shmem_float_put(dst_buf, src_buf, N, target_pe);
    shmem_quiet();  // 等待操作完成
}

NVSHMEM的优势在于：低延迟、高带宽、编程模型简单。它特别适合那些需要频繁、细粒度通信的场景，比如动态稀疏训练、图神经网络等。

1.5 NCCL vs NVSHMEM：什么时候用哪个？

很多同学问我：NCCL和NVSHMEM到底选哪个？我的回答是：不是二选一，而是协同使用。

特性	NCCL	NVSHMEM
通信模式	集体通信（AllReduce等）	点对点、全局共享内存
适用场景	数据并行、梯度同步	模型并行、流水线并行、动态通信
编程复杂度	低（API简单）	中等（需要理解共享内存模型）
性能特点	集体通信优化极好	点对点通信延迟低
典型应用	PyTorch DDP、Horovod	自定义并行策略、GNN训练

在实际项目中，我通常的做法是：用NCCL做梯度同步，用NVSHMEM做模型分片和点对点通信。两者配合，才能发挥多卡训练的最大效能。

1.6 知识体系总览

为了让大家更直观地理解本章的知识结构，我画了一张图：

1.7 小结

好了，这一章我们聊了：

• 单卡训练的瓶颈：显存和算力不够用
• 多卡训练的挑战：通信成为新瓶颈
• NCCL：集体通信的利器，适合数据并行
• NVSHMEM：共享内存的魔法，适合模型并行
• 两者协同：NCCL做梯度同步，NVSHMEM做模型分片

说白了，分布式训练的核心就是算和传的平衡。算得快，传得也要快。NCCL和NVSHMEM就是帮你把“传”这件事做到极致的工具。

下一章，我们会深入NCCL的底层原理，看看它到底是怎么做到那么快的。嗯，到时候见。

---