CANN/cannbot-skills NPU训练内存估算模板

NPU 训练静态内存估算模板【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills本文档提供训练内存估算的通用模板和常见模型的参考值用于在训练前评估内存需求或诊断 OOM 问题。内存组成总览训练峰值内存 ≈ ① 参数内存 ② 梯度内存 ③ 优化器状态 ④ 激活内存 ⑤ 临时缓冲区各组成部分估算方法① 参数内存Dense 模型: P × dtype_bytes / (TP × PP) MoE 模型: P_dense × dtype_bytes / (TP × PP) P_expert × dtype_bytes / (TP × EP × PP)P 参数总量如 8B 80亿dtype_bytes 数据类型字节数bf162, fp324, fp81② 梯度内存≈ 参数内存 (梯度与参数同 dtype、同切分方式)③ 优化器状态内存AdamW最常见 P / 并行度 × 4B × 2 (fp32 的一阶矩 m 和二阶矩 v)若启用 FSDP (data_parallel_shard_degree 1): 状态在 DP 维度上进一步分片若启用swap_optimizer true: 优化器状态卸载到 CPUNPU 上几乎为 0④ 激活内存激活内存与batch_size (B)、seq_len (S)、hidden_dim (H)、num_layers (L)和激活重计算策略强相关重计算策略[activation_checkpoint]配置激活内存估算无重计算mode none≈B × S × H × L × dtype_bytes × KK≈10~15全量重计算mode full≈B × S × H × dtype_bytes仅保存模型输入选择性重计算mode selective介于两者之间取决于selective_ac_option其中K是每个 Transformer 层保存的中间张量数量系数取决于模型架构。⑤ 临时缓冲区≈ 总内存 × 5%~15%包括HCCL 通信缓冲、NPUWorkspaceAllocator、梯度聚合桶、临时计算空间等。常见模型估算参考Llama3 8B (bf16, AdamW)组成部分计算方式单卡 (无并行)TP8参数内存8B × 2B16 GB2 GB梯度内存≈ 参数内存16 GB2 GB优化器状态 (mv)8B × 4B × 264 GB8 GB激活内存 (selective AC)取决于 B×S变动变动静态总计 (不含激活)96 GB12 GB配置示例llama3-8b参数示意实际配置见config_registry.pylocal_batch_size 2,seq_len 2048activation_checkpoint.mode selectiveswap_optimizer true→ 优化器内存卸载到 CPUDeepSeek V32 671B (bf16, AdamW, MoE)组成部分说明估算值Dense 参数~14B dense 参数28 GB / (TP × PP)Expert 参数~657B expert 参数 (256 experts)1314 GB / (TP × EP × PP)优化器状态取决于 FSDP 分片、swap_optimizer取决于具体并行策略激活内存强依赖 AC 策略、B、S需实际测量配置示例DeepSeek-V3.2 671B参数示意实际配置见config_registry.pylocal_batch_size 4,seq_len 2048activation_checkpoint.mode selective256 experts,expert_parallel_degree/expert_tensor_parallel_degree需合理设置[!NOTE] 以上为理论估算参考值。实际内存占用受内存碎片、框架开销、通信缓冲等因素影响建议结合 memory snapshot 实测验证。内存预算检查公式估算总内存 ≤ NPU 总显存 × 0.95 (预留 5% 安全余量)若估算结果超出预算按以下优先级调整参见 SKILL.md Step 4 优化建议矩阵启用 activation checkpoint启用 swap_optimizer增大并行度TP/PP/EP减小 batch_size / seq_len【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考