第4章初始化与终结

4. 初始化与终结：nvshmem_init、nvshmem_finalize、PE查询函数

好，咱们正式开始写代码前的第一件事——把环境搭起来，把程序跑通。

NVSHMEM 的初始化与终结，说白了就是两件事：让所有 GPU 互相认识，以及用完以后体面地退出。我见过不少新手，上来就调 nvshmem_put，结果程序直接崩了——原因很简单，根本没初始化。

4.1 nvshmem_init：一切通信的起点

这个函数没有参数，没有返回值。你只需要在程序最开头调用一次：

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

int main() {
    nvshmem_init();  // 所有 PE 从这里开始同步

    // ... 后续通信代码 ...

    nvshmem_finalize();
    return 0;
}

调用 nvshmem_init() 之后，NVSHMEM 内部会做几件事：

初始化 MPI 或 UCX 底层通信层
分配内部管理用的内存和队列
让所有参与的 PE 进入一个全局同步屏障

嗯，这里要注意：所有 PE 必须同时调用 nvshmem_init。如果你有 4 块 GPU，其中一块没调，那整个程序就会卡死在初始化阶段。我在项目中遇到过这种情况——调试了半天，最后发现是某块卡的 MPI 环境没配好。

⚠️ 避坑指南
我曾经在集群上跑 64 卡的任务，结果初始化阶段就报错。排查下来，是其中一台机器的 CUDA_VISIBLE_DEVICES 没设置对，导致 NVSHMEM 找不到 GPU。记住：初始化前，先确认每块卡都能被正常识别。

4.2 nvshmem_finalize：优雅地退出

有始有终。程序结束前，必须调用 nvshmem_finalize()。它会：

释放所有 NVSHMEM 分配的内部资源
断开 PE 之间的连接
等待所有未完成的通信操作完成

我个人习惯把 nvshmem_finalize 放在程序的最后，就像关水龙头一样自然。如果你忘了调，程序退出时可能会报一些奇怪的错误——比如 CUDA 上下文没释放干净。

// 正确的生命周期
nvshmem_init();
// ... 你的所有通信逻辑 ...
nvshmem_finalize();

💡 小技巧
如果你在写一个长期运行的服务，比如模型训练循环，建议把 nvshmem_finalize 放在异常处理的 finally 块里。这样即使中间报错，也能保证资源释放。

4.3 PE 查询函数：你是谁？我在哪？

初始化之后，每个进程（或者说每个 GPU）都需要知道两件事：

我是谁？——我的 PE ID 是多少
有多少人？——总共有多少个 PE

NVSHMEM 提供了两个最常用的查询函数：

函数	返回值	说明
`nvshmem_my_pe()`	int	当前 PE 的编号，从 0 开始
`nvshmem_n_pes()`	int	总共有多少个 PE

这两个函数在初始化之后随时可以调用。我通常会在程序开头打印一下，确认环境：

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

printf("Hello from PE %d / %d\n", my_pe, npes);

你想想看，如果总共有 8 个 PE，但你的代码里写死了 nvshmem_put 的目标是 PE 7，结果实际只启动了 4 个——那程序肯定跑飞。所以，永远不要硬编码 PE 数量，用 nvshmem_n_pes() 动态获取。

4.4 完整的初始化模板

下面是我个人常用的模板，你可以直接拿来用：

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

int main() {
    // 1. 初始化 NVSHMEM
    nvshmem_init();

    // 2. 获取 PE 信息
    int my_pe = nvshmem_my_pe();
    int npes  = nvshmem_n_pes();

    // 3. 确认 GPU 可用
    int dev_count;
    cudaGetDeviceCount(&dev_count);
    printf("PE %d: GPU count = %d\n", my_pe, dev_count);

    // 4. 设置当前设备（每个 PE 绑定一块 GPU）
    cudaSetDevice(my_pe % dev_count);

    // 5. 分配 NVSHMEM 对称内存
    // ... 后续章节会讲 ...

    // 6. 终结
    nvshmem_finalize();
    return 0;
}

    🔑 核心要点

    nvshmem_init 和 nvshmem_finalize 必须成对出现
所有 PE 必须同时调用 nvshmem_init
用 nvshmem_my_pe() 和 nvshmem_n_pes() 动态获取 PE 信息
每个 PE 通常绑定一块 GPU，用 cudaSetDevice 设置

4.5 知识体系图

下面这张图展示了初始化与终结的完整流程，以及 PE 查询函数在整个生命周期中的位置：

说白了，初始化就是「开门」，终结就是「关门」，PE 查询就是「看看屋里都有谁」。这三件事搞定了，后面的通信操作才能稳稳当当。

📚 课程目录

4. 初始化与终结：nvshmem_init、nvshmem_finalize、PE查询函数

4.1 nvshmem_init：一切通信的起点

4.2 nvshmem_finalize：优雅地退出

4.3 PE 查询函数：你是谁？我在哪？

4.4 完整的初始化模板

4.5 知识体系图