第二章：压测环境搭建：硬件选型与软件配置

压测环境搭建，说白了就是给GPU集群“铺路”。路铺不好，后面跑再好的模型也是白搭。我见过太多团队，硬件买回来发现驱动不兼容，或者NCCL死活连不上，最后折腾两周还没开始压测。嗯，咱们今天就把这条路铺扎实。

2.1 硬件选型：A100 vs H100

先聊硬件。目前主流的压测卡就两款：NVIDIA A100和H100。怎么选？我个人的经验是：看你的业务场景和预算。

我个人习惯是：如果做千亿参数以下的大模型压测，A100完全够用。但如果你要压测万亿参数级别的分布式训练，H100的NVLink带宽优势就非常明显了。我在项目中遇到过，用A100跑8卡压测，NVLink带宽成了瓶颈，换成H100后吞吐直接提升了40%。

2.2 驱动与CUDA安装

硬件到位了，接下来装驱动。这里有个坑：驱动版本必须和CUDA版本匹配。我见过有人装了CUDA 12.0，却配了个470系列的驱动，结果nvcc编译通过，运行时直接报错。

我的建议是：

• 先确定CUDA版本（比如CUDA 11.8或12.1）
• 再查NVIDIA官方驱动兼容性表
• 最后安装对应驱动

举个例子，我常用的组合是：

# 安装驱动（以CUDA 12.1为例）
wget https://us.download.nvidia.com/tesla/535.154.05/NVIDIA-Linux-x86_64-535.154.05.run
chmod +x NVIDIA-Linux-x86_64-535.154.05.run
sudo ./NVIDIA-Linux-x86_64-535.154.05.run --silent

# 验证驱动
nvidia-smi
# 输出应显示驱动版本535.154.05，CUDA版本12.1

装完驱动后，再装CUDA Toolkit。我个人习惯用runfile方式安装，因为可以指定安装路径，方便多版本共存。

# 安装CUDA 12.1 Toolkit
wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
sudo sh cuda_12.1.0_530.30.02_linux.run --toolkit --silent --override

# 设置环境变量
echo 'export PATH=/usr/local/cuda-12.1/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

# 验证CUDA
nvcc --version

2.3 NCCL库配置

NCCL是分布式压测的灵魂。没有它，多卡通信就是纸上谈兵。NCCL的配置其实不复杂，但有几个关键参数必须调对。

先安装NCCL：

# 下载NCCL（以2.19.3为例）
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/nccl-repo-ubuntu2204-2.19.3-ga-cuda12.1_1-1_amd64.deb
sudo dpkg -i nccl-repo-ubuntu2204-2.19.3-ga-cuda12.1_1-1_amd64.deb
sudo apt update
sudo apt install libnccl2 libnccl-dev

# 验证NCCL版本
dpkg -l | grep nccl

配置环境变量时，我建议显式指定NCCL的通信算法和协议：

# 在~/.bashrc中添加
export NCCL_ALGO=Ring
export NCCL_PROTO=Simple
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_IB_HCA=mlx5_0,mlx5_1
export NCCL_SOCKET_IFNAME=eth0

为什么会这样配？NCCL_ALGO=Ring 是环形通信算法，适合8卡以内的压测。如果卡数超过8张，可以考虑 Tree 算法。我在项目中遇到过，用默认的 Auto 算法，结果NCCL自己选了 NVLS 算法，反而在小batch size下性能更差。

2.4 Docker环境准备

最后一步，把环境打包进Docker。为什么用Docker？说白了，就是为了环境一致性。你想想看，团队里10个人，每个人装驱动的方式都不一样，压测结果怎么对比？

我推荐用NVIDIA官方提供的NGC容器作为基础镜像：

# 拉取NGC PyTorch镜像（包含CUDA 12.1 + NCCL 2.19）
docker pull nvcr.io/nvidia/pytorch:23.12-py3

# 验证镜像中的环境
docker run --rm --gpus all nvcr.io/nvidia/pytorch:23.12-py3 nvidia-smi
docker run --rm --gpus all nvcr.io/nvidia/pytorch:23.12-py3 python -c "import torch; print(torch.cuda.nccl.version())"

如果你需要自定义环境，可以写Dockerfile：

FROM nvcr.io/nvidia/pytorch:23.12-py3

# 安装压测工具
RUN pip install --upgrade pip && \
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# 设置NCCL环境变量
ENV NCCL_ALGO=Ring
ENV NCCL_PROTO=Simple
ENV NCCL_DEBUG=INFO
ENV NCCL_IB_DISABLE=0

# 复制压测脚本
COPY benchmark/ /workspace/benchmark/

WORKDIR /workspace

启动容器时，有几个参数必须注意：

# 启动容器（8卡压测示例）
docker run -it --rm \
    --gpus all \
    --network host \
    --ipc=host \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    -v /mnt/data:/data \
    my-benchmark:latest \
    bash

2.5 环境验证：跑一个简单的AllReduce测试

环境搭好了，怎么知道对不对？跑一个NCCL的AllReduce测试就行。我一般用 nccl-tests 这个工具：

# 编译nccl-tests
git clone https://github.com/NVIDIA/nccl-tests.git
cd nccl-tests
make

# 跑8卡AllReduce测试
mpirun -np 8 \
    --allow-run-as-root \
    --hostfile hostfile \
    -x NCCL_DEBUG=INFO \
    -x NCCL_ALGO=Ring \
    ./build/all_reduce_perf -b 8 -e 128M -f 2 -g 1

如果输出中看到 # Avg bus bandwidth 接近理论值（A100约600 GB/s，H100约900 GB/s），说明环境没问题。如果带宽只有几十GB/s，那就要检查NVLink是否连通、IB网卡是否配置正确。

嗯，到这里，压测环境就搭好了。硬件选型、驱动CUDA、NCCL配置、Docker环境，每一步都有坑，但踩过去之后，后面的压测就会非常顺畅。