第二章：压测环境搭建：硬件选型与软件栈配置
好，咱们直接进入正题。压测环境搭建这事儿，说简单也简单，说复杂……嗯，我踩过的坑能写满一页A4纸。今天就把这些经验掰开揉碎了讲给你听。
2.1 硬件选型：GPU / TPU / NPU 怎么挑？
先聊硬件。很多人一上来就问「哪个最好？」。我的回答是：没有最好，只有最合适。
2.1.1 GPU：最通用的选择
目前90%的AI训练压测都在GPU上跑。我个人习惯用NVIDIA的A100或H100做基准测试。为什么？因为生态最成熟，CUDA、cuDNN、TensorRT这些工具链你基本绕不开。
核心指标看这几点：
- 显存容量：决定了你能塞多大的batch size。我建议至少40GB起步，80GB更稳。
- 显存带宽：直接影响数据搬运速度。H100的3.35TB/s带宽，比A100快了近一倍。
- TFLOPS：理论算力，但别太迷信。实际吞吐受很多因素影响。
我的经验：如果你做LLM压测，显存带宽比算力更重要。我曾经在A100上跑175B模型，batch size死活上不去，换了H100后吞吐直接翻倍。
2.1.2 TPU：Google的专属利器
TPU这东西，说白了就是为TensorFlow量身定做的。如果你用PyTorch，我个人建议别碰TPU——兼容性问题会让你怀疑人生。
但如果你已经在Google Cloud上跑TensorFlow，TPU v4是个不错的选择。它的互联带宽极高，多卡扩展性比GPU强。我记得有一次帮客户做压测，8颗TPU v4的线性加速比能达到7.6，而同等规模的GPU集群只有6.8。
2.1.3 NPU：国产化的新选择
最近几年NPU（比如华为昇腾）开始冒头。说实话，我刚开始也不太信任，直到有一次在昇腾910B上跑ResNet-50压测……嗯，结果出乎意料地好。
注意：NPU的软件栈还不够成熟。我曾经遇到过算子兼容性问题，一个简单的LayerNorm在NPU上跑出了CPU的速度。建议先做小规模验证，再上大规模压测。
2.2 软件栈配置：CUDA、PyTorch/TensorFlow、Docker
硬件选好了，接下来就是搭软件栈。这部分我踩过的坑最多，你仔细听。
2.2.1 CUDA版本：别乱升级
很多人觉得CUDA越新越好。错！稳定压倒一切。
我建议用CUDA 11.8或12.1这两个版本。为什么？因为PyTorch和TensorFlow的主流版本都对这些CUDA版本做过充分测试。
# 检查CUDA版本
nvcc --version
# 推荐版本组合
CUDA 11.8 + PyTorch 2.0.1
CUDA 12.1 + PyTorch 2.1.0避坑指南：我曾经在CUDA 12.0上跑PyTorch 2.0，结果Dataloader莫名其妙卡死。后来发现是CUDA 12.0的bug，降级到11.8就解决了。所以，别追新。
2.2.2 PyTorch vs TensorFlow：选哪个？
这个问题我经常被问到。我的回答是：看团队。
| 对比项 | PyTorch | TensorFlow |
|---|---|---|
| 调试体验 | 动态图，方便调试 | 静态图，性能更优 |
| 分布式支持 | torch.distributed 成熟 | tf.distribute 也不错 |
| 社区生态 | 学术界主流 | 工业界更广 |
| 压测工具 | torch.profiler 好用 | TensorBoard 强大 |
我个人习惯用PyTorch做压测。原因很简单：torch.profiler太好用了。你想想看，能直接看到每个算子的耗时、显存占用、带宽利用率，调试起来多方便。
2.2.3 Docker：环境隔离的利器
压测环境最怕什么？环境不一致。今天在这台机器上跑得好好的，换一台就崩了。Docker就是来解决这个问题的。
# 我的标准Dockerfile
FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
RUN apt-get update && apt-get install -y \
python3.10 \
python3-pip \
&& rm -rf /var/lib/apt/lists/*
RUN pip install torch==2.0.1+cu118 \
torchvision==0.15.1+cu118 \
--extra-index-url https://download.pytorch.org/whl/cu118
RUN pip install tensorflow==2.12.0
# 压测工具
RUN pip install nvidia-ml-py3 psutil
CMD ["/bin/bash"]注意：Docker镜像别太大。我见过有人把整个conda环境打包进去，结果镜像20GB+，拉取就要半小时。建议只装必要依赖，其他用volume挂载。
2.3 网络与存储检查
硬件和软件都配好了，但还有一个隐形杀手——网络和存储。这两个地方出问题，压测结果基本废了。
2.3.1 网络：别让带宽成为瓶颈
多卡训练时，网络带宽直接决定通信效率。我建议至少用100Gbps的InfiniBand或RoCE。
怎么检查？用这个命令：
# 检查网络带宽
ib_write_bw -a -d mlx5_0
# 检查延迟
ibping -c 100 -s 1024如果发现带宽低于理论值的80%，嗯，你得排查一下了。我曾经遇到过网卡固件版本不匹配的问题，升级后带宽从60Gbps飙到了95Gbps。
2.3.2 存储：IOPS和吞吐量
数据加载速度跟不上训练速度，GPU就得干等着。这就是所谓的IO瓶颈。
我建议用NVMe SSD做本地缓存，或者用高性能分布式存储（比如Lustre、GPFS）。
# 测试存储性能
fio --name=test --ioengine=libaio --rw=randread \
--bs=4k --size=1G --numjobs=4 --runtime=60 \
--group_reporting我的经验：如果IOPS低于10万，或者吞吐量低于2GB/s，你的存储大概率是瓶颈。我曾经在HPC集群上跑压测，发现数据加载占了30%的时间，后来加了本地NVMe缓存，训练时间直接缩短了25%。
2.4 知识体系总览
说了这么多，我画张图帮你理清思路：
这张图把整个压测环境搭建的核心要素都串起来了。你仔细看，硬件、软件、网络存储这三块缺一不可。任何一个环节出问题，压测结果都会失真。
最后说一句：环境搭建别图快。我见过太多人急着跑压测，结果发现是环境问题导致数据不准，白白浪费几天时间。稳扎稳打，把每一步都验证到位，后面才能跑得顺畅。