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 单机多卡的核心技术栈拆解一套完整的单机多卡训练环境涉及的技术栈从上到下大概是这样的层级组件作用常见坑点硬件层GPU、NVLink、PCIe提供算力和卡间通信拓扑结构影响通信效率驱动层NVIDIA Driver硬件与上层软件的桥梁版本与CUDA强绑定运行时CUDA、cuDNN并行计算和深度学习算子版本匹配极其严格通信层NCCL多卡集合通信超时、拓扑感知配置框架层PyTorch、DeepSpeed模型训练和并行策略版本与CUDA对应调度层进程管理、监控资源分配和可观测显存泄漏、僵尸进程这张表是我自己梳理的每次出问题我都会从上往下排查。经验是越底层的問題表现越诡异。驱动和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.xcuDNN 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.8PyTorch 2.1.x → CUDA 11.8、12.1PyTorch 2.2.x → CUDA 11.8、12.1安装时用官方提供的index-url不要用pip默认源否则可能装到CPU版本。命令示例# 安装支持CUDA 12.1的PyTorch 2.1 pip install torch2.1.0 torchvision0.16.0 torchaudio2.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卡间高速互联带宽远高于PCIeA100的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 modeset0 | 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 libnccl22.19.3-1cuda12.1 libnccl-dev2.19.3-1cuda12.1 # 安装PyTorch RUN pip3 install torch2.1.0 torchvision0.16.0 torchaudio2.1.0 \ --index-url https://download.pytorch.org/whl/cu121 # 安装训练相关库 RUN pip3 install transformers4.36.0 datasets2.15.0 accelerate0.25.0 \ deepspeed0.12.0 wandb0.16.0启动容器时要注意挂载GPU和共享内存docker run -it --gpus all \ --shm-size64g \ --ulimit memlock-1 \ --ulimit stack67108864 \ -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(backendnccl) rank dist.get_rank() world_size dist.get_world_size() # 每张卡创建一个tensor tensor torch.ones(1).cuda() * rank # AllReduce求和 dist.all_reduce(tensor, opdist.ReduceOp.SUM) expected sum(range(world_size)) assert tensor.item() expected, fRank {rank}: expected {expected}, got {tensor.item()} if rank 0: print(fNCCL test passed with {world_size} GPUs) dist.destroy_process_group() if __name__ __main__: test_nccl()用torchrun启动torchrun --nproc_per_node8 test_nccl.py如果8张卡都能正确输出说明NCCL通信正常。如果卡住或者报错检查NCCL版本、网络接口配置NCCL_SOCKET_IFNAME、以及是否有防火墙拦截。4.5 训练脚本的并行配置单机多卡训练我用的是DDPDistributedDataParallel比DataParallel效率高得多。核心配置import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): dist.init_process_group( backendnccl, init_methodenv://, world_sizeworld_size, rankrank ) 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_replicasworld_size, rankrank) dataloader DataLoader(dataset, batch_size32, samplersampler, num_workers4) optimizer torch.optim.AdamW(model.parameters(), lr1e-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_node8 --nnodes1 \ --master_addr127.0.0.1 --master_port29500 \ train.py这里有个细节sampler.set_epoch(epoch)必须调用否则每个epoch的数据顺序都一样影响训练效果。这个坑我踩过loss曲线看起来正常但模型收敛比预期慢查了两天才发现是数据没shuffle。5. 常见问题与排查技巧实录5.1 训练卡死与NCCL超时现象训练跑一段时间后卡死GPU利用率掉到0日志无报错。排查思路先看nvidia-smi确认GPU是否还在被占用。用py-spy dump --pid PID看进程卡在哪个调用栈。如果卡在NCCL相关调用检查NCCL超时设置。解决方法# 增加NCCL超时时间 export NCCL_TIMEOUT1800 # 开启NCCL调试日志 export NCCL_DEBUGINFO export NCCL_DEBUG_SUBSYSALLNCCL默认超时是30分钟大模型训练时某些集合通信可能超过这个时间导致误判超时。我一般设成1800秒。5.2 显存OOM的排查与优化现象训练到一半报CUDA out of memory。排查思路用torch.cuda.memory_summary()看显存分配情况。检查是否有显存泄漏比如在循环里不断创建tensor而不释放。检查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倍甚至比单卡还慢。排查思路用nvidia-smi topo -m检查拓扑确认卡间通信是否走NVLink。检查DataLoader的num_workers是否足够数据加载是否成为瓶颈。检查是否有卡在等待其他卡用NCCL_DEBUGINFO看通信耗时。常见原因数据加载瓶颈num_workers设得太小CPU预处理跟不上GPU计算。我一般设成4 * num_gpus。通信瓶颈卡间走PCIe而不是NVLink或者跨NUMA节点。负载不均某些卡分到的数据多某些少导致等待。5.4 常见问题速查表问题现象可能原因排查命令解决方法训练卡死NCCL超时NCCL_DEBUGINFO增大NCCL_TIMEOUT显存OOMbatch过大/泄漏torch.cuda.memory_summary()梯度累积/混合精度多卡变慢通信瓶颈nvidia-smi topo -m调整进程绑定数据加载慢num_workers不足htop看CPU增大num_workers结果不对随机种子未固定-设置所有随机种子容器内无GPU未挂载GPUdocker run --gpus all加--gpus参数5.5 几个我踩过的坑坑一共享内存不足。前面提过Docker默认64MB共享内存多卡训练必崩。改成--shm-size64g。坑二NCCL_SOCKET_IFNAME未设置。多网卡机器上NCCL可能选错网卡导致通信走慢速网络。用export NCCL_SOCKET_IFNAMEeth0指定。坑三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/metricsGrafana里导入DCGM的dashboard模板。这套下来GPU的历史状态一目了然。6.2 训练指标监控除了GPU训练本身的指标也要监控。我用Weights Biaseswandb记录loss、学习率、梯度范数、吞吐量。关键是要记录每张卡的吞吐而不只是平均值。如果某张卡明显慢说明负载不均或者硬件有问题。import wandb wandb.init(projectai-infra, name7b-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显存接近上限、通信耗时异常。ERROROOM、NCCL超时、checkpoint失败。日志格式用JSON方便后续用ELK或者Loki检索。关键字段包括timestamp、rank、level、message、step、loss。7. 写在最后一些个人体会这套单机多卡环境我从零搭到现在稳定运行前后花了大概三个月中间踩的坑比我前三年加起来都多。最大的体会是AI-Infra没有银弹只有不断排查和迭代。你今天解决了一个OOM明天可能遇到一个更诡异的通信问题。重要的不是记住所有命令而是建立一套排查方法论从硬件到驱动从驱动到框架从框架到代码逐层往下查。另外不要迷信“最新版本”。我见过太多人为了用新特性升级CUDA结果整个环境崩掉。稳定压倒一切版本锁定、镜像管理、回滚机制这些才是基础设施的根基。最后分享一个小技巧每次环境变更前先docker commit保存当前状态出问题可以秒回滚。这个习惯帮我省了至少十次重装环境的时间。下一章我打算聊聊多机多卡的组网和调度那又是另一个深坑感兴趣的话可以持续关注。
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 单机多卡的核心技术栈拆解一套完整的单机多卡训练环境涉及的技术栈从上到下大概是这样的层级组件作用常见坑点硬件层GPU、NVLink、PCIe提供算力和卡间通信拓扑结构影响通信效率驱动层NVIDIA Driver硬件与上层软件的桥梁版本与CUDA强绑定运行时CUDA、cuDNN并行计算和深度学习算子版本匹配极其严格通信层NCCL多卡集合通信超时、拓扑感知配置框架层PyTorch、DeepSpeed模型训练和并行策略版本与CUDA对应调度层进程管理、监控资源分配和可观测显存泄漏、僵尸进程这张表是我自己梳理的每次出问题我都会从上往下排查。经验是越底层的問題表现越诡异。驱动和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.xcuDNN 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.8PyTorch 2.1.x → CUDA 11.8、12.1PyTorch 2.2.x → CUDA 11.8、12.1安装时用官方提供的index-url不要用pip默认源否则可能装到CPU版本。命令示例# 安装支持CUDA 12.1的PyTorch 2.1 pip install torch2.1.0 torchvision0.16.0 torchaudio2.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卡间高速互联带宽远高于PCIeA100的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 modeset0 | 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 libnccl22.19.3-1cuda12.1 libnccl-dev2.19.3-1cuda12.1 # 安装PyTorch RUN pip3 install torch2.1.0 torchvision0.16.0 torchaudio2.1.0 \ --index-url https://download.pytorch.org/whl/cu121 # 安装训练相关库 RUN pip3 install transformers4.36.0 datasets2.15.0 accelerate0.25.0 \ deepspeed0.12.0 wandb0.16.0启动容器时要注意挂载GPU和共享内存docker run -it --gpus all \ --shm-size64g \ --ulimit memlock-1 \ --ulimit stack67108864 \ -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(backendnccl) rank dist.get_rank() world_size dist.get_world_size() # 每张卡创建一个tensor tensor torch.ones(1).cuda() * rank # AllReduce求和 dist.all_reduce(tensor, opdist.ReduceOp.SUM) expected sum(range(world_size)) assert tensor.item() expected, fRank {rank}: expected {expected}, got {tensor.item()} if rank 0: print(fNCCL test passed with {world_size} GPUs) dist.destroy_process_group() if __name__ __main__: test_nccl()用torchrun启动torchrun --nproc_per_node8 test_nccl.py如果8张卡都能正确输出说明NCCL通信正常。如果卡住或者报错检查NCCL版本、网络接口配置NCCL_SOCKET_IFNAME、以及是否有防火墙拦截。4.5 训练脚本的并行配置单机多卡训练我用的是DDPDistributedDataParallel比DataParallel效率高得多。核心配置import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): dist.init_process_group( backendnccl, init_methodenv://, world_sizeworld_size, rankrank ) 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_replicasworld_size, rankrank) dataloader DataLoader(dataset, batch_size32, samplersampler, num_workers4) optimizer torch.optim.AdamW(model.parameters(), lr1e-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_node8 --nnodes1 \ --master_addr127.0.0.1 --master_port29500 \ train.py这里有个细节sampler.set_epoch(epoch)必须调用否则每个epoch的数据顺序都一样影响训练效果。这个坑我踩过loss曲线看起来正常但模型收敛比预期慢查了两天才发现是数据没shuffle。5. 常见问题与排查技巧实录5.1 训练卡死与NCCL超时现象训练跑一段时间后卡死GPU利用率掉到0日志无报错。排查思路先看nvidia-smi确认GPU是否还在被占用。用py-spy dump --pid PID看进程卡在哪个调用栈。如果卡在NCCL相关调用检查NCCL超时设置。解决方法# 增加NCCL超时时间 export NCCL_TIMEOUT1800 # 开启NCCL调试日志 export NCCL_DEBUGINFO export NCCL_DEBUG_SUBSYSALLNCCL默认超时是30分钟大模型训练时某些集合通信可能超过这个时间导致误判超时。我一般设成1800秒。5.2 显存OOM的排查与优化现象训练到一半报CUDA out of memory。排查思路用torch.cuda.memory_summary()看显存分配情况。检查是否有显存泄漏比如在循环里不断创建tensor而不释放。检查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倍甚至比单卡还慢。排查思路用nvidia-smi topo -m检查拓扑确认卡间通信是否走NVLink。检查DataLoader的num_workers是否足够数据加载是否成为瓶颈。检查是否有卡在等待其他卡用NCCL_DEBUGINFO看通信耗时。常见原因数据加载瓶颈num_workers设得太小CPU预处理跟不上GPU计算。我一般设成4 * num_gpus。通信瓶颈卡间走PCIe而不是NVLink或者跨NUMA节点。负载不均某些卡分到的数据多某些少导致等待。5.4 常见问题速查表问题现象可能原因排查命令解决方法训练卡死NCCL超时NCCL_DEBUGINFO增大NCCL_TIMEOUT显存OOMbatch过大/泄漏torch.cuda.memory_summary()梯度累积/混合精度多卡变慢通信瓶颈nvidia-smi topo -m调整进程绑定数据加载慢num_workers不足htop看CPU增大num_workers结果不对随机种子未固定-设置所有随机种子容器内无GPU未挂载GPUdocker run --gpus all加--gpus参数5.5 几个我踩过的坑坑一共享内存不足。前面提过Docker默认64MB共享内存多卡训练必崩。改成--shm-size64g。坑二NCCL_SOCKET_IFNAME未设置。多网卡机器上NCCL可能选错网卡导致通信走慢速网络。用export NCCL_SOCKET_IFNAMEeth0指定。坑三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/metricsGrafana里导入DCGM的dashboard模板。这套下来GPU的历史状态一目了然。6.2 训练指标监控除了GPU训练本身的指标也要监控。我用Weights Biaseswandb记录loss、学习率、梯度范数、吞吐量。关键是要记录每张卡的吞吐而不只是平均值。如果某张卡明显慢说明负载不均或者硬件有问题。import wandb wandb.init(projectai-infra, name7b-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显存接近上限、通信耗时异常。ERROROOM、NCCL超时、checkpoint失败。日志格式用JSON方便后续用ELK或者Loki检索。关键字段包括timestamp、rank、level、message、step、loss。7. 写在最后一些个人体会这套单机多卡环境我从零搭到现在稳定运行前后花了大概三个月中间踩的坑比我前三年加起来都多。最大的体会是AI-Infra没有银弹只有不断排查和迭代。你今天解决了一个OOM明天可能遇到一个更诡异的通信问题。重要的不是记住所有命令而是建立一套排查方法论从硬件到驱动从驱动到框架从框架到代码逐层往下查。另外不要迷信“最新版本”。我见过太多人为了用新特性升级CUDA结果整个环境崩掉。稳定压倒一切版本锁定、镜像管理、回滚机制这些才是基础设施的根基。最后分享一个小技巧每次环境变更前先docker commit保存当前状态出问题可以秒回滚。这个习惯帮我省了至少十次重装环境的时间。下一章我打算聊聊多机多卡的组网和调度那又是另一个深坑感兴趣的话可以持续关注。
