大模型内存占用调优

模型运行时如果host内存不断增长,则可能发生了内存泄漏。发生内存泄漏的位置可能是数据集读取部分和网络部分,可以分别跑数据集和网络来定位内存泄漏的位置。

内存分析工具:memory_profiler

memory_profiler是python的第三方库,其功能是逐行代码分析内存占用,可以通过装饰器或命令行的方式使用。具体使用方式可以参

host内存溢出的场景

报错信息如下:

  • LLVM ERROR: out of memory
  • Memory consumption is more than …, which may cause oom error.
  • dataset/util/arena.cc … cudaHostAlloc failed, ret[2], out of memory

出现类似的错误信息,则说明host内存溢出,需要优化内存占用。

调优技巧

可通过以下checklist进行数据处理的内存占用调优:

数据集操作 调整方式
预取 set_prefetch_size 配置预取
数据加载类(**Dataset)、batch、map操作的工作线程数 num_parallel_workers 减少工作线程数,但会影响数据读取性能
map操作 减少/合并map操作
shuffle操作 减少buffer_size大小
batch操作 减少batch_size大小