引子先看看面试官会怎么问在进入正文之前先来看几道真实面试中出现过的高频题。读完本章后你再回头看它们会发现每一道都指向分布式训练的核心本质。面试题 1你的 batch size 设了 8192单卡显存放不下除了换更大显存的卡还有什么办法面试题 28 张卡训练loss 震荡得比单卡严重得多batch size 明明一模一样为什么面试题 3PyTorch DDP 启动之后训练速度跟单卡差不多甚至更慢你会从哪些角度排查面试题 4你知道 ZeRO-3 和模型并行有什么区别吗为什么 ZeRO-3 被称为内存最优3.1 为什么需要分布式训练3.1.1 单卡的瓶颈现代深度学习模型参数量呈指数级增长。GPT-3 有 1750 亿参数按照混合精度训练FP16计算仅参数就需要 350 GB 显存175B x 2 bytes而这还不包括优化器状态和梯度。目前消费级 GPU 显存上限约为 80 GBH100单卡完全无法承载。分布式训练的核心动机可以归纳为三个维度维度单卡瓶颈分布式解决方式显存容量模型放不进单卡显存将模型参数 / 梯度 / 优化器状态分散到多卡训练速度单卡算力有限训练周期过长多卡并行计算线性加速batch size单卡可承载 batch size 太小BN 统计不准确多卡累积更大全局 batch size3.2 三大并行策略数据并行、模型并行、流水线并行3.2.1 数据并行Data Parallelism定义将训练数据切分成多个 micro-batch每个 GPU 持有完整的模型副本各自独立前向 反向然后通过All-Reduce同步梯度。GPU 0: [data_0] - Forward - Backward - grad_0 GPU 1: [data_1] - Forward - Backward - grad_1 GPU 2: [data_2] - Forward - Backward - grad_2 GPU 3: [data_3] - Forward - Backward - grad_3 | All-Reduce | 平均梯度 - 更新模型优点实现简单PyTorch DDP 两三行代码搞定计算负载天然均衡。缺点每张卡都需要完整模型副本显存开销随模型尺寸线性增长不适合超大模型。适用场景模型能放进单卡显存、但需要加速训练的场景。3.2.2 模型并行Model Parallelism / Tensor Parallelism定义将单个算子的计算拆分到多张卡上执行。最常见的实现方式是将矩阵乘法按行或列切分。以全连接层 Y XW 为例列切分Column-wise Partition# 单卡版本YX W# X: [B, D], W: [D, H]# 模型并行 - 列切分2卡W_0W[:,:H//2]# [D, H/2]W_1W[:,H//2:]# [D, H/2]Y_0X W_0# GPU 0 计算Y_1X W_1# GPU 1 计算Yconcat(Y_0,Y_1,dim-1)# All-Gather 合并优点能训练单卡放不下的超大层如 FFN hidden dim 为 65536 的 Transformer。缺点通信极其密集——每层 forward/backward 都需要 All-Reduce 或 All-Gather通信量和计算量成正比。注意在实际工程中模型并行通常指张量并行Tensor Parallelism, TP与下文流水线并行区分。3.2.3 流水线并行Pipeline Parallelism定义将模型按层切分成多个 stage每个 GPU 负责一个连续的子网络段数据以 micro-batch 形式在 stage 之间流水传递。GPU 0: Layers 1-4 - Forward - (输出传给 GPU 1) GPU 1: Layers 5-8 - Forward - (输出传给 GPU 2) GPU 2: Layers 9-12 - Forward - (输出传给 GPU 3) GPU 3: Layers 13-16- Forward - loss朴素流水线Naive Pipeline存在严重的 GPU 空闲问题——当 GPU 0 计算时GPU 1/2/3 都在空等。现代流水线并行使用微批次micro-batch技术来缓解# 以 GPipe 风格为例defpipeline_forward(model_stages,micro_batches):outputs[]formbinmicro_batches:hiddenmbforstageinmodel_stages:hiddenstage(hidden)outputs.append(hidden)returnoutputs改进方案 1 - GPipeGoogle将大 batch 拆成 micro-batch依次推入流水线。缺点仍然存在流水线气泡pipeline bubble——流水线填充和排空阶段的 GPU 利用率不满。改进方案 2 - 1F1BOne-Forward-One-BackwardNVIDIA前向计算和反向计算交替进行尽早释放中间激活内存降低显存峰值。Google 的 T5 和 NVIDIA 的 Megatron-LM 均采用此策略。3.2.4 三大策略对比总表维度数据并行 (DP)张量/模型并行 (TP)流水线并行 (PP)拆分维度按数据拆分按算子拆分按层拆分通信频率每步一次 All-Reduce每层多次 All-Reduce每个 micro-batch 一次 P2P通信量O(模型大小)O(激活值大小)O(激活值大小)GPU 利用率高无气泡高同步执行中有气泡实现难度低高中显存节省不节省每卡全量线性降低单卡参数线性降低单卡层数典型框架PyTorch DDP, HorovodMegatron-LM, DeepSpeedGPipe, PipeDream面试官追问链Q为什么模型并行TP的通信量比数据并行大得多A数据并行每步只通信一次梯度大小为模型参数量通信和计算是串行的。TP 每层的前向和反向都要做 All-Reduce对于一个 L 层 Transformer通信次数是数据并行的 O(L) 倍。在千卡集群上TP 的通信开销往往成为瓶颈因此 TP 通常只在节点内node-level8 卡 NVLink使用跨节点使用 DP 或 PP。Q能不能三种并行一起用A可以而且这是大模型训练的标配——3D 并行3D Parallelism。典型的部署方案节点内用 TP利用 NVLink 高带宽节点间用 PP减少跨节点通信最外层用 DPZeRO 优化。NVIDIA Megatron DeepSpeed 的组合就是这种思路。3.3 Ring-AllReduce分布式通信的基石3.3.1 为什么需要 All-Reduce在数据并行中每张卡反向传播后都有自己的梯度 g_i我们需要计算所有卡的梯度平均值并广播回每一张卡。这个过程就叫 All-Reduce。最朴素的实现是Parameter Server3.4 节详述——每个 worker 向 server 发送梯度server 求平均再广播。这种方式有严重的通信热点和带宽瓶颈。Ring-AllReduce由 Baidu 的 Andrew Gibiansky 在 2017 年提出彻底解决了这个问题。3.3.2 两阶段算法Ring-AllReduce 将 N 张 GPU 组成一个逻辑环形拓扑分两步完成阶段一Reduce-Scatter归约散射每张卡将梯度等分为 N 份N 卡数。在每一步每张卡向环中的下一个邻居发送自己的一个分块同时接收上一个邻居的分块做累加。经过 N-1 步后每张卡持有某一份的全局求和结果。初始状态N4每卡梯度等分4块 GPU 0: [g0_0, g0_1, g0_2, g0_3] GPU 1: [g1_0, g1_1, g1_2, g1_3] GPU 2: [g2_0, g2_1, g2_2, g2_3] GPU 3: [g3_0, g3_1, g3_2, g3_3] Step 1: GPU 0 发送 g0_0 - GPU 1GPU 1 发送 g1_1 - GPU 2 GPU 0 接收 g3_0累加得到 Sum_g_0 g0_0 g3_0 经过 N-13 步后 GPU 0 持有: Sum_g_0 - 所有卡 g_i_0 的累加 GPU 1 持有: Sum_g_1 GPU 2 持有: Sum_g_2 GPU 3 持有: Sum_g_3阶段二All-Gather全收集每张卡将自己持有的全局求和结果沿环广播出去。同样经过 N-1 步每张卡拿到所有分块的全局求和得到完整的梯度总和除以 N 即为均值。Step 1: GPU 0 发送 Sum_g_0 - GPU 1GPU 1 发送 Sum_g_1 - GPU 2 经过 N-13 步后 GPU 0: [Sum_g_0, Sum_g_1, Sum_g_2, Sum_g_3] GPU 1: [Sum_g_0, Sum_g_1, Sum_g_2, Sum_g_3] GPU 2: [Sum_g_0, Sum_g_1, Sum_g_2, Sum_g_3] GPU 3: [Sum_g_0, Sum_g_1, Sum_g_2, Sum_g_3] 最后每卡grad Sum_g / N3.3.3 通信复杂度分析总通信量每张卡发送和接收的数据量为2 x (N-1) x (参数量/N)约为2 x 模型参数量。总通信量与卡数 N 无关这是 Ring-AllReduce 最优雅的性质。对比 Parameter ServerParameter Server 的通信量为2 x N x 模型参数量随卡数线性增长。Ring-AllReduce 将通信热点均匀分散到所有 GPU 间链路理论上可以实现线性加速。3.3.4 实现示例伪代码defring_all_reduce(gradients,world_size,rank,send_buf,recv_buf):flat_gradflatten(gradients)chunk_sizelen(flat_grad)//world_size chunks[flat_grad[i*chunk_size:(i1)*chunk_size]foriinrange(world_size)]# Reduce-Scatterrecv_chunkrecv_buf[:chunk_size]forstepinrange(world_size-1):send_rank(rank-step)%world_size recv_rank(rank-step-1)%world_size send_buf[:]chunks[(rank-step)%world_size]send_recv(send_buf,send_rank,recv_chunk,recv_rank)chunks[rank]chunks[rank]recv_chunk# All-Gatherforstepinrange(world_size-1):send_rank(rank-step)%world_size recv_rank(rank-step-1)%world_size send_buf[:]chunks[(rank-step)%world_size]send_recv(send_buf,send_rank,recv_chunk,recv_rank)chunks[(rank-step-1)%world_size]recv_chunkreturnunflatten(concatenate(chunks),gradients)面试官追问链QRing-AllReduce 和 Tree-AllReduce 有什么区别为什么现在主流用 RingATree-AllReduce 采用树形拓扑父节点负责累加子节点的梯度。它的通信量是O(N logN)而 Ring 是O(N)。当卡数达到几百上千时Ring 的带宽优势明显。此外 Ring 的拓扑简单只需要相邻节点连通对物理网络拓扑的适应性强。QRing-AllReduce 在异构网络不同带宽下会有什么问题A环形的总带宽取决于最慢的一跳slowest link。如果集群中某两台机器的跨节点带宽远低于 NVLink 带宽整个 Ring 就会被拖慢到跨节点带宽的水平。这也是实践中会将节点内用 Ring高速、跨节点用分层 All-ReduceHierarchical All-Reduce的原因。3.4 Parameter Server 架构3.4.1 核心思想Parameter ServerPS是由 Alex Smola 和李沐等人在 2010 年提出的分布式训练架构在深度学习早期2013-2018被广泛使用。它的核心思想是中心化-去中心化结合Server 节点维护全局参数接收 worker 的梯度更新参数分发最新参数。Worker 节点从 server 拉取最新参数在本地数据上做前向和反向将梯度推送给 server。3.4.2 异步 vs 同步模式工作方式优点缺点同步 (Sync)所有 worker 算完梯度后一起更新梯度一致性高收敛稳定慢节点拖慢整体straggler problem异步 (Async)worker 算完立即 push不等待其他 worker训练速度快无 straggler梯度 stale陈旧收敛不稳定可能不收敛半同步 (SSP)允许 worker 之间相差 k 步以内的异步在速度和稳定性之间取得平衡实现复杂k 值难调3.4.3 优缺点分析优点架构直观易于理解支持大规模稀疏模型推荐系统等server 可灵活分片异步模式在大规模推荐系统中仍有应用缺点Server 节点成为通信瓶颈带宽和吞吐量同步模式下 straggler 问题严重在稠密模型如大语言模型上Ring-AllReduce 通信效率远超 PS现状在 CV / NLP 等稠密模型领域PS 已被 Ring-AllReduce 全面取代。但在推荐系统如 Click-Through Rate 预测这类大规模稀疏模型领域PS 的灵活性和异步能力仍有优势。面试官追问链QParameter Server 的异步梯度更新为什么会导致收敛问题A假设 worker 0 在时间 t 拉取了参数 w_t 开始计算而 worker 1 在时间 t1 用 w_t 算完梯度 push 更新成 w_{t1}。当 worker 0 在 t2 时刻算完并 push 时它的梯度是基于 w_t 计算的但此时参数已经是 w_{t2}。这个陈旧梯度stale gradient可能与当前参数方向完全不同导致 loss 震荡甚至发散。Q如何缓解异步 PS 的梯度陈旧问题A常见方法有①梯度裁剪限制单步更新幅度②动量校正对陈旧梯度降低学习率③SSPStale Synchronous Parallel限制异步步数上限④DC-ASGDDelay-Compensated ASGD对陈旧梯度做泰勒展开补偿。3.5 ZeRO内存优化的革命性方案3.5.1 ZeRO 要解决什么问题在数据并行中每张卡都持有一份完整的模型状态包括组件容量估算以 1.5B 参数、Adam、FP16 为例参数 (Parameters)1.5B x 2 bytes 3 GB梯度 (Gradients)1.5B x 2 bytes 3 GB优化器状态 (Optimizer States)Adam 维护 two momentum1.5B x 4 x 2 12 GB总计单卡18 GB如果使用数据并行训练 1.5B 模型每张卡都需要 18 GB 来存储模型状态加上激活值activation和临时缓冲区轻松超过 40 GB。这极大地限制了数据并行的扩展能力。ZeROZero Redundancy Optimizer由微软 DeepSpeed 团队在 2019 年提出核心思想是既然每张卡存的东西一模一样为什么不让每张卡只存一部分用的时候再取动态 Loss Scaling训练开始时 scale 设为一个较大值如 65536。如果连续 N 步未出现梯度溢出inf/nan则 scale 翻倍如果出现溢出则跳过该步并减半 scale。这使得 scale 自适应地保持在最优范围。机制三计算精度控制Tensor Core 要求NVIDIA V100 及以后的 GPU 的 Tensor Core 要求输入为 FP16。混合精度训练可以自动利用 Tensor Core 获得 8-16 倍的矩阵乘法加速。逐操作精度策略矩阵乘法GEMM/ 卷积FP16 计算Tensor Core逐点操作ReLU, BN 等FP16梯度更新 / loss 计算 / reduction 操作FP32Softmax 中 exp 求和FP32防止精度溢出3.6.3 BF16FP16 的进化版特性FP16BF16 (bfloat16)指数位5 bits8 bits同 FP32尾数位10 bits7 bits表示范围/-65,504/-3.4e38同 FP32精度高尾数多低尾数少适用有 Loss Scaling 可训练无需 Loss ScalingBF16 由 Google Brain 提出保留了 FP32 的指数范围因此天然的不会梯度下溢。代价是尾数位少 3 位舍入误差略大。目前 H100 和 A100 原生支持 BF16。面试官追问链QLoss Scaling 的 scale 值怎么确定为什么不是越大越好Ascale 太大虽然可以防止下溢但过大的 scale 会使梯度在反向传播中溢出 FP16 上限65504产生 inf。动态 loss scaling 就是在这种 trade-off 中自动寻找平衡的策略以指数增长逼近不溢出的最大值。Q混合精度训练中BN 层计算时需要注意什么ABN 层的统计量mean 和 variance需要足够精度来保证数值稳定实践中通常强制 BN 在 FP32 下计算。PyTorch 的autocast会自动为 BN 层切换到 FP32。QBF16 的尾数只有 7 位会不会影响大模型的收敛质量A目前的研究和实践表明对于 LLM 训练BF16 的精度损失几乎可以忽略。原因是参数更新量grad x lr通常远小于参数本身每次更新只改变较低的有效位数。而 BF16 保留了 FP32 的指数范围不会出现梯度消失问题。实际上GPT-3、LLaMA、PaLM 等主流大模型均使用 BF16 训练。3.7 PyTorch DDP 实现详解3.7.1 DDP 的核心设计PyTorch DistributedDataParallelDDP是目前最流行的分布式训练框架。它之所以优于老版DataParallel关键在于两个设计选择对比项DataParallel (DP)DistributedDataParallel (DDP)进程模型单进程多线程多进程每个 GPU 一个进程梯度同步每步串行 gather-scatter后台异步 All-ReduceGIL 影响受 Python GIL 限制无 GIL 问题扩展性4 卡以上性能急剧下降可扩展到数百卡启动方式自动封装需 torchrun 或 mp.spawn3.7.2 DDP 的标准使用流程importosimporttorchimporttorch.distributedasdistimporttorch.multiprocessingasmpfromtorch.nn.parallelimportDistributedDataParallelasDDPdefsetup(rank,world_size):os.environ[MASTER_ADDR]localhostos.environ[MASTER_PORT]12355dist.init_process_group(backendnccl,init_methodenv://,rankrank,world_sizeworld_size)defcleanup():dist.destroy_process_group()deftrain_ddp(rank,world_size):setup(rank,world_size)modelMyModel().to(rank)ddp_modelDDP(model,device_ids[rank],find_unused_parametersFalse,bucket_cap_mb25)datasetMyDataset()samplertorch.utils.data.distributed.DistributedSampler(dataset,num_replicasworld_size,rankrank,shuffleTrue)dataloadertorch.utils.data.DataLoader(dataset,batch_size32,samplersampler)optimizertorch.optim.Adam(ddp_model.parameters(),lr1e-3)forepochinrange(10):sampler.set_epoch(epoch)fordata,targetindataloader:data,targetdata.to(rank),target.to(rank)optimizer.zero_grad()outputddp_model(data)losstorch.nn.functional.cross_entropy(output,target)loss.backward()optimizer.step()cleanup()if__name____main__:world_sizetorch.cuda.device_count()mp.spawn(train_ddp,args(world_size,),nprocsworld_size)# 推荐使用 torchruntorchrun--nproc_per_node8train_ddp.py3.7.3 DDP 的工作流程梯度同步是如何实现的DDP 在构造时对模型参数注册了 autograd hook梯度钩子。反向传播时hook 在参数的梯度计算完成后自动触发执行异步 All-Reduce。为减少通信次数DDP 会将参数分组为梯度桶gradient bucket每个桶包含约 25 MB 的参数。当一个桶中所有参数的梯度就绪后DDP 立即对该桶启动 All-Reduce而不是等到所有梯度都算完。反向传播过程 Layer 3 梯度就绪 - 放入 Bucket 3 - Bucket 3 满 - 异步 All-Reduce Layer 2 梯度就绪 - 放入 Bucket 2 - Bucket 2 满 - 异步 All-Reduce Layer 1 梯度就绪 - 放入 Bucket 1 - Bucket 1 满 - 异步 All-Reduce All-Reduce 与后续计算结果重叠隐藏通信延迟3.7.4 DDP 性能调优要点# 1. 调整 Bucket 大小DDP(model,device_ids[rank],bucket_cap_mb25)# 2. 将 BN 替换为 SyncBN跨卡同步 BN 统计量ddp_modeltorch.nn.SyncBatchNorm.convert_sync_batchnorm(ddp_model)# 3. 禁用不必要的梯度同步gradient accumulation no_syncwithddp_model.no_sync():for_inrange(accumulation_steps-1):outputddp_model(data)loss.backward()outputddp_model(data)loss.backward()optimizer.step()3.7.5 DDP 常见问题与排查现象可能原因解决方案训练速度与单卡持平数据加载瓶颈CPU 来不及喂数据增大 num_workers使用 pin_memoryTrue多卡 loss 曲线与单卡不一致全局 batch size 变化未调整 lr遵循线性缩放法则lr x sqrt(N) 或 lr x NGPU 利用率不足 50%通信开销太大模型小 / 带宽低增大 bucket_cap_mb减少通信次数NCCL timeout某卡掉队慢节点或 NCCL 错误检查 NCCL_IB_TIMEOUT排查网络健康OOM 但单卡可以跑DDP 需要额外显存存 Bucket 缓冲区开启 gradient_checkpointing面试官追问链QDDP 中find_unused_parametersTrue是用来做什么的A当模型的某些参数在反向传播中没有接收到梯度时如使用了冻结层、条件分支、模型输出的部分不参与 loss 计算等DDP 的梯度同步 hook 会因找不到这些梯度而报错。设置find_unused_parametersTrue会让 DDP 预先遍历计算图找出哪些参数会被使用跳过未使用的参数的梯度同步。代价是第一次反向传播时会多做一次图遍历略微拖慢速度。QPyTorch 2.0 的torch.compile和 DDP 能一起用吗A可以。PyTorch 2.0 引入的torch.compile会通过 TorchDynamo 将模型编译为优化后的 FX 图DDP 的 bucket 通信和 autograd hook 不受影响。实际测试表明torch.compile DDP 可以在 GPT 类模型上获得 10-30% 的额外加速。推荐做法先torch.compile再包 DDP。modeltorch.compile(model)ddp_modelDDP(model,device_ids[rank])Q如果我有 4 台机器每台 8 卡DDP 怎么部署A使用torchrun多节点启动torchrun--nnodes4--nproc_per_node8--node_rank$RANK--master_addr192.168.1.100--master_port12345train_ddp.py其中 RANK 分别为 0, 1, 2, 3。NCCL 会自动检测网络拓扑在节点内使用 NVLink跨节点使用 InfiniBand 或 RoCE。本章总结知识点核心要点面试考法三大并行策略数据并行数据拆分、张量并行算子拆分、流水线并行层拆分对比它们的通信量、显存节省、适用场景Ring-AllReduceReduce-Scatter All-Gather 两步通信量 O(2psi) 与卡数无关推导通信量公式对比 PS 和 Tree-AllReduceParameter Server中心化架构同步/异步/半同步三种模式讨论异步梯度陈旧问题和缓解方法ZeRO 三段Stage 1: 优化器状态分片; Stage 2: 梯度分片; Stage 3: 参数分片解释内存节省原理和通信量变化混合精度训练FP16 计算 FP32 主副本 Loss Scaling 动态 scale分析梯度下溢问题对比 FP16/BF16/FP32PyTorch DDP多进程 NCCL All-Reduce 梯度桶 autograd hook写出完整训练脚本排查性能瓶颈3D 并行节点内 TP 节点间 PP 全局 DP/ZeRO设计千卡集群的并行方案核心公式速记公式含义psi 模型参数量基础计量单位单卡显存约等于 2psi (参数梯度) 4psi (Adam states) 激活值混合精度训练显存估算Ring-AllReduce 通信量 2psi与卡数无关ZeRO-3 通信量 3psi比基线多 50%线性缩放法则: lr_new lr_base x sqrt(N)全局 batch size 翻 N 倍时调整 lr流水线气泡比约等于 (P-1) / (MP-1)GPipe 的效率上限进阶阅读推荐Megatron-LMShoeybi et al., 2019模型并行 数据并行的工业级实现DeepSpeed ZeRORajbhandari et al., 2019ZeRO 系列论文内存优化的奠基之作Mixed Precision TrainingMicikevicius et al., 2018混合精度训练的原始论文Ring AllReduceGibiansky, 2017Baidu 的经典博客文章PyTorch DDP 源码解析GitHub pytorch/torch/distributed理解 Reducer.cpp