175B 参数的大模型不能放在一张 NPU 上——需要分布式。三种并行策略各有优劣数据并行DP简单但显存不降、张量并行TP通信密集但显存降得最多、流水线并行PP显存也降但有 bubbleGPU 空闲等待。torchtitan-npu 的组合策略DP 跨节点、TP 跨卡、PP 跨层三合一最大化吞吐。三种并行的本质数据并行 (DP) 每张卡有完整模型的副本数据分成 N 份 前向独立计算 → AllReduce 梯度 → 更新参数 显存1× model不降 通信每层 AllReduce 一次梯度O(model_size) 张量并行 (TP) 每张卡有模型权重的一部分按列或行切分 前向需要 AllReduce 中间激活 显存1/N × model降最多 通信每层 2× AllReduce 激活O(batch_size × hidden) 流水线并行 (PP) 每张卡有连续的几层layer 0-7 在卡 08-15 在卡 1... 前向micro-batch 流水线 显存1/PP × model按层降 通信层间传递激活O(batch_size × hidden)很少 缺点Pipeline Bubble——微批次之间的空闲等待组合策略DP8 × TP4 × PP2torchtitan-npu 的推荐配置DP 跨 8 节点、TP 跨 4 卡节点内、PP 分 2 段。总共 8×4×2 64 张 NPU。物理分布64 NPU 8 节点 × 8 卡/节点 节点 0 (卡 0-3)Layer 0-15, TP shard 0-3 节点 0 (卡 4-7)Layer 16-31, TP shard 0-3 节点 1 (卡 0-3)Layer 0-15, TP shard 0-3 节点 1 (卡 4-7)Layer 16-31, TP shard 0-3 ... 节点 7 (卡 0-3)Layer 0-15, TP shard 0-3 节点 7 (卡 4-7)Layer 16-31, TP shard 0-3 DP 组8 个节点各有一份完整模型 TP 组每 4 张卡切分一层 PP 组Layer 0-15 和 Layer 16-31 分成两段为什么这样切TP 的通信是卡间节点内 NVLink/NVSwitch800GB/sDP 的通信是节点间RoCE/IB 网络200GB/s。TP 放节点内高速、DP 放节点间低速——匹配物理带宽。Ascend C 实现第 1 步TP 的权重切分与激活通信TP 的核心把权重矩阵分片存储前向时每个 TP rank 算自己那部分然后通信合并。# torchtitan-npu/parallelism/tensor_parallel.pyclassTensorParallelLayer(nn.Module):TP 包装器自动切分权重并做激活通信def__init__(self,layer,tp_size,tp_rank):super().__init__()self.layerlayer self.tp_sizetp_size self.tp_ranktp_rankdefforward(self,x):# ColumnParallel输入广播到所有 TP rank# 输出 每 rank 算自己的 weight 部分# 不需要通信输入相同 → 各自独立算# RowParallel每 rank 独立输入# 输出需要 AllReduce所有 rank 的部分贡献相加# 通信量 hidden_dim × tp_size × batch_sizeifself.layer.is_column_parallel:returnself.layer(x)# 不需要通信else:local_outself.layer(x)# AllReduce 在 NCCL/HCCL 层面执行returndist.all_reduce(local_out,opdist.ReduceOp.SUM)完整的 TP 前向以 Attention 为例# Attention with TPQ/K/V 投影矩阵按列切分# Q_proj: [D, D] → [D, D/TP] — 每 rank 存储 1/TP 的列# 输出每 rank 算 Q_local x × Q_proj_local → shape [B, S, D/TP]defattention_tp(x,tp_size,tp_rank):# Q/K/V 投影每 rank 独立算Q_locallinear_tp(x,Q_weight_shard[tp_rank])# [B, S, D_head/TP]K_locallinear_tp(x,K_weight_shard[tp_rank])V_locallinear_tp(x,V_weight_shard[tp_rank])# Attention 计算每 rank 独立——Q/K/V 都分片但双线性 QK^T 不跨 TP 通信attn_localscaled_dot_product_attention(Q_local,K_local,V_local)# O 投影RowParallel → 输出需要 AllReduceO_locallinear_tp(attn_local,O_weight_shard[tp_rank])# [B, S, D]Odist.all_reduce(O_local)# ← 唯一一次通信returnO第 2 步PP 的 Micro-batch 调度PP 的问题卡 0 算完 micro-batch 0 后传给卡 1卡 1 才能开始——这段时间卡 0 空闲了。调度器用轮转填充减少 bubble。# torchtitan-npu/parallelism/pipeline_parallel.pyclassPipelineScheduler:1F1B (One-Forward-One-Backward) 调度器def__init__(self,num_microbatches,pp_rank,pp_size):self.num_microbatchesnum_microbatches# 微批次数量如 8self.pp_rankpp_rank self.pp_sizepp_size self.warmup_microbatchespp_size-pp_rank-1# 预热期微批次defschedule(self):返回 (step, is_forward, microbatch_id) 的序列schedule[]step0# 阶段 1Warm-up前向预热逐层传递# PP rank 0F0, F1, F2, F3, ..., F7# PP rank 1 F0, F1, F2, ..., F7formbinrange(self.num_microbatches):forppinrange(self.pp_size):ifmbpp:schedule.append((step,forward,mb-pp,pp))step1# 阶段 21F1B前向和反向交替# 前向填充全用完后F7, B0, F8(完), B1, B2, ...formbinrange(self.num_microbatches-1,-1,-1):forppinrange(self.pp_size-1,-1,-1):schedule.append((step,backward,mb,pp))step1returnschedule# 示例PP2, num_microbatches4# 时间轴# PP0: F0 F1 F2 F3 B3 B2 B1 B0# PP1: F0 F1 F2 F3 B3 B2 B1 B0# ↑ warmup bubble 1 个 micro-batch 时间Pipeline Bubble 的计算bubble_ratio (PP - 1) / num_microbatches PP2, microbatches4: bubble 1/4 25% PP2, microbatches8: bubble 1/8 12.5% PP4, microbatches8: bubble 3/8 37.5% PP8, microbatches8: bubble 7/8 87.5%经验法则PP 不能太大≤4micro-batch 不能太少≥PP×2。PP4 且 mb8 → bubble37.5% → 吞吐是 PP1 的 62.5%。第 3 步3D 并行的完整启动脚本# torchtitan-npu/examples/llama_3d_parallel.pyimporttorch_npufromtorchtitan_npu.parallelismimport(TensorParallelLayer,PipelineScheduler,FSDPWrapper)defsetup_3d_parallel():# DP: 节点间8 个节点dp_sizedist.get_world_size()//(tp_size*pp_size)# 64 / (4*2) 8dp_rankdist.get_rank()%dp_size# TP: 节点内卡间4 张卡tp_size4tp_rank(dist.get_rank()//dp_size)%tp_size# PP: 模型分两段Layer 0-15, Layer 16-31pp_size2pp_rank(dist.get_rank()//(dp_size*tp_size))%pp_size# 创建通信组dp_groupcreate_group(dp_rank,dp_size)# AllReduce 梯度tp_groupcreate_group(tp_rank,tp_size)# AllReduce 激活pp_groupcreate_group(pp_rank,pp_size)# P2P 激活传递returndp_group,tp_group,pp_groupdefcreate_llama_3d_parallel():dp_group,tp_group,pp_groupsetup_3d_parallel()# 构建模型2 段 PP每段 16 层每层用 TPifpp_rank0:layers[build_llama_layer(i)foriinrange(16)]else:layers[build_llama_layer(i)foriinrange(16,32)]# TP 包装每层wrapped_layers[TensorParallelLayer(layer,tp_sizetp_group.size(),tp_ranktp_group.rank())forlayerinlayers]# FSDP 包装DP 层面modelFSDPWrapper(nn.Sequential(*wrapped_layers),process_groupdp_group,sharding_strategyShardingStrategy.FULL_SHARD# ZeRO-3: 全部卸载)returnmodel# 训练循环modelcreate_llama_3d_parallel()optimizertorch.optim.AdamW(model.parameters(),lr3e-4)forstep,batchinenumerate(dataloader):# Pipeline 调度schedulerPipelineScheduler(num_microbatches8,pp_rankpp_rank,pp_size2)forstep_type,microbatch_idinscheduler.schedule():ifstep_typeforward:withtorch.npu.amp.autocast(dtypetorch.float16):lossmodel(batch[microbatch_id])else:loss.backward()# DP 的梯度同步FSDP 自动处理 AllReduceoptimizer.step()optimizer.zero_grad()性能分析LLaMA-13B on 64× Ascend 910 NPUhidden5120, layers40, seq2048, bs64 | 策略 | NPU 数 | 每卡显存 | 吞吐 (tokens/s/NPU) | 通信占比 | |------|--------|---------|--------------------|---------| | DP64 | 64 | OOM | — | — | | TP8, DP8 | 64 | 18.4 GB | 5,120 | 18% | | PP8, DP8 | 64 | 31.2 GB | 3,840 | 5% | | TP4, PP2, DP8 | 64 | 22.1 GB | 6,850 | 12% | | TP4, DP16 | 64 | OOM | — | — | 最优TP4, PP2, DP8 吞吐 6,850 tokens/s/NPU × 64 NPU 438,400 tokens/s 单卡显存22.1 GB在 32GB HBM 内安全 通信占比12%TP 的激活同步 9% DP 的梯度同步 3%为什么 TP4,PP2,DP8 最优TP8 通信太多18%PP8 bubble 太大吞吐低。DP64 OOM 卡放不下。TP/PP/DP 三元组在通信TP 同步激活、显存PP 分片模型、bubblePP 空闲三个维度中找平衡。踩坑一TP PP 的激活通信重复TP 的 AllReduce 调用在 PP 的每一段都会触发——PP rank 0 和 PP rank 1 各自做自己的 AllReduce。64 张 NPU 下TP4 的 AllReduce 组每组只有 4 张卡。但如果 TP 组的边界和 PP 段的边界不一致——就会出现跨段的通信浪费。# ❌ TP 组跨越 PP 段# PP0: Node0[Card0,1,2,3] → 算 Layer 0-15# PP1: Node1[Card0,1,2,3] → 算 Layer 16-31# TP 组 [Node0.Card0, Node0.Card1, Node0.Card2, Node0.Card3]# → TP AllReduce 在 PP0 完成后做PP1 也在自己的 TP 组做# → 重合PP0 的 TP AllReduce 不是 PP1 需要的# ✅ TP 组和 PP 段在同一节点内# PP0: Node0[Card0,1,2,3] → 算 Layer 0-15 → TP AllReduce (本地 NVSwitch)# PP1: Node0[Card4,5,6,7] → 算 Layer 16-31 → TP AllReduce (本地 NVSwitch)# TP 组分别独立不需要跨节点发送踩坑二梯度累积跨越 PP 段的边界PP 模式下loss.backward()只在最后一个 PP rank 上调用只有它持有实际的 loss。但 DP 的梯度同步需要所有参数都在——包括前面的 PP 段。# ❌ DP 的 AllReduce 跨越 PP 段# PP0 有 Layer 0-15 的梯度# PP1 有 Layer 16-31 的梯度# DP AllReduce 组包含 PP0 和 PP1 → 梯度不匹配有些参数缺梯度# ✅ FSDP 内部维护了跨 PP 段的梯度通信# torchtitan-npu 的 FSDPWrapper 处理这个——# 它只在持有完整层的 PP rank 上做梯度 AllReduce# PP0 的 Layer 0-15 梯度 → DP AllReduce 在 PP0 内部# PP1 的 Layer 16-31 梯度 → DP AllReduce 在 PP1 内部踩坑三PP 的微批次大小和工作集不匹配1F1B 调度中每次前向后马上反向的微批次需要存储前向的中间激活。如果 micro-batch8 但 tp_size4→每卡的中间激活 8× (hidden/tp) × seq_len × 每层中间结果。# ❌ micro-batch16 → 保存 16 份中间激活 → OOM# 每 micro-batch 的中间激活# Attention 的 QKV: 3 × [B/tp, S, d_head] × 2 bytes# FFN 的中间: [B/tp, S, ffn_dim] × 2 bytes# total per micro-batch ≈ 500MB (for 13B model, tp4)# 16 micro-batches → 8GB → 显存碎片化 → OOM# ✅ micro-batch4 → 保存 4 份 → 2GB → 安全# bubble 变大但显存安全torchtitan-npu 的 3D 并行不是三个独立策略的和——它们的通信和显存是耦合的。TP 的 AllReduce 在节点内 NVSwitch800GB/s 快但占带宽、DP 的 AllReduce 在节点间 RoCE200GB/s 慢但不频繁、PP 的 bubble 随 PP 大小线性增长PP2 时 12.5% 可接受。最优配置DP 跨节点 × TP 跨卡≤4× PP 跨层≤2——64 张 NPU 下吞吐 438,400 tokens/s。三个关键TP 组不和 PP 段交叉减少跨段通信、FSDP 的梯度同步只在持有完整层的 PP rank 上做、micro-batch 数不能太大PP2 时 ≤ 8→显存安全且 bubble12.5%。## 混合精度与通信压缩的叠加3D 并行的 AMP混合精度通信加倍了问题FP16 下每次 TP 的 AllReduce 传输是 FP32 的一半但梯度需要在 FP32 汇总再做更新——DP 的梯度 AllReduce 还是 FP32。# torchtitan-npu/amp.pyclassMixedPrecision3DParallel:AMP 3D 并行的通信优化def__init__(self,tp_group,dp_group,pp_group):self.tp_grouptp_group self.dp_groupdp_group self.pp_grouppp_group self.grad_scalertorch.npu.amp.GradScaler()deftrain_step(self,model,batch,optimizer):# 前向FP16TP 激活 AllReduce FP16省一半带宽withtorch.npu.amp.autocast():lossmodel(batch)# 反向梯度在 FP16 下传播self.grad_scaler.scale(loss).backward()# 梯度同步关键优化TP 的激活梯度已在 FP16 内部同步# DP 的权重梯度需要 Unscale → FP32 → AllReduce → Scale backself.grad_scaler.unscale_(optimizer)# FP32 梯度 AllReduceDPforparaminmodel.parameters():ifparam.gradisnotNone:# 只在 DP 组内做 AllReduce不在 TP 组重复做dist.all_reduce(param.grad,groupself.dp_group)param.grad/self.dp_group.size()self.grad_scaler.step(optimizer)self.grad_scaler.update()AMP 3D 通信的总结通信类型 数据量 精度 频率 TP 激活 AllReduce O(B×S×hidden/tp) FP16 每层 × 2 DP 梯度 AllReduce O(model_params/dp) FP32 每个 step × 1 PP 激活 P2P O(B×S×hidden) FP16 每 micro-batch × 1 总通信量 TP: 2 × 32 layers × 2 × B×S×D/tp × 2 bytes (FP16) DP: model_params/dp × 4 bytes (FP32) PP: num_microbatches × B×S×D × 2 bytes (FP16)Gradient Checkpoint 与 PP 的互动Gradient Checkpoint梯度检查点省显存但增加计算——和 PP 结合时checkpoint 的边界必须对齐 PP 的段边界。# ❌ Checkpoint 边界跨越 PP 段# PP0: Layer 0-15 → checkpoint(0-15)# PP1: Layer 16-31 → checkpoint(16-31)# 反向时 PP0 和 PP1 各自重算 checkpoint 段——# 但 PP1 需要 PP0 提供的激活→ PP1 不会重算 PP0 的部分# → 必须等 PP0 重算完再传激活给 PP1 → 延迟翻倍# ✅ PP 段内部做 checkpoint不跨越forpp_rankinrange(pp_size):first_layerpp_rank*layers_per_pp last_layer(pp_rank1)*layers_per_pp-1foriinrange(first_layer,last_layer,checkpoint_every):# 每 checkpoint_every 层做一个 checkpointtorch.utils.checkpoint.checkpoint(lambdax:layers[i:icheckpoint_every](x),intermediate_activation)HBM 快照与故障恢复3D 并行中 64 张 NPU 同时训练的故障概率每张 NPU 的 MTBF平均故障间隔≈ 5000 小时 → 64 张 NPU 的联合 MTBF ≈ 78 小时。训练 LLaMA-13B 需要 ≈ 200 小时 → 期望故障 ≈ 2.6 次。HBM 快照机制周期性地保存所有 NPU 的 HBM 状态——包括 TP shard、PP 段、DP 副本。# torchtitan-npu/snapshot.pydefsave_hbm_snapshot(model,optimizer,step):保存所有 NPU 的 HBM 状态到持久化存储snapshot{step:step,model_shards:{},# 每个 TP shard 的完整状态optimizer_state:{},# AdamW 的 momentum variancerng_state:{},# 随机数状态NPU 的 Philox 状态pp_microbatch_state:{}# PP 中间微批次状态}# 每张卡独立保存自己的状态并行 IOfortp_rankinrange(tp_size):forpp_rankinrange(pp_size):keyftp{tp_rank}_pp{pp_rank}snapshot[model_shards][key]model.state_dict(tp_rank,pp_rank)snapshot[optimizer_state][key]optimizer.state_dict()snapshot[rng_state][key]torch.npu.get_rng_state()# 写入文件异步不阻塞训练torch.save(snapshot,fsnapshot_step_{step}.pt)defrestore_hbm_snapshot(path):从快照恢复训练snapshottorch.load(path)# 恢复模型每个 TP/PP shard 恢复到对应 NPUfortp_rankinrange(tp_size):forpp_rankinrange(pp_size):keyftp{tp_rank}_pp{pp_rank}model.load_state_dict(snapshot[model_shards][key],tp_rank,pp_rank)optimizer.load_state_dict(snapshot[optimizer_state][key])torch.npu.set_rng_state(snapshot[rng_state][key])returnsnapshot[step]快照频率每 1000 步一次 → 200 小时训练 × 1000 步/hour 200 个快照 × 64 cards × 22GB 280TB 存储。实际实现用增量快照只存变化的部分降到 ≈ 5TB。调试3D 并行的常见异常定位defdebug_3d_parallel(model,dp_group,tp_group,pp_group):3D 并行的自检工具rankdist.get_rank()# 检查 1TP 的 AllReduce 是否正确iftp_group:xtorch.ones(1,devicenpu)*rank dist.all_reduce(x,grouptp_group)expectedsum(range(tp_group.size()))assertabs(x.item()-expected)1e-5,fTP all_reduce failed:{x}# 检查 2DP 的参数是否同步ifdp_group:forname,paraminmodel.named_parameters():gathered[torch.zeros_like(param)for_inrange(dp_group.size())]dist.all_gather(gathered,param,groupdp_group)# 所有 DP rank 的参数必须相同foriinrange(1,dp_group.size()):asserttorch.allclose(gathered[0],gathered[i]),\fDP param mismatch:{name}# 检查 3PP 的激活传递是否正确ifpp_groupandrank%pp_group.size()!pp_group.size()-1:# 不是最后一个 PP rank → 应该发送激活给下一段xtorch.randn(1,4096,devicenpu)dist.send(x,dstrank1,grouppp_group)# 接收方验证由下一个 rank 做torchtitan-npu 的 3D 并行是一个通信-显存-bubble 三角平衡问题。DP 跨节点RoCE 200GB/s 慢但每个 step 只同步一次梯度、TP 跨节点内卡NVSwitch 800GB/s 快但每层都要同步、PP 跨层最小通信但引入 bubble。LLaMA-13B 的最优配置 TP4, PP2, DP8 在 64 张 NPU 上实现 438,400 tokens/s单卡内存 22.1GB通信占比 12%。关键约束TP 组不和 PP 段交叉避免跨段通信浪费、PP 的 micro-batch 数量不能超过显存承载的中间激活上限、梯度检查点边界对齐 PP 段边界、AMP 下 TP 的 FP16 AllReduce 省一半带宽而 DP 的梯度同步保持 FP32 精度。
