大数跨境

匠心时刻丨MindStudio-MemScope片上内存调优实战指南

匠心时刻丨MindStudio-MemScope片上内存调优实战指南 昇腾AI开发者
2026-07-10
2

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


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



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


通用内存调优技巧



Python场景:善用垃圾回收


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


代码示例:


import gcgc.collect(0) # 回收第0代(新创建的短期对象,最常用)gc.collect(1) # 回收第1gc.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 msmemscopemsmemscope.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 msmemscopemsmemscope.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 msmemscopemsmemscope.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





扫描二维码
一键直达msMemScope开源社区





扫描二维码
一键直达MindStudio开源技术交流群


【声明】内容源于网络
0
0
昇腾AI开发者
昇腾社区
内容 977
粉丝 0
昇腾AI开发者 昇腾社区
总阅读5.1k
粉丝0
内容977