从零搭建AI基础设施：单机多卡训练环境搭建与调优实战

1. 为什么我要从零搭建一套AI基础设施

三年前我从一家传统互联网公司跳到一家做AI应用的中型团队，面试时聊的是模型微调和推理优化，入职第一天Leader丢给我一句话：“训练环境你来搭一下，下周要跑一个7B的微调。”我当时心想，不就是装个CUDA、配个PyTorch吗，能有多难。结果那一周我几乎没合眼——驱动版本冲突、CUDA和cuDNN对不上、多卡通信跑不起来、数据集加载把内存打爆、训练到一半OOM、checkpoint存到一半磁盘满了。等我把这一套折腾完，我才真正理解什么叫AI-Infra：它不是简单的“装环境”，而是把算力、存储、网络、调度、框架、监控这一整条链路捏合成一个能稳定产出模型的基础设施。

这篇是我个人这一路踩坑、复盘、重构的总结，算是“一线工程师的AI-Infra之路”的第一章。我打算把最底层、最容易被忽视、但一旦出问题就让你通宵的部分讲透：单机多卡训练环境的搭建与调优。为什么从这个点切入？因为无论你后面要做分布式训练、推理服务化还是MLOps平台，单机多卡都是最小可用的验证单元。你连一台8卡机都跑不稳，谈什么千卡集群。

这篇文章适合谁看？如果你是会写Python、懂一点深度学习、但一碰到驱动、CUDA、NCCL、显存这些底层东西就头大的算法工程师，那这篇就是写给你的。如果你已经是有经验的Infra工程师，也可以看看我在参数选择、故障排查上的思路，或许能帮你省下几个通宵。全文我会尽量说人话，把“为什么这么配”讲清楚，而不是甩一堆命令让你抄。

2. 整体设计思路：为什么单机多卡是AI-Infra的第一道坎

2.1 从“能跑”到“跑得稳”的认知转变

很多人对AI-Infra的理解停留在“把环境装好，代码能跑起来”。我一开始也是这么想的，直到我遇到一个真实场景：一个7B模型在8卡A100上微调，前100个step一切正常，loss平稳下降，到第300个step突然卡死，GPU利用率掉到0，日志没有任何报错。我重启了三次，每次都在差不多的时间点卡住。最后定位到是NCCL的通信超时，因为其中一张卡的PCIe带宽被另一个进程占用了。

这件事让我意识到，AI-Infra的核心不是“能跑”，而是“可复现、可观测、可恢复”。能跑只是及格线，跑得稳、出问题能快速定位、挂了能自动恢复，才是基础设施该有的样子。所以我在设计单机多卡环境时，遵循三个原则：

• 版本锁定：驱动、CUDA、cuDNN、NCCL、PyTorch、Python，全部锁定版本，用容器或虚拟环境隔离，绝不依赖“系统默认”。
• 可观测优先：从第一天就接入GPU监控、显存监控、通信监控，不要等出问题才想起来加日志。
• 故障可恢复：checkpoint策略、断点续训、异常退出自动重启，这些不是“高级功能”，而是基础设施的标配。

2.2 为什么选择容器化而不是裸机环境

我试过两种方案：裸机直接装驱动和框架，以及用Docker容器封装整个环境。裸机的优点是性能损耗最小，GPU直通没有额外开销；缺点是环境一旦被污染，回滚成本极高。我曾经因为一个同事误装了另一个版本的CUDA，导致整个团队的训练任务全部失败，排查了两天才找到原因。

容器化方案我最终选择了NVIDIA Container Toolkit + 自定义基础镜像。核心考量是：

• 环境隔离：每个项目可以用不同的CUDA和框架版本，互不干扰。
• 可复现：镜像即环境，换一台机器拉起来就能跑，不用重新配。
• 版本管理：镜像打tag，出问题可以快速回滚到上一个稳定版本。

性能损耗方面，我实测下来，容器化相比裸机在训练吞吐上差距在2%以内，对于绝大多数场景完全可以接受。这个2%换来的环境稳定性和可复现性，我认为非常值。

2.3 单机多卡的核心技术栈拆解

一套完整的单机多卡训练环境，涉及的技术栈从上到下大概是这样的：

这张表是我自己梳理的，每次出问题我都会从上往下排查。经验是：越底层的問題，表现越诡异。驱动和CUDA的问题往往不会直接报错，而是表现为训练变慢、随机卡死、结果不对，非常难查。

3. 核心细节解析：驱动、CUDA与框架的版本匹配逻辑

3.1 驱动与CUDA的兼容关系到底怎么看

