【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsname: model-train-oom-analysis description: 用于诊断 PyTorch on NPU 大模型训练中的 NPU OOMOut of Memory问题。当用户报告训练因内存不足崩溃、出现 OOM 相关错误OutOfMemoryError / NPU out of memory / workspace allocator / HCCL memory、或需要进行训练内存优化时触发本技能。按照日志分类 → 静态估算 → snapshot 深度分析 → 优化建议的流程定位和解决问题。model-train-oom-analysis 技能用于 NPU 大模型训练 OOM 问题的系统化诊断与解决。适用范围与约定本技能的方法论适用于 torchtitan-npu 等 PyTorch on NPU 训练框架。文中出现的torchtitan_npu/...路径与配置项均为torchtitan-npu 框架的具体示例。torchtitan-npu 采用Python config registry配置是torchtitan_npu/models/model/config_registry.py中返回TrainerConfig(...)的函数各项为TrainingConfig/ProfilingConfig/ActivationCheckpointConfig/OptimizerConfig/ParallelismConfig等 dataclass 字段用--config 注册名选择并支持--section.fieldvalue形式的 CLI override不再使用 TOML 配置文件。使用其他框架时请将上述路径、配置项与启动方式映射到对应框架的训练入口、配置项与并行接口。内存估算公式、snapshot 分析维度、优化决策矩阵与框架无关可直接复用。参考资料分析脚本scripts/analyze_snapshot.py内存估算模板references/memory-estimation-templates.mdMosaic 工具指南references/mosaic-analysis-guide.mdPyTorch Mosaic 教程https://docs.pytorch.org/tutorials/beginner/mosaic_memory_profiling_tutorial.html脚本路径均相对本技能目录。若以源码方式调用请补全到model/model-train-oom-analysis/scripts/analyze_snapshot.py。适用场景训练过程中出现 OOM 崩溃OutOfMemoryError、NPU out of memory等。需要评估当前模型/配置的内存需求是否合理。训练前需要预估内存占用选择合适的并行策略和 batch_size。不适用场景精度问题loss 偏离、NaN/Inf→ 使用model-train-accuracy-debug技能。纯性能瓶颈训练慢但不 OOM。非 NPU 设备的内存问题。所需输入OOM 错误的完整训练日志必需。训练配置torchtitan-npu 为 config registry 注册名经--config传入其他框架为对应配置定义必需。NPU 设备显存规格如 Ascend 910B: 64GB。工作流Step 0日志分析 — 确定 OOM 类型与时机必做这是所有后续步骤的前提必须首先完成。分析训练日志中的 OOM 报错信息# 搜索 OOM 相关关键字 rg -n -i out.of.memory|OOM|OutOfMemory|NPU.*memory|allocator|HCCL.*memory train_log # 确认 OOM 发生在第几步 rg -n -i step|iteration train_log | tail -5根据报错关键字判定 OOM 来源根据发生时机判定是配置问题还是内存泄漏OOM 来源特征关键字后续步骤NPUWorkspaceAllocatorworkspace allocator、workspace memory→ Step 1 → Step 2HCCL 通信HCCL、hccl、communication→ Step 1 → Step 2NPUCachingAllocatorCachingAllocator、torch.OutOfMemoryError、NPU out of memory→ Step 1 → Step 3无法判断不含上述明确关键字→ Step 1 → Step 3[!IMPORTANT] OOM 类型判定直接决定后续方向。Workspace/HCCL OOM 需调水线Step 2CachingAllocator OOM 需先估算配置是否超限Step 1再决定是否需要 snapshot 深度分析Step 3。[!WARNING]渐进型 OOM 识别如果 OOM 发生在训练后期如 step 10而非第 1-2 步很可能是内存泄漏而非配置不足。OOM 在 step 0-2 → 配置问题 → Step 1 → Step 4OOM 在 step NN 较大→ 疑似泄漏 → Step 1 → Step 3含 3.5 泄漏检测 3.6 代码审查Step 1静态内存估算 — 判断配置是否可行必做在采集 snapshot 之前先做静态估算快速判断当前配置下内存是否本身超限。1.1 收集参数模型参数dim、n_layers、n_heads、vocab_size、inter_dim、MoE 参数专家数、共享专家、激活专家数等从模型配置/定义获取。torchtitan-npu 示例从torchtitan_npu/models/model/__init__.pyflavors 函数如_make_dsv3_model_config和torchtitan_npu/models/model/model.pyConfig继承上游 torchtitan 的 ModelArgs获取。训练参数从配置文件获取local_batch_size、seq_len、TP/PP/DP/EP/CP 各并行度、activation_checkpoint策略、优化器类型。1.2 估算与判断详细公式参见 references/memory-estimation-templates.md核心训练峰值内存 ≈ 参数内存 梯度内存 优化器状态 激活内存 临时缓冲 估算总内存 ≤ NPU 总显存 × 0.95 (预留 5% 安全余量)1.3 决策路由估算结果判断后续步骤峰值估算 NPU 显存 × 0.95配置本身超限无需 snapshot→ 直接 Step 4 调整配置Workspace/HCCL OOM 且估算合理临时内存不足→ Step 2 调水线CachingAllocator OOM 且估算接近上限需确认实际分布→ Step 3 snapshot 分析估算远低于上限但仍 OOM可能碎片化或泄漏→ Step 3 snapshot 分析[!TIP]配置超限是最常见的 OOM 根因——参数量过大、batch_size/seq_len 过大、未启用激活重计算、优化器状态未分片/卸载等都会导致显存需求超出单卡上限。此时无需采集 snapshot直接根据 Step 4 优化矩阵调整配置即可。Step 2临时内存 OOM — 配置 PTA 内存水线当 OOM 来自NPUWorkspaceAllocator或 HCCL 通信且 Step 1 估算表明模型本身可以放下时问题根因是 PTA 缓存分配器占用了过多显存导致临时工作空间或通信缓冲区无法分配。通过限制 PTA 显存占用上限为临时内存留出空间。其底层基于torch.npu.set_per_process_memory_fraction()各框架暴露的配置项不同。torchtitan-npu 示例在配置函数的trainingTrainingConfig(...)中设置或启动时 CLI overridetrainingTrainingConfig(torch_npu_memory_ratio0.9), # 限制 PTA 最多使用 90% 显存留 10% 给 workspace/HCCL--config 注册名 --training.torch_npu_memory_ratio0.9该功能在torchtitan_npu/train.py的_patch_for_train_npu_memory()中实现。其他框架可直接在训练入口调用torch.npu.set_per_process_memory_fraction(ratio)。首次尝试可设0.9若仍 OOM 降低到0.85。若过低导致 CachingAllocator 本身 OOM说明模型内存需求过高需回到 Step 1 重新估算并进入 Step 4 优化配置。Step 3CachingAllocator OOM — 采集并分析 Memory Snapshot当 Step 1 估算无法确认问题根因估算接近上限、或估算合理但仍 OOM时采集 memory snapshot 做详细分析。3.1 启用 Memory Snapshot 采集通过torch.npu.memory._record_memory_history()_dump_snapshot()采集多数训练框架已封装该能力。torchtitan-npu 示例在torchtitan_npu/models/model/config_registry.py对应配置函数的profilingProfilingConfig(...)中开启profilingProfilingConfig( enable_memory_snapshotTrue, save_memory_snapshot_foldermemory_snapshot, ),或启动时用 CLI override无需改注册文件--config 注册名 --profiling.enable_memory_snapshot --profiling.save_memory_snapshot_foldermemory_snapshottorchtitan 在 OOM 发生时也会自动转储快照。3.2 脚本化分析使用 scripts/analyze_snapshot.py 对.pickle文件进行多维度分析[!NOTE] snapshot 基于 pickle 反序列化加载请仅分析可信来源自己训练产出的.pickle文件不要加载来历不明的快照。# 单快照分析 python scripts/analyze_snapshot.py snapshot.pickle # JSON 输出便于程序化处理 python scripts/analyze_snapshot.py snapshot.pickle --json # 自定义 Top-N python scripts/analyze_snapshot.py snapshot.pickle --top-n 20脚本输出包含6 个维度的分析① 内存分配概览Reserved / Allocated / Active / 碎片化率 / 设备利用率② 语义类别分类按 Mosaic 分类体系将内存细分为激活、梯度、优化器状态、参数、通信缓冲、注意力层、MoE 层等类别每项给出大小和占比③ 峰值内存分析定位峰值时刻展示峰值各类别占比和 Top 调用栈④ 时间线分析追踪 alloc/free 事件序列识别训练阶段增长/释放/稳定检测潜在内存泄漏⑤ 碎片化深度分析空闲块分布、最大空闲块、小块占比并给出碎片化等级评估⑥ Top-N 大块分配最大内存块的大小、类别和完整调用栈分析判断矩阵指标阈值对应优化措施激活内存占比 40%启用 activation checkpoint (--activation_checkpoint.modefull或selective)优化器状态占比 40%启用优化器状态卸载/分片如 swap optimizer、FSDP 分片梯度内存占比 30%增大 TP/PP 并行度通信缓冲占比 20%调整并行策略、检查 HCCL 配置碎片化率 20%调用torch.npu.empty_cache()、调整分配策略其他/未知占比 20%使用 Mosaic 深度分析定位来源潜在泄漏标记触发对比多 step snapshot 确认3.3 深度分析 — Mosaic 工具可选当基础分析无法充分定位问题时尤其是 其他/未知 类别占比高使用 Mosaic 做更精确的分析。详见 references/mosaic-analysis-guide.md。# 峰值栈追踪 mosaic_get_memory_usage_peak --snapshot snapshot.pickle # 分类可视化 Profile mosaic_get_memory_profile --snapshot snapshot.pickle \ --out-path profile.html \ --profile categories \ --preserve-allocation-order # 自定义模式匹配追踪特定操作的内存占用 mosaic_get_memory_profile --snapshot snapshot.pickle \ --out-path custom_profile.html \ --profile custom \ --custom-profile {hccl: hccl, expert: expert, moe: moe}3.4 对比分析 — 双快照差异对比对比两个场景如开启 AC 前后、不同 batch_size、不同并行策略的内存差异python scripts/analyze_snapshot.py \ baseline.pickle modified.pickle \ --label-a 无AC --label-b 有AC输出包含各指标的并排对比表格、各类别的内存变化增/减量和百分比、关键洞察。3.5 泄漏检测 — 多步骤 Snapshot 对比当 Step 0 判断为渐进型 OOMOOM 发生在训练后期时执行此步。# 传入 3 个不同 step 的 snapshot自动分析增长趋势 python scripts/analyze_snapshot.py \ step5.pickle step10.pickle step20.pickle现象结论下一步内存单调增长某类别持续增长确认泄漏定位到具体类别→ Step 3.6 代码审查内存增长但后期稳定warmup 阶段正常增长→ 排除泄漏回到 Step 4内存未增长但仍 OOM非泄漏内存本身不足→ Step 1 估算 Step 4 优化3.6 代码审查 — 泄漏根因定位当 Step 3.5 确认存在泄漏时执行此步。审查 Checklist异步通信 重计算交互当activation_checkpoint TP 同时启用时redistribute(async_opTrue)产生的异步 tensor 在重计算过程中可能未被正确wait导致中间张量无法释放。检查RowwiseParallel/ColwiseParallel的输出处理函数是否显式调用了wait_tensor()Tensor 引用持有检查是否有调试代码、全局变量或日志记录意外持有了 tensor 引用自定义 autograd Function检查save_for_backward保存的 tensor 是否在backward中被正确释放Hook 注册检查是否有未清理的 forward/backward hook 持有 tensor# 检查并行化代码中的异步通信路径以实际框架的并行化模块为准 rg -n async_op|redistribute|_prepare_output_fn parallelize_module # 检查是否有 wait_tensor 配对 rg -n wait_tensor|async_opTrue source_root[!TIP] 通信类泄漏的典型修复方式是在redistribute后显式调用torch.ops._c10d_functional.wait_tensor()确保异步操作完成。Step 4优化建议矩阵根据 Step 0~3 的分析结果按照以下矩阵给出针对性优化建议配置项以 torchtitan-npu 为例其他框架对应调整配置项以 torchtitan-npu 为例给出 CLI override 写法等价于在 config registry 配置函数对应 dataclass 字段中修改。问题诊断优化措施配置修改torchtitan CLI override 示例预期效果临时内存不足Workspace/HCCL OOM配置内存水线--training.torch_npu_memory_ratio0.85~0.95为临时内存预留空间激活内存过大启用全量激活重计算--activation_checkpoint.modefull激活内存降至最低激活内存过大折中选择性激活重计算--activation_checkpoint.modeselective --activation_checkpoint.selective_ac_optionop激活内存适度降低激活内存过大减小 batch_size--training.local_batch_size减半激活内存近似线性降低激活内存过大减小 seq_len--training.seq_len减半激活内存近似线性降低优化器状态过大启用 swap optimizer--optimizer.swap_optimizer --optimizer.swap_optimizer_times16优化器状态卸载到 CPU参数梯度内存过大增大 TP 并行度--parallelism.tensor_parallel_degree↑参数/梯度按 TP 切分参数梯度内存过大增大 PP 并行度--parallelism.pipeline_parallel_degree↑参数按层切分到不同卡MoE Expert 内存过大增大 EP 并行度--parallelism.expert_parallel_degree↑Expert 在更多卡间分布内存碎片化严重清理缓存在代码中适时调用torch.npu.empty_cache()释放碎片化内存长序列内存过大启用 Context Parallel--parallelism.context_parallel_degree↑序列维度切分内存泄漏通信类异步通信加 wait修改并行化代码在redistribute后调用wait_tensor()消除泄漏内存泄漏其他类审查引用持有检查并修复意外的 tensor 引用持有消除泄漏决策优先级按推荐顺序首先确认 OOM 类型一次性 vs 渐进型若为渐进型先做泄漏检测若为泄漏通过多步骤 snapshot 定位泄漏类别审查对应代码若为临时内存 OOM先调水线然后启用 activation checkpoint对训练效果无影响仅牺牲少量性能接着启用 swap_optimizer减少 NPU 内存压力增加少量 CPU 通信开销再考虑调整并行策略可能需要更多卡或修改训练逻辑最后减小 batch_size / seq_len影响训练效率和收敛行为Step 5补充工具 — msmemscope可选当上述步骤无法充分定位问题时可使用 msmemscopeMindStudio MemScope进行更深层次的内存分析。项目地址https://gitcode.com/Ascend/msmemscopemsmemscope 提供以下高级分析能力内存泄漏检测识别训练过程中随 step 增长的异常内存分配内存拆解分析按组件细粒度分解显存使用比 memory snapshot 更精细内存对比监测对比不同训练阶段warmup vs 稳定训练的内存差异使用方式参考 msmemscope 官方文档中的 PyTorch 采集示例和分析指南。输出要求最终诊断报告必须包含OOM 类型判定来源及日志证据静态内存估算结果参数/梯度/优化器/激活各占多少memory snapshot 分析结果若已采集各类别占比、碎片化率、Top 内存分配具体优化建议含配置修改示例优化后的预期内存占用是否需要进一步分析的建议关键路径torchtitan-npu 示例以下路径为 torchtitan-npu 框架的定位锚点示例使用其他框架时映射到对应的训练入口、模型定义与并行模块。类别路径训练入口torchtitan_npu/entry.py训练 patch含内存水线torchtitan_npu/train.py自定义配置torchtitan_npu/config/configs.py模型定义torchtitan_npu/models/model/model.py模型参数 flavorstorchtitan_npu/models/model/__init__.py训练配置torchtitan_npu/models/model/config_registry.py--config传注册名非文件路径并行逻辑torchtitan_npu/models/model/parallelize.py【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考