1. 从“调参侠”到“基建工”AI-Infra到底在修什么路刚入行那会儿我对AI工程师的理解特别朴素不就是调调包、跑跑模型、盯着loss曲线祈祷它别炸嘛。直到有一次我负责的一个图像分类任务单卡跑得好好的领导说“上四卡并行吧速度快”。结果一上多卡吞吐量不升反降GPU利用率从85%掉到30%出头日志里全是NCCL timeout。那天晚上我对着nvidia-smi里那几条几乎躺平的利用率曲线第一次意识到模型代码写得再漂亮底层跑不起来全是白搭。这就是AI-Infra人工智能基础设施要解决的问题。它不负责发明新的注意力机制也不负责刷榜SOTA它负责的是让算法工程师写的那堆PyTorch代码能在真实的硬件集群上跑得动、跑得快、跑得稳。你可以把AI-Infra工程师理解成F1车队里的底盘和动力总成团队——车手算法工程师负责开得快我们负责让车在高速过弯时不散架。具体来说AI-Infra的日常工作横跨这么几个层面最底下是硬件层你得知道A100和H100的显存带宽差多少、NVLink和PCIe在通信上的瓶颈在哪往上是集群调度层Kubernetes怎么调、Slurm怎么配、GPU怎么切分和隔离再往上是训练框架层分布式数据并行DDP、张量并行TP、流水线并行PP怎么组合DeepSpeed和Megatron-LM各自的坑在哪最上面还有推理服务层模型怎么量化、怎么批处理、怎么在延迟和吞吐之间找平衡点。这篇文章适合谁看如果你是个算法工程师发现自己写的模型一上规模就各种玄学问题想搞清楚底层到底发生了什么那这篇东西能帮你补上缺失的那一环。如果你是个刚转行做AI-Infra的新人面对一堆缩写和配置项不知道从哪下手那这篇东西能给你一条相对清晰的路径。如果你是个后端或运维工程师想往AI方向靠那更好你已有的系统功底在这里非常值钱缺的只是对AI工作负载特性的理解。我打算把这个系列写成一个“一线工程师的踩坑与填坑记录”不讲空泛的架构图只讲我实际遇到过的问题、排查的过程、以及最后怎么解决的。第一章先聊聊最基础也最容易被忽视的一块单机多卡环境下的通信与显存管理。这是所有分布式训练的起点也是坑最密集的地方。2. 单机多卡的第一道坎NCCL通信为什么总在超时2.1 从一次真实的NCCL timeout说起那是我第一次正经配多卡训练。四张V100NVLink互联硬件条件不差。代码用的是PyTorch的DistributedDataParallel启动脚本照抄的官方示例torchrun --nproc_per_node4 train.py。前几个epoch一切正常到第7个epoch左右突然卡住然后报错NCCL timeout after 1800000 ms三十分钟的超时说明通信完全断了。我第一反应是硬件问题换了线、换了卡、甚至换了台机器问题依旧。后来把日志级别调到NCCL_DEBUGINFO才发现真正的问题四张卡在AllReduce时有一张卡的通信数据量和另外三张对不上。原因出在我的数据加载器上。我用了DistributedSampler但数据集长度不能被4整除最后一个batch在每个rank上的大小不一致。PyTorch的DDP默认要求所有进程的输入张量形状完全一致否则AllReduce操作就会hang住。这个问题在单卡上永远不会出现一上多卡就暴露了。提示DistributedSampler有个drop_last参数设为True可以丢弃最后一个不完整的batch保证各rank数据量一致。但如果你不想丢数据就得手动补齐或者用join模式。2.2 NCCL到底在干什么为什么它这么容易出问题NCCLNVIDIA Collective Communications Library是NVIDIA搞的一套集合通信库专门为GPU之间的通信做优化。它支持的通信原语包括AllReduce、Broadcast、Reduce、AllGather、ReduceScatter等。在DDP训练里最核心的是AllReduce——每个GPU算完自己的梯度后需要把所有GPU的梯度加起来求平均然后再分发给每个GPU。这个过程听起来简单但实现起来极其复杂。NCCL要根据你的硬件拓扑NVLink、PCIe、InfiniBand自动选择最优的通信路径。比如四张卡通过NVLink全互联那AllReduce就可以走ring或者tree算法直接在GPU之间传数据不经过CPU。但如果拓扑是混合的比如两张卡NVLink互联另外两张通过PCIe那NCCL就得做分层通信先组内reduce再跨组通信。问题就出在这个“自动选择”上。NCCL的拓扑探测有时候会误判尤其是在虚拟化环境或者容器里。我遇到过好几次容器里nvidia-smi topo -m显示的拓扑和宿主机上不一致导致NCCL选了错误的通信路径性能直接腰斩。排查这类问题的标准流程是这样的先看拓扑nvidia-smi topo -m确认GPU之间的连接方式。再看NCCL实际用了什么NCCL_DEBUGINFO日志里会打印它选择的算法和通道。如果发现NCCL选了PCIe而不是NVLink可以强制指定NCCL_P2P_LEVELNVL或者NCCL_TOPO_FILE手动指定拓扑文件。容器环境下确保/dev/nvidia*设备节点和/sys下的拓扑信息都正确挂载。2.3 几个我踩过的NCCL配置坑第一个坑NCCL_SOCKET_IFNAME没设对。多机训练时NCCL需要走网络通信。如果你的机器有多张网卡NCCL可能选到一张慢的或者不通的网卡上。我见过最离谱的情况是NCCL选了一张管理网卡带宽只有1Gbps而数据网卡是100Gbps的InfiniBand。设置NCCL_SOCKET_IFNAMEib0或者NCCL_SOCKET_IFNAMEeth0可以强制指定。第二个坑NCCL_IB_DISABLE的默认值。在某些环境下NCCL会尝试用InfiniBand但实际没有IB网络导致初始化超时。这时候需要设NCCL_IB_DISABLE1强制走Socket。但反过来如果你确实有IB网络千万别设这个否则性能损失巨大。第三个坑共享内存不够。NCCL在单机多卡通信时会用共享内存做中转。如果/dev/shm太小比如Docker默认只有64MB大模型训练时就会报错。解决办法是启动容器时加--shm-size1g或者更大。第四个坑NCCL_ALGO和NCCL_PROTO的强制指定。大多数时候让NCCL自动选就行但在某些特定拓扑下手动指定Ring算法或者LL协议能带来10%以上的性能提升。这个需要结合具体硬件做benchmark没有万能公式。2.4 一个实用的NCCL性能测试方法在正式跑训练之前我习惯先用NCCL自带的nccl-tests跑一遍带宽测试。这个工具能告诉你在当前硬件拓扑下AllReduce的实际带宽是多少。如果测出来的带宽远低于理论值比如NVLink理论带宽300GB/s实测只有50GB/s那说明配置有问题先别急着跑训练。# 编译nccl-tests git clone https://github.com/NVIDIA/nccl-tests.git cd nccl-tests make MPI1 MPI_HOME/usr/lib/x86_64-linux-gnu/openmpi # 跑AllReduce带宽测试8张卡从8M到128M的消息大小 ./build/all_reduce_perf -b 8M -e 128M -f 2 -g 8输出里重点看algbw算法带宽和busbw总线带宽。busbw更能反映实际通信效率因为它考虑了AllReduce的通信量放大。如果busbw只有理论值的30%以下那基本可以确定通信配置有问题。3. 显存不够用先搞清楚钱花在哪了3.1 显存占用的四大块你算对了吗很多人一看到OOMOut of Memory就想着加卡或者减batch size但从来没算过显存到底被什么吃掉了。一个典型的训练任务显存占用可以拆成四块占用类型说明典型占比模型参数权重和偏置10%-20%梯度反向传播算出来的梯度10%-20%优化器状态Adam的momentum和variance20%-40%激活值前向传播的中间结果30%-60%以Adam优化器为例每个参数需要存参数本身4字节、梯度4字节、一阶矩4字节、二阶矩4字节加起来是16字节。一个10亿参数的模型光这四项就是16GB。再加上激活值一张32GB的V100可能连batch size8都跑不起来。激活值是大头也是最容易被忽视的。Transformer类模型的激活值占用和序列长度、batch size、层数都成正比。我见过一个BERT-large的训练batch size16、序列长度512激活值占了将近20GB。这时候减batch size是最直接的办法但会影响收敛。3.2 梯度累积用小batch模拟大batch梯度累积Gradient Accumulation是个很实用的技巧。它的思路是用小的micro-batch跑多次前向和反向把梯度攒起来等攒够了再更新一次参数。这样显存占用是按micro-batch算的但效果上等价于大batch。# 梯度累积示例 accumulation_steps 4 optimizer.zero_grad() for i, (inputs, labels) in enumerate(dataloader): outputs model(inputs) loss criterion(outputs, labels) loss loss / accumulation_steps # 损失要除以累积步数 loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意那个loss / accumulation_steps很多人忘了除导致梯度被放大了accumulation_steps倍训练直接发散。另外如果你用了DDP梯度累积和DDP的配合需要小心因为DDP默认每次backward都会触发AllReduce。正确的做法是用model.no_sync()上下文管理器只在最后一次backward时才同步梯度。for i, (inputs, labels) in enumerate(dataloader): with model.no_sync() if (i 1) % accumulation_steps ! 0 else contextlib.nullcontext(): outputs model(inputs) loss criterion(outputs, labels) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()3.3 混合精度训练省显存还能提速混合精度Mixed Precision是我最推荐的一个优化手段。它的核心思想是前向和反向用FP16或BF16计算但参数更新用FP32。这样显存占用能降30%-50%速度还能提升20%-30%。PyTorch里用torch.cuda.amp就能实现from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, labels in dataloader: optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler的作用是防止FP16下梯度下溢。FP16的最小正数是6e-8左右很多小梯度直接变成0了。GradScaler会先把loss放大一个系数算完梯度再缩回去保证梯度不丢失。但混合精度也有坑。某些操作在FP16下数值不稳定比如softmax、layer norm、loss计算。这些地方需要强制用FP32。PyTorch的autocast会自动处理大部分情况但如果你自己写了自定义的CUDA kernel就得手动指定。注意BF16比FP16更稳因为它的指数位和FP32一样多动态范围大基本不会下溢。如果你的卡支持BF16A100及以上优先用BF16可以省掉GradScaler。3.4 激活重计算用时间换空间激活重计算Activation Checkpointing是另一个省显存的大招。它的思路是前向传播时不保存中间激活值反向传播时重新算一遍。这样显存占用能从O(n)降到O(sqrt(n))代价是计算量增加30%左右。PyTorch里用torch.utils.checkpointfrom torch.utils.checkpoint import checkpoint class MyModel(nn.Module): def forward(self, x): x checkpoint(self.layer1, x) x checkpoint(self.layer2, x) return x哪些层适合做checkpoint计算量大但显存占用也大的层比如Transformer的attention层。计算量小的层做checkpoint不划算重算的开销可能比省下的显存更值。我实测过一个13B参数的模型不开checkpoint需要8张A100 80GB开了之后4张就够了。虽然单步训练时间增加了25%但能跑起来比什么都强。4. 分布式策略选型DDP、FSDP还是DeepSpeed4.1 DDP最简单也最常用的起点DDPDistributed Data Parallel是PyTorch原生的数据并行方案。它的逻辑很直接每个GPU持有一份完整的模型副本各自跑不同的数据反向传播时通过AllReduce同步梯度。DDP的优点是简单、稳定、调试方便。代码改动量极小基本上把model nn.DataParallel(model)换成model DDP(model)就行。但它的缺点也很明显每个GPU都要存一份完整的模型、梯度和优化器状态。一个10B参数的模型用Adam训练光模型相关状态就要160GB单卡根本放不下。所以DDP的适用场景是模型能单卡放下只是想加速训练。一般来说参数量在1B以下的模型用DDP就够了。4.2 FSDPPyTorch的亲儿子但坑不少FSDPFully Sharded Data Parallel是PyTorch后来推出的思路是借鉴了DeepSpeed的ZeRO。它把模型参数、梯度和优化器状态都切分到各个GPU上每个GPU只存一部分。这样显存占用能大幅降低理论上可以训练任意大的模型。FSDP的使用比DDP复杂不少需要配置auto_wrap_policy、sharding_strategy等参数。我踩过的一个坑是auto_wrap_policy设得不对导致模型被切得太碎通信开销爆炸训练速度比DDP还慢。from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy auto_wrap_policy functools.partial( size_based_auto_wrap_policy, min_num_params1e6 ) model FSDP(model, auto_wrap_policyauto_wrap_policy)min_num_params1e6的意思是参数量超过100万的层会被单独wrap成一个FSDP unit。这个值需要根据模型结构调太大则切分不够细显存省不下来太小则通信太频繁速度上不去。另一个坑是FSDP和混合精度的配合。FSDP默认用FP32存参数如果你要用混合精度需要额外配置MixedPrecision策略。而且FSDP的checkpoint保存和加载也和DDP不一样需要用FSDP.state_dict_type上下文管理器。4.3 DeepSpeed功能最全但学习曲线最陡DeepSpeed是微软搞的功能极其丰富ZeRO的三个阶段、混合精度、梯度累积、激活重计算、MoE支持等等。它的ZeRO-3阶段和FSDP类似也是把什么都切分。但DeepSpeed的配置是通过一个JSON文件驱动的灵活度很高但也意味着配置项多到让人头大。{ train_batch_size: 64, gradient_accumulation_steps: 4, fp16: { enabled: true }, zero_optimization: { stage: 3, offload_optimizer: { device: cpu }, offload_param: { device: cpu } } }上面这个配置开了ZeRO-3还把优化器状态和参数都offload到CPU。这样显存占用极低但速度会慢很多因为CPU和GPU之间的数据传输成了瓶颈。offload适合那种“模型太大不加offload根本跑不起来慢就慢点吧”的场景。DeepSpeed的坑主要集中在配置文件的参数上。比如train_batch_size是全局batch size不是单卡的gradient_accumulation_steps要和train_batch_size配合好否则会报错。还有zero_optimization的stage 2和stage 3在通信模式上完全不同stage 3的通信量更大对网络要求更高。4.4 选型建议从简单到复杂别一步到位我的建议是先用DDP跑通再根据显存瓶颈逐步升级。如果DDP能跑就别折腾FSDP和DeepSpeed。如果DDP显存不够先试混合精度和激活重计算这两个改动小、收益大。还不够再上FSDP。FSDP也搞不定再考虑DeepSpeed的ZeRO-3加offload。别一上来就DeepSpeed ZeRO-3加CPU offload那样调试成本极高而且性能往往不如预期。我见过太多团队花了两周配DeepSpeed最后发现用DDP加混合精度就能解决问题。5. 那些文档里不会写的实操细节5.1 随机种子分布式训练下的可复现性单卡训练时设个torch.manual_seed(42)就能复现。但多卡训练下事情没这么简单。每个进程的随机种子必须不同否则所有GPU会生成相同的随机数数据增强就失效了。但不同又要可复现所以通常用seed rank的方式def set_seed(seed, rank): random.seed(seed rank) np.random.seed(seed rank) torch.manual_seed(seed rank) torch.cuda.manual_seed_all(seed rank)另外DataLoader的worker_init_fn也要设种子否则每个epoch的数据顺序会变。还有CUDA的卷积算法选择torch.backends.cudnn.deterministic True可以保证每次选同样的算法但会牺牲一点性能。5.2 日志与监控别等挂了才去看分布式训练最怕的就是某个rank挂了但其他rank还在傻等。所以日志一定要每个rank单独打并且带上rank编号。我习惯用logging模块每个rank写自己的日志文件import logging def setup_logging(rank): logger logging.getLogger(frank_{rank}) handler logging.FileHandler(ftrain_rank_{rank}.log) logger.addHandler(handler) logger.setLevel(logging.INFO) return logger监控方面nvidia-smi只能看个大概真正有用的是nvitop或者gpustat能看到每个进程的显存和利用率。如果要长期监控可以上Prometheus加Grafana把GPU指标采集起来。还有一个容易被忽视的点NCCL的日志。设NCCL_DEBUGWARN可以在出问题时看到警告信息设NCCL_DEBUGINFO能看到详细的通信过程。但INFO级别日志量很大平时别开出问题了再开。5.3 checkpoint保存别让几天的训练白费分布式训练的checkpoint保存比单卡复杂。DDP下只需要rank 0保存就行因为所有rank的模型是一样的。但FSDP和DeepSpeed下每个rank只存了一部分参数需要特殊的保存逻辑。FSDP用FSDP.state_dict_typefrom torch.distributed.fsdp import StateDictType, FullStateDictConfig with FSDP.state_dict_type( model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(rank0_onlyTrue, offload_to_cpuTrue) ): state_dict model.state_dict() if rank 0: torch.save(state_dict, checkpoint.pt)DeepSpeed更简单它有自己的save_checkpoint方法会自动处理分片保存和恢复。我踩过的一个坑是checkpoint保存时没设offload_to_cpuTrue结果rank 0的显存直接爆了。因为把所有分片的参数gather到rank 0上显存需求瞬间翻倍。加上offload到CPU就没事了。5.4 故障恢复训练中断了怎么续大规模训练跑几天几夜是常事中间难免遇到硬件故障、网络抖动、甚至机房断电。所以断点续训是必须的。除了保存模型参数还要保存优化器状态、学习率调度器的状态、当前的epoch和step。checkpoint { model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, step: global_step, loss: loss.item(), } torch.save(checkpoint, fcheckpoint_epoch_{epoch}.pt)恢复的时候把这些状态都load回去然后从断点继续。注意DataLoader的状态也要恢复否则数据顺序会乱。PyTorch的DataLoader没有内置的状态保存需要自己记录sampler的epoch。还有一个经验checkpoint别只存一个存最近三个。有时候最新的checkpoint可能是在模型发散时保存的回滚到上一个能救命。6. 从单机多卡到多机多卡下一步该准备什么单机多卡跑通了下一步自然是多机多卡。但多机多卡引入的新问题比单机多卡多得多。首先是网络单机内是NVLink和PCIe多机之间是InfiniBand或者RoCE。IB的配置极其繁琐从子网管理器到IP over IB每一步都可能出问题。RoCE虽然简单点但对网络设备的要求高丢包率稍微大一点性能就崩。其次是存储多机训练时数据加载成了瓶颈。单机时数据在本地SSD上多机时得走网络文件系统NFS或者分布式存储如Lustre、Ceph。NFS在小规模下还行规模一大就成瓶颈。我见过一个集群100张卡同时读NFS上的小文件直接把NFS服务器打挂了。最后是调度多机多卡通常意味着你要和別人共享集群。这时候Kubernetes或者Slurm就派上用场了。K8s的GPU调度需要装NVIDIA的设备插件Slurm则需要配GRESGeneric Resource Scheduling。这两个我后面会单独开一章讲坑同样不少。但不管多机多卡怎么复杂单机多卡都是基础。如果单机多卡的通信和显存都没搞明白上多机只会让问题更难排查。所以我的建议是先把单机多卡吃透把NCCL调优、显存管理、分布式策略选型这三件事做扎实再往多机走。这个系列我会继续写下去下一章打算聊多机多卡环境下的网络配置和存储优化包括IB和RoCE的选型对比、NFS的性能调优、以及K8s和Slurm的实战配置。如果你也在搞AI-Infra或者正准备往这个方向转欢迎一起交流踩坑经验。毕竟这条路一个人走容易掉坑里一群人走至少能互相拉一把。