显存碎片优化
问题分析
在大模型训练阶段,系统需大量显存以承载模型参数、中间计算结果、优化器状态及批量输入数据等关键信息。频繁的显存分配与释放操作可能导致显存碎片化现象,从而降低资源利用率并影响整体性能。
解决方案
通过精细化管理不同生命周期的张量(Tensor),有效缓解显存碎片化问题,提升训练效率。
解决思路
- 识别与分类:区分长短期生命周期的张量,针对性优化显存管理策略
- 长生命周期张量主要包括:模型初始化时生成的权重与梯度、前向传播过程中的激活值、未被及时清理的张量(保留至下一迭代步骤)、以及梯度收敛后生成的优化器状态张量。
- 短生命周期张量主要包括:即时计算张量、辅助计算张量等。
- 隔离与优化:采用多显存池技术,对不同生命周期的张量进行隔离管理
- 将识别出的长短生命周期张量放入不同的显存池分别管理。
- 对于长生命周期的大张量,实施精确匹配策略,为其分配与自身尺寸相匹配的显存块,结合多级分配策略,防止相关显存池产生碎片。
使用场景
此优化机制专为训练场景设计,特别适用于显存资源紧张的情况。当用户遇到“显存不足”(out of memory)错误,且reserved与allocated显存之间存在显著差距(例如,差距超过1GB)时,表明系统显存碎片化严重。启用本特性有助于减少显存碎片,避免显存不足问题,确保训练流程稳定高效。
Tried to allocated 3384.00 MiB (NPU 2; 61.22 GiB total capacity; 53.87 GiB already allocated; 53.87 GiB current activate; 1.59 GiB free; 56.60 GiB reserved in total by PyTorch)
上述输出显示,reserved与allocated显存之间的差值高达2.73GB,表明显存碎片化程度较高。此时,启动显存碎片优化机制显得尤为重要。
使用方法
设置如下环境变量,即开启显存碎片优化特性。
export MEMORY_FRAGMENTATION=1
使用效果
启用显存碎片优化后,系统将有效减少显存碎片,提升资源利用效率,确保大模型训练任务的顺利进行。
注意事项
- 原因分析
在显存资源充足的情况下,该特性倾向于为新的计算任务分配连续的显存空间,而非复用已分配但未完全使用的显存片段。这一策略可能导致在特定情形下与HCCL(Huawei Collective Communication Library,华为集合通信库)发生显存资源竞争。尤其当HCCL面临显存不足时,由于无法有效释放由torch预留的额外闲置空间,可能会触发“HCCL显存不足”的异常错误。
- 解决方法
为规避上述问题,建议通过调用torch_npu.npu.set_per_process_memory_fraction接口,预先设定PyTorch进程所能占用显存的上限比例。例如,设置x=0.94意味着PyTorch最多可占用单张显卡总显存的94%,以此实现对显存资源的有效管控。
配置示例:
# MindSpeed/mindspeed/core/memory/memory_fragmentation/memory_recorder.py import torch_npu torch_npu.npu.set_per_process_memory_fraction(0.94)