PYTORCH_NPU_ALLOC_CONF

功能描述

PyTorch训练或在线推理场景，可通过此环境变量控制缓存分配器行为。

可选参数：

• max_split_size_mb:<value>，内存块允许切分上限。

  大于设定值的内存块在使用过程中不会进行切分，这有助于减少内存碎片。此选项主要应用于当模型由于oom(内存不足)而中断并显示存在大量非活动的拆分内存块场景的优化。<value>默认值为无限大，单位MB，即所有大小的内存块都可以拆分，最小设置值大于20MB。

• garbage_collection_threshold:<value>，垃圾回收阈值。

  主动回收未使用的NPU内存块，不会触发同步机制。在设置value阈值（例如0.8）后，如果NPU内存容量使用超过阈值（即分配给NPU应用程序的总内存的80%），缓存分配器将开始回收NPU内存块，优先释放最先申请和长时间未复用的内存块，避免释放积极复用的内存块。其中<value>取值范围为(0.0,1.0)，默认值为0.0。

• expandable_segments：使能内存池扩展段功能。

  实验特性，默认为False。如果设置为True，此设置将指示缓存分配器创建特定的内存块分配，这些内存块后续可以扩展，以便能更好地处理内存使用中频繁变更使用内存大小的情况。

参数配置格式PYTORCH_NPU _ALLOC_CONF=<option1>:<value1>,<option2>:<value2>。

配置示例

示例一：

export PYTORCH_NPU_ALLOC_CONF=max_split_size_mb:32,garbage_collection_threshold:0.6

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

使用约束

expandable_segments特性在驱动不支持此功能模块下无法使用。

支持的型号

Atlas 训练系列产品

Atlas A2 训练系列产品