总体思路

  1. 分析内存申请情况。
    从模型结构维度，以下四个阶段都会申请设备侧内存：
    • 模型初始化。模型权重等长生命周期内存，保存到训练结束。
    • 前向计算。参与反向计算的激活值等中长生命周期内存，保存到同step的反向结束。
    • 反向计算。首次反向计算产生的模型梯度等长生命周期内存，保存到训练结束。
    • 优化器梯度更新。优化器状态等长生命周期内存，保存到训练结束。
    • 模型初始化、前/反向计算和优化器梯度更新可能存在格式转换/精度转换/非连续转连续等，会申请临时内存。
  2. 分析内存使用情况。
    从内存使用维度，以下几个方面都会影响内存使用峰值：
    • 内存池缓存算法。不同的缓存算法，碎片控制水平存在差异。
    • 内存池个数。空闲内存不能跨内存池复用，内存池个数越多，空闲内存利用率就越低。
    • 算子实现。不同的实现算法/策略（例：小块数据多次计算vs大块数据几次计算），内存开销也不同。
    • PyTorch API/OP适配。不同的适配方式（例：大kernel vs 小算子拼接），内存开销也不同。
    • 非连续转连续。某些算子不支持非连续输入，需要插入格式转换算子，产生额外临时内存开销。
    • 私有格式。某些场景，私有格式比ND格式多占用内存，存在私有格式流转的网络内存峰值会高于ND格式整网流转。
  3. 采集内存数据。

    GPU和NPU各需要采两份内存profiling：

    1. 训练开始到首个step结束。（训练开始是指训练进程拉起后，不是首个step的开始，包含模型权重初始化等长生命周期内存申请阶段。）
    2. 采集一个稳定训练后的step。（需要包含优化器梯度更新阶段，注意在fp16或混合精度训练场景下，可能存在某次迭代梯度溢出导致跳过梯度更新的情况，影响内存行为分析。）
  4. 通过tensorboard查看内存使用情况，内存分析可参见《PyTorch 网络模型迁移和训练》中“性能调优 > Profiling数据采集与分析 > （推荐）Ascend PyTorch Porfiler数据采集与分析”章节中“性能分析（Memory View）”部分。显示效果如图1所示。
    图1 界面显示图
    • 绿色折线表示内存池向驱动申请的内存总量。
    • 蓝色折线表示内存池当前被使用的内存总量。

    左键选中一个区域，可以放大显示。右键返回全图模式。下方的表格列出了申请内存的算子，可以按照大小、申请时间、或释放时间来排序。可以通过表格左上方的搜索栏搜索相关算子，也可以通过表格右上方的滑轮拖动控制表格显示所选区域申请特定大小内存的算子。

  5. 判断内存占用问题基础步骤。
    1. 使用驱动命令获取内存总的使用情况。GPU使用nvidia-smi，NPU使用npu-smi info。
    2. 如果NPU和GPU总的内存使用相同，可以认为NPU上不存在由于软硬件限制导致的额外内存占用。
    3. 如果NPU和GPU总的内存使用不同，需要采集双方的内存profiling，方法见前面的描述。
    4. 采集到内存profiling后，使用tensorboard查看GPU内存池和NPU内存池向驱动申请的内存总量。
    5. 如果GPU内存池和NPU内存池向驱动申请的内存总量不同，查找导致内存不一致的地方。依次检查训练初始化、前向、反向以及优化器的内存使用情况。
    6. 如果GPU内存池和NPU内存池向驱动申请的内存总量相同，NPU还要检查CANN的内存申请情况，Tensorboard单独有一页显示CANN内存使用。
父主题： 内存分析与优化