MindStudio支持PyTorch Memory Snapshot,内存分析更丝滑

0 阅读5分钟

一、强化学习内存问题定位难题

在强化学习场景下存在多个不同的任务阶段,其内存变化比较大,可以使用 PyTorch Profiler 中的 profile_memory 数据或者 Snapshot 数据来分析内存使用,特别是Snapshot,其可以查看内存使用细节,在深入分析时非常有用。

PyTorch 官方社区提供了memory_viz用于可视化 Snapshot数据。然而使用起来有以下几个问题:

  1. memory_viz 仅能支持数十MB的 Snapshot文件分析。当采集时间较长,采集得到的 Snapshot 数据超过 100MB时,导入memory_viz 明显卡顿,甚至崩溃,难以完成分析任务。
  2. memory_viz支持展示的两个关键图像(活跃内存生命周期图和内存分配器历史状态图)分别在不同的页面中,当需要联合分析时,必须不断切换页面重新加载,降低了定位问题的效率。

二、MindStudio支持PyTorch Memory Snapshot

MindStudio 可视化工具MindStudio Insight支持图表联动,前端操作步骤少,而且对大文件查询展示等操作进行了优化,因此 操作用时大幅减小,相较当前社区能力,耗时缩短80%以上10秒内就能实现功能。具体优势如下:

  1. 支持导入和分析 Snapshot 数据,使用体验与社区基本一致。
  2. 支持的Snapshot数据上限从官方100MB提升至1G,保障大数据量场景(如:强化学习)内存问题分析定位任务顺利完成。
  3. 实现活跃内存生命周期图(Active Memory Timeline)与内存分配器历史状态图(Allocator State History)的双图联动可视化,方便对比分析。

活跃内存生命周期图(Active Memory Timeline):表示内存块的生命周期堆叠情况。核心数据源来自device_traces。通过事件序列解析出内存块的生命周期。

内存分配器历史状态图(Allocator State History):在Snapshot数据中仅有dump时刻的内存分配器的状态。而该页签依赖于内存回放机制展示任意事件发生后的内存分配器状态。

图1  双图联动演示.gif

图1 双图联动演示

三、实战案例

3.1 案例一:分析内存泄漏问题

3.1.1 问题描述

某模型 Demo 训练,长时间运行后出现 OOM(Out Of Memory)。

3.1.2 排查过程

第一步:使用 Snapshot 数据分析

采集 Snapshot 数据导入 MindStudio Insight。观察内存申请/释放折线图&内存块图(Memory Request/Release Line Graph & Memory Block
Graph)。发现图中每步训练都创建了一段明显的内存,但在之后一直保留没有释放。这是内存泄漏问题,很可能就是导致长时间运行后出现OOM(Out Of Memory)的原因。

图片2.png

图2 三步训练过程中内存块图存在明显未释放内存块

接下来查看这段内存是由哪段代码申请创建的,定位根因代码。

图片3.png

图3 看创建未释放内存块的代码调用栈

第二步:查看代码

打开代码文件memory_leaks_demo.py 并定位到第 12 行,发现变量leak_mem_list是全局变量,而且没有释放的代码,确认是这里导致的内存泄漏,只需将无用变量leak_mem_list 删除即可。

4.png

图4 在内存泄漏问题的代码片段

3.2 案例二:分析内存碎片问题

3.2.1 问题描述

强化学习训练多模态模型,但训练效率很慢。

3.2.2 排查过程

第一步:使用 Profiling 数据初步分析

采集 Profiling 数据导入 MindStudio
Insight。发现无明显快慢卡现象,与卡间通信无关。查看内存折线图,观察发现模型训练过程中,进程/算子保留(APP/Operator Reserved)远大于算子分配(Allocated),且稳定后并未发现明显内存泄露,实际使用(Allocated)峰值远低于保留(APP/Operator Reserved)值,首先考虑是否为存在较多/较大内存碎片。

5.png

图5 训练过程中内存折线图进程/算子保留值远大于算子分配值

通过修改代码开启虚拟内存,手动触发内存重整,让程序回收内存碎片,再观察内存使用情况。发现每个Step仍在生命周期中被顶到与之前接近的峰值。这说明可能代码运行过程中,存在导致单步内短效内存占用较大的场景,如变长序列场景中输入较长序列。

6.png

图6 原始三步训练过程中的内存折线图

7.png

图7 开启虚拟内存后三步训练过程中的内存折线图

第二步:使用 Snapshot 数据深入分析

采集 Snapshot 数据导入 MindStudio Insight 查看,发现在一个 Step 的最后一个 segment alloc 时,已出现了约 10.6G的差距。右边图中空白块表示未使用内存,彩色块表示已使用内存。

8.png

图8内存池状态图中存在大块空白内存块

发现存在较大量的碎片和大块的未被复用的内存,确定是内存碎片问题。结合训练代码/逻辑进一步分析,得到最终结论:这是输入序列长度变化太大导致的内存碎片问题。训练数据中个别极端长度的数据会把显存顶上去,之后的代码中没有做释放操作,因此碎片不断累积,拖慢训练效率。

修改代码文件中处理变长序列尾部拼接功能的逻辑,并将 batch_size 从 3提升到 4,再次运行效率提升显著,未出现明显内存碎片。

四、总结

内存问题一直是训练和推理中常见的问题,使用 PyTorch 社区提供的memory_viz工具在小数据场景下可以分析定位内存问题。但是随着模型参数规模的不断增加,分析内存问题需要采集的数据量也在飞速增长。MindStudio Insight工具,在 2026 年 3 月上线的版本开始支持 Snapshot格式数据,并允许导入 1G大小的数据,联动功能也更方便定位内存问题,欢迎下载使用。

MindStudio Insight代码仓gitcode.com/Ascend/msin…

公众号二维码.jpg

扫描二维码,关注微信公众号,学习昇腾前沿科技知识