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。垃圾回收阈值需与内存因子配合使用，内存因子可参考《API 参考》的“torch_npu.npu.set_per_process_memory_fraction”章节。

• expandable_segments：使能内存池扩展段功能，即虚拟内存特性。

  默认为False。如果设置为True，此设置将指示缓存分配器创建特定的内存块分配，这些内存块后续可以扩展，以便能更好地处理内存使用中频繁变更使用内存大小的情况。如果设置为False，关闭内存池扩展段功能，使用原有的内存申请方式。

• base_addr_aligned_kb：取值范围为0~16，设置值需为整数，单位KB，默认值为16。仅在expandable_segments设置为true的时候生效。若此参数配置为16，在申请大块内存（大于等于2M）的时候，会尽量保持申请内存的基地址16KB对齐；若配置为0时，申请内存的基地址512B对齐。此外配置此参数会对内存占用量、性能造成波动。

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

max_split_size_mb和garbage_collection_threshold中任意一个为非默认值时，expandable_segments必须设置为False。

用户使用Ascend Extension for PyTorch 6.0.rc3版本配套的驱动（Ascend HDK 24.1.RC3），开启虚拟内存特性时，可以使用单进程多卡特性；用户使用Ascend Extension for PyTorch 6.0.rc3之前版本配套的驱动（Ascend HDK 24.1.RC3之前版本），开启虚拟内存特性时，不能使用单进程多卡特性。

配置示例

示例一：

export PYTORCH_NPU_ALLOC_CONF=max_split_size_mb:32,garbage_collection_threshold:0.6

使用约束

expandable_segments特性需在Ascend HDK 23.0.0及以上版本上使用。

支持的型号

Atlas 训练系列产品

Atlas A2 训练系列产品