当前位置：首页 > news >正文

CANN 显存管理与内存优化：NPU 存储体系的深度剖析

news 2026/5/25 2:11:31

一、NPU 存储架构全景

1.1 三级存储体系

理解 NPU 的存储架构是做好内存优化的前提。昇腾 NPU 有三级存储，每一级的容量、带宽、延迟差异巨大：

HBM（High Bandwidth Memory）是 NPU 的主存，类似于 GPU 的显存。Ascend 910B 配备 64GB HBM，带宽约 1.6TB/s。所有模型参数、中间结果、输入输出数据最终都存在这里。HBM 的容量决定了你能跑多大的模型。

L2 Cache是片上缓存，容量在几 MB 到十几 MB 之间，带宽远高于 HBM。L2 Cache 主要用于向量计算单元（Vector Core）的临时数据存放。当向量算子需要反复读写同一块数据时，把数据放在 L2 里比反复访问 HBM 快得多。

L1 Cache（SRAM）是片上静态随机存储器，容量最小（通常 128KB 左右），但延迟最低、带宽最高。L1 Cache 直接服务 Cube 核（矩阵计算单元），矩阵乘法的中间累加结果就存在这里。

三级存储之间的数据搬运通过 DMA（Direct Memory Access）引擎完成。DMA 搬运不占用计算单元，可以和计算并行执行——这就是流水线优化的物理基础。

1.2 数据搬运的代价

很多人低估了数据搬运的开销。以一次典型的卷积计算为例：

从 HBM 读取输入数据到 SRAM，假设数据量是 1MB，HBM 带宽 1.6TB/s，理论耗时 0.6 微秒。但实际上，HBM 的有效带宽受 bank conflict、刷新周期等因素影响，实际只能达到理论值的 60-70%。再加上 DMA 引擎的调度开销，实际搬运时间大约是理论值的 2-3 倍。

也就是说，一次 1MB 数据的搬运实际需要 1-2 微秒。如果一个算子被拆成 10 次小搬运，总开销就变成了 10-20 微秒。而一次 16x16x16 的矩阵乘计算只需要 0.01 微秒左右。搬运开销远大于计算开销——这就是为什么算子融合如此重要。

1.3 为什么手动管理内存

操作系统有自动内存管理（malloc/free），为什么 NPU 还需要手动管理？原因有三：

第一，NPU 上没有虚拟内存机制。操作系统管理的是 CPU 的内存，NPU 的 HBM 是独立的地址空间，malloc 无法直接分配。

第二，自动管理的开销不可接受。每次 malloc/free 都需要维护空闲链表、检查碎片、执行对齐，这些操作本身就要消耗宝贵的计算资源。

第三，确定性延迟。推理服务对延迟极其敏感，自动内存管理可能在某个时刻触发垃圾回收或碎片整理，导致延迟突增。手动管理可以保证每次分配都是确定性的。

二、显存分配策略

2.1 静态分配 vs 动态分配

静态分配在模型编译阶段就确定所有中间结果的内存地址。ATC 编译器在转换模型时，会分析每个算子的输入输出张量大小，计算出内存使用峰值，然后一次性分配。静态分配的优点是零运行时开销，缺点是不够灵活——如果 batch size 变了，可能需要重新编译。

动态分配在运行时按需分配和释放。适用于动态 shape 的场景，比如输入长度不固定。动态分配的缺点是每次分配都有开销，而且容易产生内存碎片。

实际生产中，推理服务通常用静态分配（确定性好），训练场景用动态分配（灵活性高）。

2.2 内存池化

内存池化是减少分配开销的经典方法。预先申请一大块内存（池），然后从中切分小块给各个算子使用。算子释放内存时，不是真正归还给系统，而是标记为可复用，下次分配时优先从池中取。

classNPUDeviceMemoryPool:"""NPU 显存池 预分配策略: - 默认分配 80% 的可用显存作为池 - 留 20% 给临时分配和系统开销 分配策略: - 优先使用空闲块（first-fit） - 找不到合适大小则合并相邻空闲块 - 仍然不够则向系统申请新块 释放策略: - 不真正归还，只标记为空闲 - 相邻空闲块自动合并 - 定期清理长期未使用的块 """def__init__(self,device_id=0,pool_ratio=0.8):self.device_id=device_id self.pool_size=self._get_device_memory()*pool_ratio self.allocated={}self.free_blocks=[(0,self.pool_size)]# (offset, size)self.total_allocated=0def_get_device_memory(self):"""获取设备总显存（模拟）"""return64*1024*1024*1024# 64GBdefmalloc(self,size:int,tag:str="")->int:"""从池中分配内存 返回: 虚拟偏移地址（实际地址 = base + offset） """# 对齐到 32 字节aligned_size=(size+31)//32*32# First-fit 查找fori,(offset,block_size)inenumerate(self.free_blocks):ifblock_size>=aligned_size:# 从空闲块中切分self.free_blocks.pop(i)ifblock_size>aligned_size:self.free_blocks.append((offset+aligned_size,block_size-aligned_size))self.allocated[offset]={'size':aligned_size,'tag':tag,}self.total_allocated+=aligned_sizereturnoffset# 没有足够大的空闲块，尝试合并self._merge_free_blocks()# 再试一次fori,(offset,block_size)inenumerate(self.free_blocks):ifblock_size>=aligned_size:self.free_blocks.pop(i)ifblock_size>aligned_size:self.free_blocks.append((offset+aligned_size,block_size-aligned_size))self.allocated[offset]={'size':aligned_size,'tag':tag,}self.total_allocated+=aligned_sizereturnoffsetraiseMemoryError(f"显存不足: 需要{aligned_size}bytes, 池中最大空闲块 "f"{max(b[1]forbinself.free_blocks)ifself.free_blockselse0}bytes")deffree(self,offset:int):"""释放内存"""ifoffsetnotinself.allocated:returninfo=self.allocated.pop(offset)self.total_allocated-=info['size']# 加入空闲列表self.free_blocks.append((offset,info['size']))def_merge_free_blocks(self):"""合并相邻的空闲块"""iflen(self.free_blocks)<2:returnself.free_blocks.sort(key=lambdax:x[0])merged=[self.free_blocks[0]]foroffset,sizeinself.free_blocks[1:]:prev_offset,prev_size=merged[-1]ifprev_offset+prev_size==offset:merged[-1]=(prev_offset,prev_size+size)else:merged.append((offset,size))self.free_blocks=mergeddefstats(self)->dict:"""内存使用统计"""return{'total_pool':self.pool_size,'allocated':self.total_allocated,'free':self.pool_size-self.total_allocated,'utilization':self.total_allocated/self.pool_size,'fragments':len(self.free_blocks),'largest_free':max(b[1]forbinself.free_blocks)ifself.free_blockselse0,}

三、内存碎片问题

3.1 碎片的成因

内存碎片是 NPU 显存管理中最棘手的问题之一。当频繁分配和释放不同大小的内存块时，空闲内存会被切割成很多小块，这些小块彼此不相邻，无法合并成大块。即使总空闲内存足够，也可能因为没有连续的大块而分配失败。

碎片分两种：外部碎片是空闲块总大小够用，但每个块都太小，无法满足一次大分配。内部碎片是分配的内存块比实际需要的大（因为对齐），多出来的部分浪费了。

3.2 碎片率计算

defcalc_fragmentation(free_blocks:list,total_free:int)->float:"""计算内存碎片率 碎片率 = 1 - (最大连续空闲块 / 总空闲内存) 碎片率为 0 表示完全没有碎片（所有空闲内存连在一起） 碎片率为 1 表示碎片极其严重（没有足够大的连续块） """iftotal_free==0:return0.0max_block=max(b[1]forbinfree_blocks)iffree_blockselse0return1.0-(max_block/total_free)

3.3 碎片应对策略

分桶分配是减少碎片的有效方法。按大小将内存块分成几个桶（比如 <1KB、1-16KB、16-256KB、>256KB），每个桶独立管理。分配时根据请求大小选择合适的桶，这样小块不会污染大块的空闲空间。

对齐分配虽然会增加内部碎片，但能显著减少外部碎片。将所有分配对齐到 32 字节或 64 字节，保证相邻块的地址自然对齐，便于后续合并。

定期压缩在空闲时触发，将所有已分配的块移动到连续的区域，腾出一大块完整的空闲空间。压缩的代价是需要暂停推理服务，所以通常在请求低谷期执行。

四、推理场景的显存优化

4.1 KV Cache 复用

Transformer 模型在自回归生成时，每生成一个 token 都需要之前的 KV Cache。如果每步都重新计算 KV Cache，计算量会随序列长度线性增长。实际做法是把每步计算出的 KV Cache 缓存起来，下一步直接复用。

KV Cache 的大小 = 2 × 层数 × 序列长度 × 隐藏维度 × 精度字节数。对于 LLaMA-70B，序列长度 4096 时，KV Cache 约占 40GB——比模型参数本身还大。所以 KV Cache 的管理直接影响能支持的最大序列长度和并发数。

4.2 激活值检查点

训练时，反向传播需要前向传播的中间结果（激活值）。保存所有激活值的显存开销巨大。激活值检查点（Activation Checkpointing）的思路是：只保存部分关键层的激活值，其他层的激活值在反向传播时重新计算。

这本质上是用计算换显存。重新计算前向传播的开销大约增加 30-40%，但显存占用可以减少 60-80%。对于显存受限的大模型训练，这是非常划算的交易。

4.3 原地操作

原地操作（In-place Operation）让算子直接覆盖输入张量的内存，而不是分配新的输出张量。比如 ReLU 操作，输出和输入形状完全一样，完全可以直接在输入的内存上修改。

原地操作的显存节省 = 输出张量大小。对于一个 1x3x224x224 的 float32 张量，原地操作一次节省 600KB。在整个网络中累积，节省的显存非常可观。

五、训练场景的显存优化

5.1 梯度累积

当 batch size 太大放不进显存时，可以分多次前向+反向，累积梯度，最后一次性更新参数。这等价于用小 batch 模拟大 batch 的效果。

比如目标 batch size 是 128，但显存只能容纳 32，那就做 4 次前向+反向，每次用 32 个样本，梯度累加 4 次后再更新。等效的 learning rate 也需要相应调整。

5.2 混合精度的显存收益

FP16 相比 FP32，每个参数的显存占用减半。模型参数 + 梯度 + 优化器状态，总共可以节省约 50% 的显存。具体来说：

模型参数：350GB → 175GB
梯度：350GB → 175GB
优化器状态（Adam）：1400GB → 700GB（FP32 主权重 + 两个 FP16 动量）
总计：2100GB → 1050GB

显存节省了，但精度不能丢。所以需要 FP32 主权重 + FP16 计算的混合精度策略。

5.3 显存监控与预警

classMemoryMonitor:"""显存使用监控 持续监控显存使用，超过阈值时告警。 支持预测：根据增长趋势预测何时会 OOM。 """def__init__(self,warning_ratio=0.85,critical_ratio=0.95):self.warning_ratio=warning_ratio self.critical_ratio=critical_ratio self.history=[]defcheck(self,current_usage_gb:float,total_gb:float=64.0):usage_ratio=current_usage_gb/total_gb self.history.append({'usage_gb':current_usage_gb,'ratio':usage_ratio,})ifusage_ratio>=self.critical_ratio:return'CRITICAL',f'显存使用{usage_ratio:.1%}，即将 OOM'elifusage_ratio>=self.warning_ratio:return'WARNING',f'显存使用{usage_ratio:.1%}，建议减少 batch size'else:return'OK',f'显存使用{usage_ratio:.1%}'defpredict_oom_time(self,growth_rate_gb_per_sec:float,total_gb:float=64.0)->float:"""预测 OOM 时间"""ifnotself.history:returnfloat('inf')current=self.history[-1]['usage_gb']remaining=total_gb-currentifgrowth_rate_gb_per_sec<=0:returnfloat('inf')returnremaining/growth_rate_gb_per_sec

六、常见问题

问题	原因	解决方案
推理 OOM	batch size 太大或序列太长	减小 batch / 用 KV Cache 分页
训练 OOM	模型太大放不进单卡	用 ZeRO 分片 + 流水线并行
碎片化严重	频繁分配释放不同大小的块	用内存池 + 分桶分配
显存泄漏	张量被引用但未释放	检查引用计数，用弱引用