这是新手最容易懵的地方。NVIDIA的驱动和CUDA之间不是一一对应的，而是驱动有一个“最高支持的CUDA版本”。比如驱动版本535.x，最高支持CUDA 12.2；驱动版本525.x，最高支持CUDA 12.0。你装的CUDA版本只要不超过驱动支持的上限，就能跑。

但这里有个坑：CUDA的minor版本兼容性。CUDA 12.x系列内部是向后兼容的，也就是说驱动支持12.2，你可以跑12.0、12.1、12.2，但不能跑11.8。而CUDA 11.x和12.x之间是不兼容的，跨大版本必须升级驱动。

我一般这样操作：先确定框架需要的CUDA版本，再反推驱动版本。比如PyTorch 2.1官方推荐CUDA 12.1，那我就装支持12.1的驱动，比如535系列。具体命令：

# 查看当前驱动版本 nvidia-smi # 输出里会显示 Driver Version 和 CUDA Version # 注意这里的CUDA Version是驱动支持的最高版本，不是你装的版本

提示：nvidia-smi显示的CUDA Version是驱动能支持的最高CUDA运行时版本，实际使用的CUDA版本由你安装的CUDA Toolkit决定。这两个经常被混淆。

3.2 cuDNN和NCCL的版本选择

cuDNN是深度学习的加速库，PyTorch、TensorFlow都依赖它。cuDNN的版本必须和CUDA版本严格对应，比如cuDNN 8.9.x对应CUDA 12.x，cuDNN 8.6.x对应CUDA 11.8。装错了不会报错，但会在某些算子上报“找不到符号”或者直接段错误。

NCCL是多卡通信库，版本选择相对宽松，但建议用较新的版本，因为新版本对拓扑的感知和通信优化更好。我一般用NCCL 2.18以上，配合CUDA 12.x。

这里分享一个我踩过的坑：有一次我用了系统自带的NCCL，版本是2.10，结果在8卡A100上做AllReduce时性能只有预期的一半。换成2.19之后，带宽直接翻倍。原因是老版本NCCL没有正确识别NVLink拓扑，走了PCIe。所以NCCL版本不要用系统默认的，一定要自己装匹配的版本。

3.3 PyTorch与CUDA的对应关系

PyTorch的每个版本都明确标注了支持的CUDA版本。比如：

• PyTorch 2.0.x → CUDA 11.7、11.8
• PyTorch 2.1.x → CUDA 11.8、12.1
• PyTorch 2.2.x → CUDA 11.8、12.1

安装时用官方提供的index-url，不要用pip默认源，否则可能装到CPU版本。命令示例：

# 安装支持CUDA 12.1的PyTorch 2.1 pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 \ --index-url https://download.pytorch.org/whl/cu121

装完之后一定要验证：

import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.device_count()) print(torch.version.cuda) print(torch.backends.cudnn.version())

如果torch.cuda.is_available()返回False，八成是CUDA版本和PyTorch不匹配，或者驱动版本太低。

4. 实操过程：从裸机到可用的8卡训练环境

4.1 硬件拓扑检查与确认

在装任何软件之前，先确认硬件拓扑。这一步很多人跳过，但它决定了你后面多卡通信的效率。用nvidia-smi topo -m查看：

nvidia-smi topo -m

输出会显示每张卡之间的连接方式，常见的有：

• NVLink：卡间高速互联，带宽远高于PCIe，A100的NVLink带宽是600GB/s。
• PCIe：通过PCIe交换机连接，带宽取决于PCIe代数和通道数。
• SYS：跨NUMA节点，走CPU互联，最慢。

我遇到过一台机器，8张卡里有两张走的是SYS，导致这两张卡参与通信时整个AllReduce被拖慢。后来调整了进程绑定，把通信密集的卡放在同一NUMA节点，性能才恢复正常。

注意：如果拓扑里有SYS连接，建议在启动训练时用CUDA_VISIBLE_DEVICES把跨NUMA的卡排除，或者用numactl绑定CPU亲和性。

4.2 驱动安装与验证

驱动安装我推荐用官方runfile，比apt源更可控。步骤：

# 1. 禁用nouveau驱动 echo -e "blacklist nouveau\noptions nouveau modeset=0" | sudo tee /etc/modprobe.d/blacklist-nouveau.conf sudo update-initramfs -u # 2. 重启进入文本模式 sudo systemctl set-default multi-user.target sudo reboot # 3. 安装驱动 sudo sh NVIDIA-Linux-x86_64-535.104.05.run --silent --dkms # 4. 验证 nvidia-smi

安装完成后，nvidia-smi应该能正确显示所有GPU。如果只显示部分GPU，检查是否被其他进程占用，或者BIOS里是否禁用了某些PCIe插槽。

4.3 容器环境搭建

我用的基础镜像是nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04，这个镜像已经包含了CUDA和cuDNN，省去手动安装的麻烦。Dockerfile大概长这样：

FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04 # 安装Python和基础工具 RUN apt-get update && apt-get install -y \ python3.10 python3-pip git vim htop \ && rm -rf /var/lib/apt/lists/* # 安装NCCL RUN apt-get update && apt-get install -y libnccl2=2.19.3-1+cuda12.1 libnccl-dev=2.19.3-1+cuda12.1 # 安装PyTorch RUN pip3 install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 \ --index-url https://download.pytorch.org/whl/cu121 # 安装训练相关库 RUN pip3 install transformers==4.36.0 datasets==2.15.0 accelerate==0.25.0 \ deepspeed==0.12.0 wandb==0.16.0

启动容器时要注意挂载GPU和共享内存：

docker run -it --gpus all \ --shm-size=64g \ --ulimit memlock=-1 \ --ulimit stack=67108864 \ -v /data:/data \ -v /home/user:/workspace \ ai-infra:base bash

--shm-size这个参数非常关键。默认Docker的共享内存只有64MB，多卡训练时DataLoader的worker进程会通过共享内存传递数据，64MB根本不够，会报“Bus error”或者直接卡死。我一般设成64G，和物理内存匹配。

4.4 多卡通信验证

环境搭好后，第一件事不是跑训练，而是验证NCCL通信是否正常。用PyTorch自带的测试脚本：

import torch import torch.distributed as dist def test_nccl(): dist.init_process_group(backend='nccl') rank = dist.get_rank() world_size = dist.get_world_size() # 每张卡创建一个tensor tensor = torch.ones(1).cuda() * rank # AllReduce求和 dist.all_reduce(tensor, op=dist.ReduceOp.SUM) expected = sum(range(world_size)) assert tensor.item() == expected, f"Rank {rank}: expected {expected}, got {tensor.item()}" if rank == 0: print(f"NCCL test passed with {world_size} GPUs") dist.destroy_process_group() if __name__ == '__main__': test_nccl()

用torchrun启动：

torchrun --nproc_per_node=8 test_nccl.py

如果8张卡都能正确输出，说明NCCL通信正常。如果卡住或者报错，检查NCCL版本、网络接口配置（NCCL_SOCKET_IFNAME）、以及是否有防火墙拦截。

4.5 训练脚本的并行配置

单机多卡训练我用的是DDP（DistributedDataParallel），比DataParallel效率高得多。核心配置：

import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): dist.init_process_group( backend='nccl', init_method='env://', world_size=world_size, rank=rank ) torch.cuda.set_device(rank) def train(rank, world_size): setup(rank, world_size) model = MyModel().to(rank) model = DDP(model, device_ids=[rank]) # 数据加载要用DistributedSampler dataset = MyDataset() sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) dataloader = DataLoader(dataset, batch_size=32, sampler=sampler, num_workers=4) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) for epoch in range(num_epochs): sampler.set_epoch(epoch) # 每个epoch要重新shuffle for batch in dataloader: # 训练逻辑 pass dist.destroy_process_group()

启动命令：

torchrun --nproc_per_node=8 --nnodes=1 \ --master_addr=127.0.0.1 --master_port=29500 \ train.py

这里有个细节：sampler.set_epoch(epoch)必须调用，否则每个epoch的数据顺序都一样，影响训练效果。这个坑我踩过，loss曲线看起来正常，但模型收敛比预期慢，查了两天才发现是数据没shuffle。

5. 常见问题与排查技巧实录

5.1 训练卡死与NCCL超时

现象：训练跑一段时间后卡死，GPU利用率掉到0，日志无报错。

排查思路：

1. 先看nvidia-smi，确认GPU是否还在被占用。
2. 用py-spy dump --pid <PID>看进程卡在哪个调用栈。
3. 如果卡在NCCL相关调用，检查NCCL超时设置。

解决方法：

# 增加NCCL超时时间 export NCCL_TIMEOUT=1800 # 开启NCCL调试日志 export NCCL_DEBUG=INFO export NCCL_DEBUG_SUBSYS=ALL

NCCL默认超时是30分钟，大模型训练时某些集合通信可能超过这个时间，导致误判超时。我一般设成1800秒。

5.2 显存OOM的排查与优化

现象：训练到一半报CUDA out of memory。

排查思路：

1. 用torch.cuda.memory_summary()看显存分配情况。
2. 检查是否有显存泄漏，比如在循环里不断创建tensor而不释放。
3. 检查batch size是否过大。

优化手段：

• 梯度累积：用小batch跑多次，累积梯度再更新。
• 混合精度：用torch.cuda.amp，显存占用减少约40%。
• 梯度检查点：用torch.utils.checkpoint，用计算换显存。
• DeepSpeed ZeRO：把优化器状态、梯度、参数分片到多卡。

