PYTORCH_NPU_ALLOC_CONF

功能描述

通过此环境变量可控制缓存分配器行为。配置此环境变量会改变内存占用量，可能造成性能波动。

可选参数：

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

  大于设定值的内存块在使用过程中不会进行切分，这有助于减少内存碎片。此选项主要应用于当模型由于oom（Out of Memory，内存不足）而中断并显示存在大量非活动的切分内存块场景的优化。<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:<value>，使能内存池扩展段功能，即虚拟内存特性。

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

• base_addr_aligned_kb:<value>，内存基地址对齐粒度。

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

• page_size:<value>，设置申请大页内存的大小。

  取值仅支持1G，参考配置示例。内存申请粒度为1G，不足1G的向上对齐到1G，始终为1G的整数倍。

  未配置该选项时，内存申请粒度为2M，如需申请1G的大页内存，会占用1024/2=512个页表；设置该选项后，内存申请粒度为1G，1G大页内存只占用1个页表，能有效降低页表数量，有效扩大TLB（Translation Lookaside Buffer）缓存的地址范围，从而提升离散访问的性能。

  TLB是昇腾AI处理器中用于高速缓存的硬件模块，用于存储最近使用的虚拟地址到物理地址的映射。

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

用户使用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

示例二：

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True,base_addr_aligned_kb:16

示例三：

export PYTORCH_NPU_ALLOC_CONF=page_size:1g

使用约束

• expandable_segments特性需在Ascend HDK 23.0.0及以上版本上使用。
• max_split_size_mb和garbage_collection_threshold中任意一个为非默认值时，expandable_segments必须设置为False。
• page_size特性要求在Ascend HDK 25.0.RC1及以上版本、CANN商发8.1.RC1及以上版本使用，支持如下产品：
  • Atlas 训练系列产品
  • Atlas A3 训练系列产品
• page_size特性与其他特性不支持同时配置，且申请内存注意事项如下：
  • 当申请内存大于1M时：
    • 若配置page_size，内存申请粒度为1G。
    • 若未配置page_size，内存申请粒度为2M。
  • 当申请内存小于等于1M时：配置page_size也不生效，内存申请粒度为2M。

支持的型号

• Atlas 训练系列产品
• Atlas A2 训练系列产品
• Atlas A3 训练系列产品