第2章：单机多卡编程基础

好，咱们进入正题。单机多卡，说白了就是一台机器上插了好几块GPU，怎么让它们协同干活。这章我会从CUDA编程模型讲起，再带你看看PyTorch的分布式包怎么用。嗯，都是基本功，但很重要。

2.1 CUDA编程模型回顾

先聊聊CUDA。很多人觉得CUDA就是写核函数，其实没那么玄乎。我刚开始接触时也犯过迷糊——以为把数据扔到GPU上就完事了。结果呢？性能还不如CPU。

CUDA的核心思想是啥？异构计算。CPU负责控制流，GPU负责并行计算。你想想看，CPU就像项目经理，GPU就像一群搬砖工。项目经理分派任务，搬砖工们一起干。

关键概念：

Host：CPU及其内存
Device：GPU及其显存
Kernel：在GPU上执行的函数
Grid & Block：线程的组织结构

写CUDA代码时，我习惯先画个图。线程是怎么组织的？Grid里有Block，Block里有Thread。每个Thread执行同样的Kernel，但处理不同的数据。这就是SIMT（单指令多线程）模型。

// 一个简单的CUDA向量加法
__global__ void vecAdd(float* A, float* B, float* C, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        C[i] = A[i] + B[i];
    }
}

int main() {
    // 分配显存
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, N * sizeof(float));
    cudaMalloc(&d_B, N * sizeof(float));
    cudaMalloc(&d_C, N * sizeof(float));
    
    // 拷贝数据到GPU
    cudaMemcpy(d_A, h_A, N * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, N * sizeof(float), cudaMemcpyHostToDevice);
    
    // 启动Kernel
    vecAdd<<<(N+255)/256, 256>>>(d_A, d_B, d_C, N);
    
    // 拷贝结果回CPU
    cudaMemcpy(h_C, d_C, N * sizeof(float), cudaMemcpyDeviceToHost);
}

这段代码看着简单，但坑不少。我曾经犯过一个错——忘了检查cudaMalloc的返回值。结果显存不够时，程序直接崩了。后来我养成了习惯：每次CUDA调用后都检查错误。

避坑指南：

我曾经在生产环境遇到过一个问题：多卡训练时，某张卡的显存突然爆了。排查了半天，发现是某个Kernel的block size设得太大，导致寄存器溢出。记住：block size不是越大越好，要根据GPU的SM数量和寄存器资源来调。

2.2 PyTorch分布式包简介

好了，CUDA基础有了，咱们看看PyTorch怎么玩分布式。PyTorch的分布式包叫torch.distributed。说实话，我第一次用的时候也觉得文档写得不够清楚。但用多了就发现，核心概念就那么几个。

先说说进程组。分布式训练时，每个GPU对应一个进程。这些进程组成一个组，组内可以互相通信。初始化进程组是第一步：

import torch.distributed as dist

# 初始化进程组
dist.init_process_group(
    backend='nccl',  # GPU通信用nccl
    init_method='env://',  # 通过环境变量获取rank和world_size
    world_size=4,  # 总共4个进程
    rank=rank  # 当前进程的编号
)

这里有几个关键参数：

backend：通信后端。GPU用nccl，CPU用gloo。我建议你直接用nccl，性能最好。
init_method：初始化方式。env://最常用，还有tcp://和file://。
world_size：总进程数，通常等于GPU数量。
rank：当前进程的编号，从0到world_size-1。

小技巧：

我个人习惯用环境变量来传rank和world_size。启动脚本里设好RANK、WORLD_SIZE、MASTER_ADDR和MASTER_PORT，代码里直接用env://，省事。

2.3 点对点通信

点对点通信，就是两个进程之间互相发数据。最常用的是send和recv。你想想看，这就像两个人打电话——一个人说，一个人听。

# 进程0发送数据给进程1
if rank == 0:
    tensor = torch.tensor([1, 2, 3]).cuda()
    dist.send(tensor, dst=1)
    
if rank == 1:
    tensor = torch.zeros(3).cuda()
    dist.recv(tensor, src=0)
    print(tensor)  # 输出: tensor([1, 2, 3])

点对点通信看着简单，但有个坑——死锁。我曾经写过这样的代码：

# 错误示例：可能死锁
if rank == 0:
    dist.send(tensor, dst=1)
    dist.recv(tensor, src=1)
    
if rank == 1:
    dist.send(tensor, dst=0)
    dist.recv(tensor, src=0)

为什么？因为send是阻塞的。进程0发完数据后，等着进程1发数据。进程1也一样。结果两个都在等，谁也不发。解决办法是用isend和irecv（非阻塞版本），或者调整发送顺序。

2.4 集合通信

集合通信就厉害了。它让所有进程一起干活，而不是两两通信。常用的有：

操作	说明	类比
Broadcast	一个进程把数据发给所有人	老师发试卷
All-Reduce	所有进程的数据求和/求平均，结果分给所有人	全班统计成绩
All-Gather	收集所有进程的数据，分给所有人	每人交作业，大家互相看
Reduce-Scatter	先求和，再把结果分散到各个进程	分组统计，每人拿一部分

分布式训练中最常用的是All-Reduce。每个GPU算完梯度后，用All-Reduce把梯度平均一下，然后各自更新参数。PyTorch的DistributedDataParallel底层就是这么干的。

# All-Reduce示例
tensor = torch.tensor([rank]).cuda()
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
print(f"Rank {rank}: {tensor}")  
# 假设4个进程，输出都是 tensor([6])，因为 0+1+2+3=6

核心要点：

集合通信的性能瓶颈通常在带宽和延迟。我遇到过一个问题：8卡训练时，All-Reduce占了总时间的30%。后来改用NCCL的ncclAllReduce，并调整了通信拓扑，时间降到了10%。

2.5 知识体系总览

说了这么多，咱们用一张图来总结本章的核心内容：

这张图把本章的三个核心模块串起来了。CUDA是基础，PyTorch分布式包是工具，通信方式是核心机制。三者缺一不可。

我的建议：

刚开始学分布式时，别急着上大模型。先写个简单的All-Reduce demo，跑通4卡。然后试试点对点通信，感受一下进程间怎么协作。我当年就是这么过来的——从单卡到多卡，一步步来，稳得很。

好了，这一章就到这里。记住：分布式编程的核心不是写代码，而是理解通信。把通信搞明白了，后面的大规模训练就是水到渠成的事。