LLaMA-13B 在 8 张 NPU 上训练batch_size128、seq_len4096 → 单卡需要 28GB 显存放激活activation但 Ascend 910 只有 56GB HBM——28GB 激活 13GB 参数 8GB 优化器状态 49GB只剩 7GB 余量。如果 seq_len 涨到 8192激活翻倍到 56GB → OOM。两条路径梯度累积减少每步的激活量多步后才同步和 Gradient Checkpoint不存所有激活反向时重算。单独用任一都不够——梯度累积解决不了单层激活 剩余 HBM 的问题Checkpoint 的重算开销降低吞吐 30%。联合用才最优梯度累积把 batch 切成 micro-batch每个 micro-batch 只存 checkpoint 标记层的激活中间层反向时重算。梯度累积用时间换显存标准训练一个 step forward(128 samples) → backward → optimizer.step()。128 samples 的激活全存 HBM。梯度累积128 8 × 16 → forward(16 samples) → backward(16 samples) [不 step] → forward(16) → backward(16) [不 step] → …(重复 6 次)… → optimizer.step() [等效 batch128]。# cann-recipes-train/examples/gradient_accumulation.pyimporttorchimporttorch_npufromtorch.distributed.fsdpimportFullyShardedDataParallelasFSDPdeftrain_with_grad_accum(model,dataloader,optimizer,config): config.accum_steps 8 # 累积 8 步 config.micro_batch_size 16 # 每步 16 samples config.effective_batch 128 # 等效 batch model.train()optimizer.zero_grad()forstep,batchinenumerate(dataloader):# 把 batch 切成 micro-batchesmicro_batchesbatch.chunk(config.accum_steps)fori,micro_batchinenumerate(micro_batches):# 前向只存当前 micro-batch 的激活withtorch.autocast(npu,dtypetorch.float16):lossmodel(micro_batch)/config.accum_steps# ↑ 除以 accum_steps让 loss 反向时的梯度是平均梯度# 等价于所有 micro-batch 的 loss 求和后反向# ∇(loss/8) ∇loss/8 → 8 步累积后 ∇loss# 反向梯度留在 param.grad 里不清零loss.backward()# 释放 micro-batch 的激活让 HBM 可以给下一步用delmicro_batch,loss torch.npu.empty_cache()# 8 个 micro-batch 的梯度都累积在 param.grad → 一次 step# 梯度同步FSDP AllReduce也只做一次省 7 次通信torch.nn.utils.clip_grad_norm_(model.parameters(),max_norm1.0)optimizer.step()optimizer.zero_grad()梯度累积的通信收益8 个 micro-batch 只做 1 次梯度 AllReducevs 8 次省了 87.5% 的通信量。但显存收益有限——micro-batch 的激活在内存中不在 HBM 中只有当前 forward 的激活在 HBM。Gradient Checkpoint用计算换显存标准训练forward 把每层的激活存下来backward 直接读。LLaMA-13B 的 40 个 Transformer 层每层存 0.7GB 激活batch16, seq4096→ 28GB。Checkpoint只存 N 个「检查点」层的激活其他层在 backward 时从检查点重算 forward。比如每 2 层存一个检查点 → 存 20 层14GB省 14GB。代价每层额外一次 forward → 额外 50% 计算量40 层 forward backward 80 层的工作量with checkpoint 变成 40 40 20 100 层 25% 额外计算。# cann-recipes-train/examples/gradient_checkpointing.pyimporttorch.utils.checkpointascheckpointclassTransformerWithCheckpoint(nn.Module):def__init__(self,num_layers,hidden_dim,checkpoint_every2):super().__init__()self.num_layersnum_layers self.checkpoint_everycheckpoint_every self.layersnn.ModuleList([TransformerLayer(hidden_dim)for_inrange(num_layers)])defforward(self,x):forlayer_idx,layerinenumerate(self.layers):iflayer_idx%self.checkpoint_every0:# 检查点层forward 正常执行激活存 HBMxlayer(x)else:# 非检查点层用 checkpoint不存激活重算xcheckpoint.checkpoint(layer,x,use_reentrantFalse# PyTorch 2.0 推荐)returnx# 配置modelTransformerWithCheckpoint(num_layers40,hidden_dim5120,checkpoint_every2# 每 2 层一个检查点 → 20 个检查点)# 显存占用# 无 checkpoint: 40 层 × 0.7GB 28GB# checkpoint_every2: 20 层 × 0.7GB 14GB省 14GB# checkpoint_every4: 10 层 × 0.7GB 7GB省 21GB但重算 75%联合策略梯度累积 Checkpoint 的叠加# cann-recipes-train/examples/combined_strategy.pydeftrain_combined(model,dataloader,optimizer,config):联合策略梯度累积 8 step checkpoint_every2accum_steps8micro_batch_size16# 每个 micro-batch 的激活20 层 checkpoint × 0.7GB 14GB# 加上参数 13GB 优化器 8GB 35GB → 安全56GB HBMoptimizer.zero_grad()forstep,batchinenumerate(dataloader):micro_batchesbatch.chunk(accum_steps)fori,mbinenumerate(micro_batches):# 最后一个 micro-batch 同步梯度AllReduce# 其他 micro-batch 只在本地累积withmodel.no_sync()ifiaccum_steps-1elsecontextlib.nullcontext():withtorch.autocast(npu,dtypetorch.float16):lossmodel(mb)/accum_steps loss.backward()delmb,loss torch.npu.empty_cache()# 8 步累积后 → 1 次 optimizer.step()clip_grad_norm_(model.parameters(),1.0)optimizer.step()optimizer.zero_grad()显存分析LLaMA-13B (40 layers, seq4096, hidden5120) on 8× Ascend 910 | 策略 | 激活显存 | 参数优化器 | HBM 总计 | 通信次数/step | 吞吐 | |------|---------|-----------|---------|------------|------| | 无优化 (bs128) | 28GB | 21GB | 49GB (OK) | 1 | 8,200 t/s | | 无优化 (bs256) | 56GB | 21GB | 77GB OOM!| — | — | | GradAcc8 (bs16 ×8) | 7GB | 21GB | 28GB (OK) | 1 | 7,800 t/s | | Checkpoint (every2) | 14GB | 21GB | 35GB (OK) | 1 | 6,400 t/s | | **Checkpoint GradAcc8** | 3.5GB | 21GB | 24.5GB (OK) | 1 | 6,100 t/s | | **Checkpoint GradAcc4** | 7GB | 21GB | 28GB (OK) | 1 | 6,800 t/s |最优Checkpoint (every2) GradAcc4。吞吐 6,800 tokens/s比无优化的 8,200 低 17%但 batch 从 128 扩到 512 的等效 batch收敛快 4×。踩坑一Checkpoint 放错位置——嵌入层和输出层embedding 层和 lm_headlogits 层不能放 checkpoint——这两层是参数的入口和出口forward 不产生大显存embedding 只需存 lookup 结果0.5GB但 backward 需要梯度。checkpoint 这里→反向时 embedding 被重算→梯度丢失。# ❌ embedding 放在 checkpoint → 梯度丢失classBadModel(nn.Module):defforward(self,input_ids):# embedding 不占大显存但放 checkpoint 梯度没了xcheckpoint.checkpoint(self.embedding,input_ids)# ← 错误forlayerinself.layers:xlayer(x)returnself.lm_head(x)# → loss 不降embedding 的梯度 0# ✅ 从第一个 Transformer 层开始 checkpointclassCorrectModel(nn.Module):defforward(self,input_ids):xself.embedding(input_ids)# ← 正常 forward不 checkpointfori,layerinenumerate(self.layers):ifi%20:xlayer(x)# 检查点层else:xcheckpoint.checkpoint(layer,x)# 重算层returnself.lm_head(x)# ← lm_head 也不 checkpoint踩坑二Gradient Checkpoint 和 FSDP 的 no_sync 冲突FSDP 的no_sync()跳过 AllReduce——但 checkpoint 的反向传播走两次 forward一次重算梯度的requires_grad状态在重算 forward 时不正确 → FSDP 的no_sync判断失败 → 重复通信。# ❌ no_sync checkpoint → 重复 AllReduce吞吐降 15%withmodel.no_sync():# 期望跳过 AllReducelossmodel(mb)loss.backward()# → checkpoint 内部重算时FSDP 误判状态 → 做了 AllReduce# → 8 个 micro-batch 做了 8 次 AllReduce期望 1 次# ✅ 检查点层不放 no_sync 内手动管理 FSDP 状态fori,mbinenumerate(micro_batches):lossmodel(mb)/accum_steps loss.backward()ifiaccum_steps-1:# 不等 AllReduce攒着forparaminmodel.parameters():ifparam.gradisnotNone:param.grad._accumulatingTrue# 标记「还在累积」else:# 最后一步做 AllReduceforparaminmodel.parameters():ifparam.gradisnotNone:param.grad._accumulatingFalse踩坑三HBM Empty Cache 过多反而降低吞吐每步torch.npu.empty_cache()释放碎片——但执行太频繁如每个 micro-batch 调一次→ NPU 驱动反复释放/分配 → 额外 2-3ms 开销 × 8 micro-batches 16-24ms per step。# ❌ 每个 micro-batch 都 empty_cache → 20ms 额外开销fori,mbinenumerate(micro_batches):lossmodel(mb)loss.backward()torch.npu.empty_cache()# ← 每步 2ms → 8 步 16ms# ✅ 只在显存压力大时最后一步前调一次fori,mbinenumerate(micro_batches):lossmodel(mb)loss.backward()# 只在 OOM 风险高时清理显存使用 90%ifiaccum_steps-2:# 倒数第 2 步mem_usedtorch.npu.memory_allocated()/torch.npu.max_memory_allocated()ifmem_used0.9:torch.npu.empty_cache()# 最后一步后必然清理# 为 optimizer step 腾出空间torch.npu.empty_cache()梯度累积 Gradient Checkpoint 联合用Checkpoint 把激活从 28GB 压到 14GB每 2 层一个检查点梯度累积把 batch 切成 4 个 micro-batch每个 3.5GB 激活→ HBM 总计 24.5GB。吞吐从 8,200 tokens/s 降到 6,80017%但 batch 从 128 扩到 512 等效收敛加速 4×。通信省了 75%4 micro-batch 只做 1 次 AllReduce。三个踩坑embedding/lm_head 不 checkpoint 否则梯度丢失、checkpoint 反向重算触发 FSDP 误判导致重复 AllReduce、empty_cache 太频繁反噬吞吐每步 2ms 开销。