我实测下来，7B模型用混合精度+梯度检查点，单卡24G显存可以跑batch size 4；加上ZeRO-2，可以跑batch size 16。

5.3 多卡训练速度不升反降

现象：从单卡扩展到8卡，吞吐没有提升8倍，甚至比单卡还慢。

排查思路：

1. 用nvidia-smi topo -m检查拓扑，确认卡间通信是否走NVLink。
2. 检查DataLoader的num_workers是否足够，数据加载是否成为瓶颈。
3. 检查是否有卡在等待其他卡，用NCCL_DEBUG=INFO看通信耗时。

常见原因：

• 数据加载瓶颈：num_workers设得太小，CPU预处理跟不上GPU计算。我一般设成4 * num_gpus。
• 通信瓶颈：卡间走PCIe而不是NVLink，或者跨NUMA节点。
• 负载不均：某些卡分到的数据多，某些少，导致等待。

5.4 常见问题速查表

5.5 几个我踩过的坑

坑一：共享内存不足。前面提过，Docker默认64MB共享内存，多卡训练必崩。改成--shm-size=64g。

坑二：NCCL_SOCKET_IFNAME未设置。多网卡机器上，NCCL可能选错网卡，导致通信走慢速网络。用export NCCL_SOCKET_IFNAME=eth0指定。

坑三：checkpoint保存阻塞训练。保存大模型checkpoint时，如果直接torch.save，会阻塞训练进程几十秒。我后来改成异步保存，用一个单独线程写磁盘。

坑四：随机种子未固定。多卡训练时，每张卡的随机种子要不同，但整体要可复现。我一般用seed = base_seed + rank。

坑五：僵尸进程占用显存。训练异常退出后，有时进程没完全清理，显存还被占用。用fuser -v /dev/nvidia*找到进程并kill。

6. 监控与可观测性：让问题在发生前暴露

6.1 GPU监控的最小可用方案

训练跑起来只是第一步，能实时看到GPU状态才是关键。我用的是nvitop，比nvidia-smi直观得多：

pip install nvitop nvitop

它会实时显示每张卡的利用率、显存、温度、功耗，还能看到每个进程占用的资源。排查“哪张卡在偷懒”特别有用。

如果要长期记录，我用Prometheus + Grafana + DCGM Exporter。DCGM是NVIDIA的数据中心GPU监控工具，能采集到非常细粒度的指标，比如SM利用率、显存带宽、NVLink流量。配置大概是这样：

# 启动DCGM Exporter docker run -d --gpus all --rm -p 9400:9400 \ nvcr.io/nvidia/k8s/dcgm-exporter:3.2.5-3.1.8-ubuntu22.04

然后在Prometheus里配置抓取localhost:9400/metrics，Grafana里导入DCGM的dashboard模板。这套下来，GPU的历史状态一目了然。

6.2 训练指标监控

除了GPU，训练本身的指标也要监控。我用Weights & Biases（wandb）记录loss、学习率、梯度范数、吞吐量。关键是要记录每张卡的吞吐，而不只是平均值。如果某张卡明显慢，说明负载不均或者硬件有问题。

import wandb wandb.init(project="ai-infra", name="7b-finetune") # 训练循环里 wandb.log({ "loss": loss.item(), "lr": optimizer.param_groups[0]['lr'], "throughput": samples_per_sec, "gpu_util": get_gpu_util(rank) })

6.3 日志规范

日志要分级、要结构化、要能快速定位。我的做法是：

• INFO：每个epoch的开始和结束，loss值。
• WARNING：显存接近上限、通信耗时异常。
• ERROR：OOM、NCCL超时、checkpoint失败。

日志格式用JSON，方便后续用ELK或者Loki检索。关键字段包括：timestamp、rank、level、message、step、loss。

7. 写在最后：一些个人体会

这套单机多卡环境我从零搭到现在稳定运行，前后花了大概三个月，中间踩的坑比我前三年加起来都多。最大的体会是：AI-Infra没有银弹，只有不断排查和迭代。你今天解决了一个OOM，明天可能遇到一个更诡异的通信问题。重要的不是记住所有命令，而是建立一套排查方法论：从硬件到驱动，从驱动到框架，从框架到代码，逐层往下查。

另外，不要迷信“最新版本”。我见过太多人为了用新特性升级CUDA，结果整个环境崩掉。稳定压倒一切，版本锁定、镜像管理、回滚机制，这些才是基础设施的根基。

最后分享一个小技巧：每次环境变更前，先docker commit保存当前状态，出问题可以秒回滚。这个习惯帮我省了至少十次重装环境的时间。下一章我打算聊聊多机多卡的组网和调度，那又是另一个深坑，感兴趣的话可以持续关注。