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

环境搭建这件事，说白了就是给你的压测找个靠谱的「跑道」。跑道不行，你测出来的数据全是噪音。我见过太多团队，花大价钱买了卡，结果环境没配好，跑出来的吞吐量还不如隔壁用旧卡的同学。嗯，这里面的坑，我一个个给你说清楚。

2.1 硬件选型：GPU vs CPU，到底选谁？

先问个问题：你的算子是计算密集型，还是访存密集型？

我个人习惯，先看算子的计算密度（Compute Intensity）。如果每个元素的计算量很大，比如卷积、矩阵乘法，那GPU是首选。如果算子主要是数据搬运、条件分支多，比如某些稀疏操作，CPU反而更稳。

GPU选型建议：

  • NVIDIA A100 / H100：适合大模型、高吞吐场景。我去年压测一个Transformer算子，A100的FP16算力达到312 TFLOPS，比V100翻了3倍多。
  • NVIDIA V100 / A10：性价比之选。如果你的算子不是特别吃显存，V100的32GB显存够用，而且价格只有A100的1/3。
  • RTX 4090：个人开发机首选。24GB显存，FP16算力82 TFLOPS，关键是便宜。但注意，它没有NVLink，多卡通信会慢。

CPU选型建议：

  • AMD EPYC 7763：64核128线程，适合并行度高的CPU算子。我测过一个稀疏矩阵乘法，EPYC比Intel Xeon快40%。
  • Intel Xeon Platinum 8380：40核80线程，AVX-512指令集对某些算子有奇效。但功耗高，散热要跟上。

核心原则：压测环境要和生产环境保持一致。你想想看，如果生产用A100，你拿4090测，那数据能信吗？

2.2 软件栈配置：CUDA、PyTorch/TensorFlow

硬件选好了，接下来就是软件栈。这里我踩过不少坑，一个个说。

2.2.1 CUDA版本选择

CUDA版本不是越新越好。我建议：

  • CUDA 11.8：最稳的版本。PyTorch 2.0、TensorFlow 2.12都原生支持。我压测时90%的场景用这个。
  • CUDA 12.x：如果你用H100或者需要FP8支持，可以上。但注意，有些旧算子可能不兼容。

我曾经因为CUDA版本没对齐，导致一个自定义算子跑出来的结果比预期慢3倍。查了两天才发现是CUDA runtime和driver版本不匹配。嗯，这种低级错误，犯一次就够了。

2.2.2 PyTorch vs TensorFlow

说实话，两个框架我都用。但压测时，我倾向PyTorch：

  • PyTorch 2.0+：torch.compile能自动融合算子，对压测很友好。我测过一个LayerNorm，开启编译后吞吐提升30%。
  • TensorFlow 2.x：如果你用TF Serving做推理，那压测环境必须用TF。但注意，TF的算子执行顺序有时候会变，影响结果复现。

小技巧：压测前，用 torch.backends.cudnn.benchmark = True 让cuDNN自动选择最优算法。我测卷积算子时，这个开关能带来15%-20%的性能提升。

2.2.3 其他依赖库

除了框架本身，还有几个库必须配齐：

  • cuDNN 8.9+：卷积、池化算子的加速库。版本要和CUDA对齐。
  • NCCL 2.18+：多卡通信库。如果你做多卡压测，这个必须装。我遇到过NCCL版本太低导致通信带宽只有理论值一半的情况。
  • TensorRT 8.6+：推理优化库。如果你压测的是推理场景，用TensorRT能再压榨20%-50%的性能。

2.3 Docker环境准备：一次配置，到处运行

环境配置最怕什么？怕换台机器就崩。Docker就是来解决这个问题的。我个人的习惯是：所有压测环境都容器化，绝不直接在宿主机上跑。

2.3.1 基础镜像选择

别从零开始搭。直接用NVIDIA官方提供的镜像：

# 拉取PyTorch镜像
docker pull nvcr.io/nvidia/pytorch:23.12-py3

# 拉取TensorFlow镜像
docker pull nvcr.io/nvidia/tensorflow:23.12-tf2-py3

这些镜像已经装好了CUDA、cuDNN、NCCL，你只需要装自己的算子库就行。

2.3.2 Dockerfile示例

我一般这样写Dockerfile：

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

# 安装额外依赖
RUN pip install --no-cache-dir \
    pytest \
    psutil \
    nvidia-ml-py3

# 复制算子代码
COPY ./operators /workspace/operators

# 设置工作目录
WORKDIR /workspace/operators

# 默认命令
CMD ["python", "benchmark.py"]

注意：Docker容器默认不共享宿主机内存。如果你压测大模型，记得加 --shm-size=32g 参数。我曾经因为这个，一个Batch Size为32的模型直接OOM，排查了半天才发现是共享内存不够。

2.3.3 启动容器

启动命令我一般这样写：

docker run --gpus all \
    --shm-size=32g \
    -v /data:/data \
    -v /home/user/operators:/workspace/operators \
    -it nvcr.io/nvidia/pytorch:23.12-py3 \
    python benchmark.py --op conv --batch_size 64

这里有几个关键点：

  • --gpus all：让容器访问所有GPU。如果你只想用某张卡，可以指定 --gpus '"device=0,1"'。
  • -v：挂载数据卷。我习惯把数据集和代码挂载进去，这样镜像可以保持轻量。
  • --shm-size：前面说过了，大模型必加。

2.4 环境验证：跑个简单的算子试试

环境搭好了，别急着压测。先跑个简单的算子验证一下：

import torch
import time

# 验证CUDA是否可用
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU count: {torch.cuda.device_count()}")
print(f"GPU name: {torch.cuda.get_device_name(0)}")

# 跑一个简单的矩阵乘法
a = torch.randn(4096, 4096, device='cuda')
b = torch.randn(4096, 4096, device='cuda')

# 预热
for _ in range(10):
    c = torch.matmul(a, b)

# 正式测试
start = time.time()
for _ in range(100):
    c = torch.matmul(a, b)
torch.cuda.synchronize()
end = time.time()

print(f"Average time: {(end - start) / 100 * 1000:.2f} ms")
print(f"TFLOPS: {2 * 4096**3 / ((end - start) / 100) / 1e12:.2f}")

如果这个代码能跑通，而且TFLOPS接近理论值，那你的环境就基本没问题了。

避坑指南：我曾经在Docker里跑这个测试，发现TFLOPS只有理论值的60%。查了半天，发现是容器没开 --gpus all，导致只用了1张卡。嗯，这种低级错误，犯一次就够了。

2.5 本章知识体系

下面这张图，帮你理清压测环境搭建的核心逻辑：

压测环境搭建核心流程 硬件选型 软件栈配置 Docker环境 GPU: A100 / H100 / V100 CPU: EPYC / Xeon CUDA 11.8 / 12.x PyTorch / TensorFlow cuDNN / NCCL / TensorRT 基础镜像选择 Dockerfile编写 容器启动参数 环境验证：跑矩阵乘法测试 输出：稳定的压测环境 + 可复现的基准数据

这张图把整个流程串起来了。从硬件选型开始，到软件栈配置，再到Docker容器化，最后用矩阵乘法验证。每一步都踩过坑，希望你别再踩一遍。

环境搭好了，下一步就是设计压测用例了。不过那是下一章的事，咱们先把眼前的环境搞定。

最后提醒一句：压测环境一定要和生产环境保持一致。你想想看，如果生产用A100，你拿4090测，那数据能信吗？我见过太多团队，因为环境不一致，压测数据漂亮，上线就崩。嗯，这种坑，咱们不踩。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →