搞 AI 推理这一行很多人最后都会撞上同一堵墙算力明明还有余量显存却先爆了。我第一次在昇腾设备上调 transformer 推理时就是这个状态——算子性能、编译流程都跑通了结果正式压测第一批请求直接报内存分配失败。后来把 CANN 的内存管理机制从底到上摸了一遍才明白显存不足背后不止是容量问题还牵扯到分配策略、生命周期分析和复用效率。这篇文章就围绕 CANN 内存管理把显存花在哪、怎么压下去、出了问题怎么排查完整地捋一遍。适合正在搞昇腾推理、做模型转换或者想在低显存设备上跑大模型的同学参考。1. 先算清账AI 推理时显存到底花在了什么地方1.1 推理显存的“四本账”权重、激活、工作区、输出不管用哪套推理框架显存开销基本都能拆成四块权重内存模型参数本身占用的静态空间。7B 模型用 FP16 存就是约 7e9 × 2 字节 ≈ 14GB换成 INT8 约 7GBINT4 压到 3.5~4GB。这就是为什么 6G 显存跑 7B 级别 LLM 必须做量化压缩。激活值内存推理过程中每层算子输出的中间特征图。Transformer 场景下它和 batch、序列长度、隐藏层维度和层数强相关是动态峰值的大头。工作区内存算子执行时的临时缓冲GEMM、卷积、softmax 等算子都会申请 workspace。算子的实现策略不同workspace 可能从几十 KB 到几 GB 不等。输出内存最终结果和中间结果回拷需要的缓存通常规模不大但必须连续容易碎片化。很多同学只盯着权重占了多少忽略激活值和工作区结果模型文件明明只有 6GB上卡跑一次推理却直接吃掉 11GB。这不是框架有 bug而是你把动态部分漏算了。1.2 为什么“算力够用”却“显存不够”动态执行与静态编译的差异同一个模型用动态图框架跑和用静态图跑显存占用可以差出 30% 到一倍。原因在于管理方式完全不同。动态图是执行到哪个算子才给哪个算子的输出分配显存引用计数归零再释放。这种方式灵活但两个致命伤一是多个算子中间结果同时存活峰值是“所有同期活跃张量之和”二是内存大小变化剧烈容易碎成一片后面想找一块大连续内存就找不到。CANN 的 om 模型不是这个路子。ATC 转换时已经把算子图排好序每一个张量的“出生”和“死亡”时间都被编译器算清楚输出到同一块缓冲区的张量可以复用同一段内存。也就是说模型转换等于把显存的使用方案提前编排好运行时只是照着执行。这也是同样网络在静态图模式下更省显存的根本原因。当然静态图也不是万能的。如果开了动态 shape编译器必须按“最大可能 shape”预留内存这也是后面要重点说的隐性峰值来源。2. 拆解 CANN 内存管理从内存池到复用机制2.1 运行时内存池预分配、页大小和分配策略AscendCL 运行时并没有让每一次aclrtMalloc都直接向硬件申请物理内存而是先把显存切成一个内存池来管理。第一次aclrtMalloc时按策略取出一块返回给调用方aclrtFree也不是真的还给设备而是还给池子。这样做的好处是分配速度快、避免频繁和驱动交互、还能统一规划大页。用aclrtGetMemInfo可以随时查看可用显存#include acl/acl.h #include iostream size_t freeMem 0; size_t totalMem 0; aclError ret aclrtGetMemInfo(ACL_MEMTYPE_DEVICE, freeMem, totalMem); if (ret ACL_SUCCESS) { std::cout total: totalMem / 1024 / 1024 MB , free: freeMem / 1024 / 1024 MB std::endl; }这里ACL_MEMTYPE_DEVICE表示设备侧显存另有 host 侧内存类型对应aclrtMallocHost/aclrtFreeHost。调试时我会在模型加载前后各打一次看模型到底悄悄吃了多少。aclrtMalloc的分配策略有几种常见取值ACL_MEM_MALLOC_HUGE_FIRST优先尝试大页分配失败时降级普通页ACL_MEM_MALLOC_HUGE_ONLY只分配大页失败立刻报错ACL_MEM_MALLOC_NORMAL_ONLY只分配普通页。大页能减少 TLB miss对大块连续内存非常友好但大页资源本身有限。实战中我默认选HUGE_FIRST只有确认设备大页充足且对性能敏感时才用HUGE_ONLY。2.2 静态内存与动态内存谁决定了峰值占用ATC 转换出来的 om 模型在加载时会分成静态内存和动态内存两块。静态内存对应权重、常量、固定不动的输出缓冲区基本等于模型文件大小加载后常驻。动态内存是给激活值、中间结果、算子 workspace 预留的资源池运行期间被反复复用。峰值占用大致等于“静态内存 动态内存中被同时占用的最大值”而不是两者的简单相加。判断一个模型能不能塞进某张卡不能只看权重大小。正确流程是先用aclmdlQuerySize查一下 work 和 weight 的需求再叠加运行时那些激活值才是真实占用size_t workSize 0; size_t weightSize 0; aclError ret aclmdlQuerySize(modelPath, workSize, weightSize);workSize就是运行时工作内存需求weightSize是权重内存需求。很多人直接LoadFromFile一把梭其实这两个数字完全可以在加载前拿到用来决策“这张卡放不放得下”。2.3 内存复用与生命周期分析模型转换时发生了什么再往底层看“省显存”这件事主要靠编译期的生命周期分析。一个张量从首次被写入到最后一次被读取这段时间叫它的活跃区间。只要两个张量的活跃区间不重叠编译器就把它们安排到同一块物理内存上。算子融合的本质也是这个把多个算子的中间张量合并成内部缓冲不落显存或者把张量的活跃区间缩短让更多训练量可以复用同一块内存。ATC 提供了--buffer_optimize开关例如 L2 优化可以把部分中间结果放到片上 SRAM减少 HBM 的读写压力间接降低动态内存需求。不过这个开关不是无脑开最好。我遇到过开启后部分算子 kernel 行为异常的情况一般先关掉对比一版确认性能和显存差异再决定。3. 实操落地把显存占用压下去的组合拳3.1 显式管理模型内存aclmdlLoadFromFileWithMem的正确用法常规加载模型用aclmdlLoadFromFile内存由运行时内部处理。想要精细控制就用aclmdlLoadFromFileWithMem自己把 work 和 weight 两块内存准备好void* workMem nullptr; void* weightMem nullptr; // 1. 先查需要多大 size_t workSize 0; size_t weightSize 0; aclmdlQuerySize(modelPath, workSize, weightSize); // 2. 按需分配 aclrtMalloc(workMem, workSize, ACL_MEM_MALLOC_HUGE_FIRST); aclrtMalloc(weightMem, weightSize, ACL_MEM_MALLOC_HUGE_FIRST); // 3. 使用指定内存加载 uint32_t modelId 0; aclmdlLoadFromFileWithMem(modelPath, modelId, workMem, workSize, weightMem, weightSize);这样做的价值在于两块内存来自我自己管理的内存池可以在多次加载之间复用避免反复申请释放造成的碎片。尤其在高并发服务里进程长时间运行每次模型重载都能少产生一批不连续空洞。注意不同 CANN 版本的接口签名可能有差异实操前先看一眼头文件里的声明。3.2 模型转换期的内存优化参数ATC 转换是压缩显存的第一道关口。我总结了三个最常用的方向固定 shape。明确--input_shapeinput:1,3,224,224这样写编译器不需要为动态维度留余量动态内存会明显减小。固定不了的场景再考虑动态 batch 或动态分辨率但要设置合理的取值范围。动态维度的上下界。--dynamic_batch_size和--dynamic_image_size的值要给实数别为了灵活把上限抬得过高。上限每高一倍预留内存可能翻倍但实际用不到。内存优化开关和算子融合。--buffer_optimize控制缓冲区优化--enable_small_channel对小通道网络有一定融合和内存优化作用。每换一个模型都值得尝试三组配置全关、仅 buffer 优化、全开跑一遍实际显存曲线再决定。真实经验是灵活性和内存占用永远是矛盾的。追求极致显存就把 shape 定死追求通用就接受一部分内存浪费。3.3 框架侧设置torch_npu 与 MindSpore 的运行期配置如果你在 PyTorch 生态上走昇腾torch_npu 提供了类似 CUDA 的缓存分配器显存管理在框架层还有一层。环境变量PYTORCH_NPU_ALLOC_CONF可以控制分配器的分块大小和回收阈值常见的max_split_size_mb可以限制大块显存被拆分的粒度expandable_segments:True则允许段按需扩展对降低峰值和碎片都有帮助。MindSpore 侧可以在set_context里配置memory_optimize_level开启更积极的内存复用适合对 DVPP 和算子 workspace 都做了充分优化的静态图训练与推理。框架层优化的通用思路先测默认表现再一项项加配置看到显存曲线平滑、峰值没有反复抬升就是配置合理的信号。4. 进阶挑战低显存跑大模型与动态 Shape 优化4.1 KV Cache 有多大先算再跑大模型推理和普通 CNN 最大的区别是多了 KV Cache。它存的是历史 token 的 key 和 value供注意力机制重复读取。计算公式KV Cache 内存 2K 和 V × 层数 × batch × 序列长度 × 每 token 字节数举个具体例子隐藏维度 4096、32 层 transformerbatch1、seq_len4096、FP16KV Cache 约 2 × 32 × 4096 × 4096 × 2 字节大约 2.1GB。权重已经吃掉不少显存KV Cache 再占两块6G 卡跑长序列基本不可能。低显存跑大模型的几个现实手段权重做 INT8/INT4 量化KV Cache 用 INT8 甚至 FP8限制最大推理长度。很多人只盯着权重量化忽略 KV Cache 同样能压。我一般在固定显存预算下先算两笔账权重多少、KV Cache 在目标并发和序列长度下多少加起来接近卡容量上限就调整量化档位或并发数。4.2 分页 KV Cache 与连续批处理减少空闲浪费传统推理给每个请求预留完整长度的 KV Cache实际跑短序列时大量显存闲置。分页 KV Cache 的解决思路类似于操作系统虚拟内存物理块按需分配逻辑上连续物理上可以散落。昇腾平台上做 LLM 服务时选择 MindIE 这类支持分页 KV Cache 的推理引擎通常比手动管理缓存省下可观内存还能支持连续批处理把同一时刻的请求效率拉满。这个思路同样可以迁移到普通模型如果推理引擎支持显存按需扩容而不是一次性申请最大可能容量多路请求交错服务时的内存峰值会低很多。expandable_segments能缓解的正是这类问题。4.3 多模型并发与多卡部署的内存隔离当显存仍然不够时很多人第一反应是换大卡其实先看内存分配是否合理更划算。多模型并发部署时我给每个进程用ASCEND_RT_VISIBLE_DEVICES指定独立设备避免两个进程争抢同一设备的碎片化内存。单卡多模型则要留意两个模型都加载后静态内存是叠加的动态内存池互相隔离总碎片可能会很高。我的排查顺序是先npu-smi info看 HBM 总占用和进程分布确认有没有残留进程占着显存再检查每个模型加载前后的 free 变化定位哪一块超预算最后才是调分配策略或换卡。5. 常见问题与排查技巧实录5.1 内存不足报错先查设备、进程、页分配昇腾设备上报内存不足形式各有不同有的直接aclrtMalloc failed有的报“HBM memory exhausted”。别急着改代码先按顺序排除npu-smi info看 HBM 使用率是否有上个任务残留进程没有退出。确认可见设备是否正确。ASCEND_RT_VISIBLE_DEVICES设置不对进程可能被分到被别人占满的设备上。确认分配策略是否过强。ACL_MEM_MALLOC_HUGE_ONLY在大页不足时必挂换成HUGE_FIRST再看。确认模型转换时的动态 shape 上限是否过大把上限降低再看。这四步能解决九成“昨天还能跑今天内存不足”的问题。5.2 推理进程显存缓慢上涨泄漏排查套路显存不掉反涨通常不是算法变复杂而是资源没释放。重点查三处模型是否多次aclmdlLoad而只aclmdlUnload一次每次请求创建的 stream 是否在结束时aclrtDestroyStreamaclrtMalloc出来的内存有没有成对aclrtFree。CANN 自带的 msprof 工具可以采集设备侧内存使用曲线结合日志里的分配记录能定位到是哪个阶段在持续申请内存。还想更快的话在代码里给每块aclrtMalloc打点记录大小和调用栈跑一万个请求之后统计存活内存谁没归还一目了然。这类问题很多是 host 侧内存不是显存。别忘了检查aclrtMallocHost对应aclrtFreeHost很多“显存泄漏”其实是 pin 内存泄漏。5.3 动态 Shape 导致的“隐性峰值”与我的几个优化心得动态 shape 的坑在于编译器和运行时是按上下限提前预留的。你把 batch 上限设成 8即使平时只跑 1系统依然按 8 的规模留内存。这个预留又不会直观反应在模型文件大小上所以特别隐蔽。我的习惯是先量化实际流量再反推 shape 范围。服务端推理场景如果 99% 请求都是单条短文本就直接固定 batch1动态维度给个安全但不过分的上限。宁可偶尔重转一版模型也不要让显存长期被用不到的上限锁死。5.4 一张速查表显存排查常见问题与对策现象最可能原因优先处理方式加载模型即失败权重work 超出剩余显存aclmdlQuerySize先查必要时量化权重跑批处理时偶发失败激活值/动态内存峰值超限固定 shape降低动态上限开--buffer_optimize长期运行缓慢涨显存资源未释放检查 load/unload 和 stream 生命周期用 msprof 定位开HUGE_ONLY必失败大页资源不足改用HUGE_FIRST或关闭大页强制分配多模型并发互相影响设备隔离不彻底用ASCEND_RT_VISIBLE_DEVICES绑定固定设备最后分享一个小技巧如果显存实在压不下来先把 batch 设成 1把输入 shape 写死再看算子 workspace 的优化空间。很多模型不是真的塞不下是默认配置太浪费。我自己调一个 7B 模型在 16G 昇腾上稳定跑到 8 并发核心就是四件事固定 shape、开启 buffer 优化、显式管理权重内存、KV Cache 设置上限。你把这几个点逐个过一遍大概率也能从天天报显存不足变成看着npu-smi info气定神闲。