第三课：单机多卡环境搭建

好，咱们进入实战环节。

这一课，我带你搞定单机多卡的环境搭建。说白了，就是让一台机器上的多张GPU能协同工作。这是分布式训练的基础，也是你后面跑MirroredStrategy的前提。

硬件要求：GPU与NVLink

先聊聊硬件。我个人的习惯是，在动手装环境之前，先把硬件摸清楚。不然你折腾半天，发现卡不支持，那就尴尬了。

GPU选型

TensorFlow的MirroredStrategy，对GPU有明确要求。它需要支持CUDA的计算能力在3.5以上。不过，我建议你直接用NVIDIA的Volta架构（V100）或更新的Turing（T4、RTX 2080）、Ampere（A100、RTX 3080）架构。这些卡对分布式训练的支持更完善。

我在项目中遇到过用GTX 1080 Ti跑多卡的情况。嗯，能跑，但效率一般。因为老架构的卡，跨卡通信带宽有限，容易成为瓶颈。

NVLink：加速的秘密武器

你想想看，多卡训练时，最耗时的操作是什么？是梯度同步。每张卡算完自己的梯度，得告诉其他卡。这个通信过程，如果走PCIe总线，速度有限。

NVLink就是NVIDIA搞出来的高速互联技术。它让GPU之间可以直接通信，带宽比PCIe高好几倍。我曾经在V100集群上做过对比：用NVLink连接，梯度同步时间能缩短60%以上。

互联方式	带宽	适用场景
PCIe 3.0 x16	约16 GB/s	小模型、2卡训练
NVLink 2.0	约150 GB/s	大模型、4卡/8卡训练
NVSwitch	约600 GB/s	DGX系列、大规模集群

软件环境：CUDA、cuDNN、TensorFlow

硬件搞定了，接下来装软件。这部分我踩过不少坑，给你总结一套稳妥的流程。

CUDA Toolkit安装

CUDA是NVIDIA的并行计算平台。TensorFlow依赖它来调用GPU。

版本选择：别追新。我建议用CUDA 11.2或11.8。为什么？因为TensorFlow 2.10对这两个版本支持最稳定。我曾经手贱装了CUDA 12.0，结果TensorFlow报了一堆兼容性错误，折腾了两天才回退。

# 下载CUDA 11.8（以Ubuntu 20.04为例）
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
sudo sh cuda_11.8.0_520.61.05_linux.run

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

# 验证安装
nvcc --version

cuDNN安装

cuDNN是深度神经网络的加速库。它基于CUDA，专门优化了卷积、池化等操作。

安装时注意：cuDNN版本必须和CUDA版本匹配。比如CUDA 11.8对应cuDNN 8.6。我见过有人把cuDNN 8.0配CUDA 11.8，结果训练时各种段错误。

# 下载cuDNN 8.6（需要NVIDIA开发者账号）
# 解压并复制到CUDA目录
tar -xzvf cudnn-linux-x86_64-8.6.0.163_cuda11-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda-11.8/include
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda-11.8/lib64
sudo chmod a+r /usr/local/cuda-11.8/include/cudnn*.h /usr/local/cuda-11.8/lib64/libcudnn*

# 验证
cat /usr/local/cuda-11.8/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

TensorFlow安装

终于到主角了。我推荐用pip安装TensorFlow 2.10，这是最后一个原生支持GPU的版本（2.11开始需要自己编译）。

# 创建虚拟环境（强烈推荐）
python3 -m venv tf_env
source tf_env/bin/activate

# 安装TensorFlow
pip install tensorflow==2.10

# 验证GPU是否可用
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

验证多卡可用性

环境装好了，怎么知道多卡能不能用？写个简单的测试脚本。

检查GPU数量

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
print(f"检测到 {len(gpus)} 张GPU:")
for gpu in gpus:
    print(f"  - {gpu.name}")

if len(gpus) < 2:
    print("警告：至少需要2张GPU才能进行多卡训练！")
else:
    print("多卡环境就绪！")

测试梯度同步

这个测试能验证MirroredStrategy是否能正确分配计算任务。

import tensorflow as tf

# 创建MirroredStrategy
strategy = tf.distribute.MirroredStrategy()
print(f"设备数量: {strategy.num_replicas_in_sync}")

# 在策略范围内定义模型
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

# 生成随机数据
import numpy as np
x_train = np.random.random((1000, 784))
y_train = np.random.randint(0, 10, (1000,))

# 训练一个batch
model.fit(x_train, y_train, batch_size=64, epochs=1, verbose=1)
print("多卡训练测试通过！")

常见问题排查

如果你遇到问题，别慌。我列几个我踩过的坑：

• CUDA_ERROR_OUT_OF_MEMORY：显存不够。减小batch_size，或者用tf.config.set_memory_growth(gpu, True)开启显存动态增长。
• Could not create cudnn handle：cuDNN版本不匹配。检查一下cuDNN和CUDA的版本对应关系。
• No GPU devices found：驱动没装好。运行nvidia-smi看看能不能识别到GPU。

好了，环境搭建就到这里。下一课，我们正式进入MirroredStrategy的实战代码。到时候你会看到，前面这些准备工作，都是值得的。