调整内存参数

基本原理

详情参考PyTorch官方文档

使用场景

内存使用到极限,有oom异常退出。

操作步骤

  1. 设置内存因子,限制进程申请的内存上限,取值0~1。例:设置0.95,可使用的内存上限为60G * 0.95 = 57G。内存申请超过上限触发内存释放。设置PyTorch申请的内存上限,可以避免内存使用极限场景下,PyTorch耗尽device内存,其他组件申请内存失败导致的进程异常。

    torch_npu.npu.set_per_process_memory_fraction(0.95)

  2. 垃圾回收阈值,默认不开启,与内存因子配合使用,取值0~1,为内存上限的百分比。当内存申请到gc阈值,则触发内存池空闲内存块回收。建议由大到小配置调试。

    export PYTORCH_NPU_ALLOC_CONF="garbage_collection_threshold:0.95"

  3. 内存块允许切分上限,单位MB。例:设置50,大于等于50M的内存块不允许切分使用。设置该参数可避免大内存块被切分导致较多的内存碎片影响内存复用。调试时可以先采集内存profiling,按照算子内存申请降序排列,参数值由大到小尝试。

    export PYTORCH_NPU_ALLOC_CONF="max_split_size_mb:50"

  4. 使能内存池扩展段功能,由PyTorch自己管理虚拟地址与物理地址映射,降低内存碎片。对于动态shape场景,shape随step增加而增大,从而导致内存块不能复用内存碎片上升,对该场景有较好优化。

    export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True"

    若同时设置多个参数可以通过逗号分隔。

    export PYTORCH_NPU_ALLOC_CONF="garbage_collection_threshold:0.95,max_split_size_mb:50,expandable_segments:True"