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,此设置将指示缓存分配器创建特定的内存块分配,这些内存块后续可以扩展,以便能更好地处理内存使用中频繁变更使用内存大小的情况。如果设置为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
示例二:
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True,base_addr_aligned_kb:16
使用约束
expandable_segments特性需在Ascend HDK 23.0.0及以上版本上使用。
支持的型号
Atlas 训练系列产品
Atlas A2 训练系列产品