第2章 NVSHMEM 环境搭建

2、NVSHMEM 环境搭建：硬件要求、软件栈与第一个程序

好，咱们直接进入正题。环境搭建这事儿，说难不难，说简单吧，坑也不少。我刚开始搞 NVSHMEM 的时候，光装库就折腾了大半天，后来发现是 CUDA 版本没对齐。嗯，咱们一步步来。

2.1 硬件要求：NVLink 和 NVSwitch

先聊聊硬件。NVSHMEM 之所以快，说白了就是靠 NVLink 和 NVSwitch 这两样东西。没有它们，你跑 NVSHMEM 就跟开跑车在泥巴路上飙车一样——白搭。

NVLink：这是 GPU 之间的高速互联通道。我习惯叫它「GPU 之间的高速公路」。单条 NVLink 带宽能到 300 GB/s 以上（看具体代际），比 PCIe 4.0 x16 快好几倍。
NVSwitch：这东西像个超级交换机，能把多个 GPU 全互联起来。比如 DGX 系列里，8 张 GPU 通过 NVSwitch 组成一个全互联拓扑，任意两张卡之间都能直接走 NVLink。

核心要求：你的机器里至少要有 2 张支持 NVLink 的 GPU（比如 V100、A100、H100 或更新的型号），并且它们之间通过 NVLink 或 NVSwitch 物理连接。单卡环境也能装 NVSHMEM，但说白了只能跑单进程，体验不到真正的优势。

我在项目中遇到过一台机器，明明有 4 张 A100，结果 NVLink 没接好，跑 NVSHMEM 程序时通信延迟高得离谱。查了半天才发现是物理链路没配置对。所以，装之前先跑个 nvidia-smi topo -m 看看 GPU 之间的拓扑关系，确认 NVLink 是连通的。

2.2 软件栈：CUDA 与 NVSHMEM 库

软件部分其实不复杂，就两样东西：CUDA 工具包和 NVSHMEM 库。但版本匹配是个大坑，我吃过亏。

组件	推荐版本	备注
CUDA Toolkit	11.0 及以上	我建议用 11.8 或 12.x，稳定性好
NVSHMEM	2.x 或 3.x	跟 CUDA 版本要对应，别乱配
GPU 驱动	最新稳定版	驱动太老会导致 NVLink 识别不了

避坑指南：我曾经因为 CUDA 11.4 配了 NVSHMEM 2.5，结果编译时一堆符号找不到。后来查文档才发现，NVSHMEM 2.5 要求 CUDA 11.6 以上。所以，装之前一定去 NVIDIA 官网看兼容性矩阵，别凭感觉来。

安装步骤其实就三步：

装 CUDA Toolkit（记得选 runfile 或 deb 包，别用系统包管理器里的老版本）
下载 NVSHMEM 源码或预编译包（我习惯用源码编译，灵活一些）
设置环境变量 NVSHMEM_HOME 和 LD_LIBRARY_PATH

你想想看，如果环境变量没设对，编译时链接器找不到库，那错误信息能把你绕晕。我建议在 .bashrc 里加这么几行：

export CUDA_HOME=/usr/local/cuda-12.0
export NVSHMEM_HOME=/path/to/nvshmem
export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$CUDA_HOME/lib64:$LD_LIBRARY_PATH
export PATH=$CUDA_HOME/bin:$PATH

2.3 编译与运行第一个 NVSHMEM 程序

好，环境搭好了，咱们写个最简单的程序试试水。这个程序就做一件事：让两个 GPU 互相打个招呼——PE 0 发一个整数给 PE 1，PE 1 接收并打印出来。

先看代码：

#include <nvshmem.h>
#include <stdio.h>

int main(void) {
    int my_pe, npes;
    int *buf;
    int msg = 42;

    nvshmem_init();
    my_pe = nvshmem_my_pe();
    npes = nvshmem_n_pes();

    // 每个 PE 分配一个 int 大小的共享内存
    buf = (int*)nvshmem_malloc(sizeof(int));
    if (buf == NULL) {
        fprintf(stderr, "PE %d: nvshmem_malloc failed\n", my_pe);
        nvshmem_finalize();
        return 1;
    }

    if (my_pe == 0) {
        // PE 0 把 msg 写到 PE 1 的 buf 里
        nvshmem_int_put(buf, &msg, 1, 1);
        printf("PE %d: sent %d to PE 1\n", my_pe, msg);
    }

    // 等所有 PE 都完成 put 操作
    nvshmem_barrier_all();

    if (my_pe == 1) {
        printf("PE %d: received %d from PE 0\n", my_pe, buf[0]);
    }

    nvshmem_free(buf);
    nvshmem_finalize();
    return 0;
}

编译命令很简单：

nvcc -o hello_nvshmem hello_nvshmem.cu \
    -I$NVSHMEM_HOME/include -L$NVSHMEM_HOME/lib -lnvshmem

运行的时候，要用 mpirun 或者 nvshmrun 来启动多进程。我个人习惯用 nvshmrun，因为它对 NVSHMEM 的支持更直接：

nvshmrun -n 2 ./hello_nvshmem

如果一切顺利，你会看到类似这样的输出：

PE 0: sent 42 to PE 1
PE 1: received 42 from PE 0

小技巧：第一次跑的时候，如果报错说「找不到设备」，八成是 nvshmrun 没找到你的 GPU。试试加个 --gpu-per-node 2 参数，显式指定每节点用几张卡。

这里有个细节要注意：nvshmem_barrier_all() 是全局同步点。没有它，PE 1 可能在 PE 0 还没写完数据时就读取了，那结果就是乱码。我在调试一个多 GPU 程序时，就因为少了个 barrier，数据对不上，查了整整一下午。

2.4 知识体系总览

下面这张图把本章的核心逻辑串起来了，你一看就明白：

嗯，到这步，你的 NVSHMEM 环境就算搭好了。第一个程序跑通之后，你会觉得「原来就这么回事」。但别急，后面还有原子操作、同步机制这些硬骨头等着咱们。慢慢来，先把基础打牢。

📚 课程目录

2、NVSHMEM 环境搭建：硬件要求、软件栈与第一个程序

2.1 硬件要求：NVLink 和 NVSwitch

2.2 软件栈：CUDA 与 NVSHMEM 库

2.3 编译与运行第一个 NVSHMEM 程序

2.4 知识体系总览