NPU内存持续劣化、内存占用居高不下，该如何有效优化？NPU内存资源本就十分宝贵，但不少开发者往往难以掌握真实内存使用状态：不清内存资源的消耗去向，无法定位OOM报错根因，也缺少系统化的调优手段。针对以上痛点，本文将基于昇腾MindStudio-MemScope内存分析工具，全面梳理并讲解内存调优的完整实践方法。

在开始调优之前，需要先搞清楚问题的类型。内存问题主要分为以下几类：

搞清楚问题类型，才能对症下药。接下来，本文分场景介绍内存调优方法。

通用内存调优技巧

Python场景：善用垃圾回收

Python的垃圾回收机制（GC）是内存管理的好帮手。主动调用gc.collect()可以回收不可达对象，但要注意：频繁调用可能影响性能。

代码示例：

import gc
gc.collect(0) # 回收第0代（新创建的短期对象，最常用）
gc.collect(1) # 回收第1代
gc.collect(2) # 回收第2代（长期存活对象，尽量少触发）

小贴士

建议在关键节点（如每个epoch结束后）手动触发GC，而不是频繁调用。

PyTorch NPU场景：这些环境变量设置了吗？

在昇腾AI基础软硬件平台上使用PyTorch，有几个关键的环境变量可以显著优化内存使用：

内存缓存回收阈值

import torch_nputorch_npu.npu.set_per_process_memory_fraction(0.95)

这个参数设置触发内存缓存回收的使用率（0~1）。建议从高到低尝试，太高可能在大块内存申请时OOM，太低则频繁触发影响性能。

垃圾回收阈值

export PYTORCH_NPU_ALLOC_CONF="garbage_collection_threshold:0.95"

内存达到阈值时自动触发GC，同样建议由高到低尝试。

虚拟内存开关

export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True"

开启后可以有效减少内存碎片，强烈推荐！

内存块切分上限

export PYTORCH_NPU_ALLOC_CONF="max_split_size_mb:50"

设置合适的值可以减少碎片，但设置太大会推高内存峰值。建议采集内存数据后根据实际情况调整。

多流内存复用

export MULTI_STREAM_MEMORY_REUSE=1

多流场景下的神器，可以显著提高内存回收效率。

ATB算子场景：选择合适的内存分配算法

export ATB_WORKSPACE_MEM_ALLOC_ALG_TYPE=3

ATB提供了4种workspace内存分配算法（0~3），推荐使用算法3（引入block合并），在大多数场景下表现最优。

进阶调优：用msMemScope透视内存黑盒

环境变量调优只是第一步，当需要深入分析内存问题时，就需要专业的工具了。msMemScope是昇腾AI基础软硬件平台上的内存分析利器，具备以下特点：

便捷易用：Python API方式，功能自由配置，动态启停

多维标记：支持权重、激活值、梯度、优化器状态等tensor标记

全量采集：支持对所有内存事件的申请、释放、使用情况进行采集

下面介绍两个核心功能：内存拆解和低效内存识别。

内存拆解

内存拆解功能可以将内存占用拆分为多个模块（如权重、激活值、梯度、优化器状态等），帮助用户快速定位内存占用异常的模块。

快速上手代码：

import msmemscope
msmemscope.config(    
	data_format='db',
	analysis='decompose',  # 开启内存拆解
)
msmemscope.start()
train_model()  # 训练/推理代码
msmemscope.stop()

采集结果可以用通过MindStudio Insight可视化，示例如下：

实战案例：vLLM推理场景内存拆解

以vLLM Ascend服务化推理为例，我们可以按权重、激活值、KV Cache等维度进行拆解：

Step 1：在关键位置添加标记

在vllm_ascend::worker_v1.py添加数据采集操作，导入msMemScope模块，并设置采集的范围。

def load_model(self) -> None:
        import msmemscope
msmemscope.config(    
	data_format='db',
	analysis='decompose',  # 开启内存拆解
        )
        rank = torch_npu.npu.current_device()
        if rank == 0:
            msmemscope.start()
        if self.vllm_config.model_config.enable_sleep_mode:
            allocator = CaMemAllocator.get_instance()
            assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
            context = allocator.use_memory_pool(tag="weights")
        else:
            from contextlib import nullcontext
            context = nullcontext()  # type: ignore
        with context, set_current_vllm_config(self.vllm_config):
            self.model_runner.load_model()

vllm_ascend：model_runner_v1.py添加内存标记操作

（1）kv_cache，函数名称：initialize_kv_cache，在这个函数开头和结束插入内存标记即可：

describer.describe(“UserDefined@vllm@kvcache”)describer.undescribe(“UserDefined@vllm@kvcache”)

（2）modelweight，在函数load_model中添加下图标记。

describer.describe(“UserDefined@vllm@modelweight”)describer.undescribe(“UserDefined@vllm@modelweight”)

（3）forward_activation，在函数execute_model中添加标记：

describer.describe(“UserDefined@vllm@activation”)describer.undescribe(“UserDefined@vllm@activation”)

（4）profile_run，在函数profile_run中添加标记。

describer.describe(“UserDefined@vllm@profilerun”)describer.undescribe(“UserDefined@vllm@profilerun”)

采集结果示例如下：

Step 2：可视化分析

将采集的DB文件导入MindStudio Insight，即可看到各模块的内存占用情况。重点关注峰值时刻的内存分布，找出异常模块，然后在内存详情列表中筛选该时间点、该模块申请的内存，即可得到内存详细信息。

低效内存识别

低效内存是指内存块申请后没有立即使用，或者使用结束后未及时释放的情况。简单说就是"占着茅坑不拉屎"，白白浪费内存资源。

msMemScope支持识别三种低效内存：

快速检测代码：

import msmemscope
msmemscope.config(    
	data_format='db',
	analysis='inefficient',  # 开启低效内存识别
)
msmemscope.start()
train_model()
msmemscope.stop()

识别结果会标注在输出文件的MALLOC事件中，通过inefficient字段标识具体类型。

总结内存调优三步走

1、定位问题：先搞清楚是内存劣化、OOM还是内存过高

2、基础调优：设置合适的环境变量，开启虚拟内存等优化选项

3、深度分析：使用msMemScope进行内存拆解和低效内存识别

内存调优是一门"省钱"的艺术。掌握这些技巧，不仅能解决OOM问题，更能提升内存利用率，让你的每一分算力投入都物有所值。

msMemScope工具已开源，欢迎访问项目主页了解更多：https://gitcode.com/Ascend/msmemscope

推荐阅读

msInsight工具：https://gitcode.com/Ascend/msinsight