一、NPU 存储架构全景1.1 三级存储体系理解 NPU 的存储架构是做好内存优化的前提。昇腾 NPU 有三级存储每一级的容量、带宽、延迟差异巨大HBMHigh Bandwidth Memory是 NPU 的主存类似于 GPU 的显存。Ascend 910B 配备 64GB HBM带宽约 1.6TB/s。所有模型参数、中间结果、输入输出数据最终都存在这里。HBM 的容量决定了你能跑多大的模型。L2 Cache是片上缓存容量在几 MB 到十几 MB 之间带宽远高于 HBM。L2 Cache 主要用于向量计算单元Vector Core的临时数据存放。当向量算子需要反复读写同一块数据时把数据放在 L2 里比反复访问 HBM 快得多。L1 CacheSRAM是片上静态随机存储器容量最小通常 128KB 左右但延迟最低、带宽最高。L1 Cache 直接服务 Cube 核矩阵计算单元矩阵乘法的中间累加结果就存在这里。三级存储之间的数据搬运通过 DMADirect Memory Access引擎完成。DMA 搬运不占用计算单元可以和计算并行执行——这就是流水线优化的物理基础。1.2 数据搬运的代价很多人低估了数据搬运的开销。以一次典型的卷积计算为例从 HBM 读取输入数据到 SRAM假设数据量是 1MBHBM 带宽 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_id0,pool_ratio0.8):self.device_iddevice_id self.pool_sizeself._get_device_memory()*pool_ratio self.allocated{}self.free_blocks[(0,self.pool_size)]# (offset, size)self.total_allocated0def_get_device_memory(self):获取设备总显存模拟return64*1024*1024*1024# 64GBdefmalloc(self,size:int,tag:str)-int:从池中分配内存 返回: 虚拟偏移地址实际地址 base offset # 对齐到 32 字节aligned_size(size31)//32*32# First-fit 查找fori,(offset,block_size)inenumerate(self.free_blocks):ifblock_sizealigned_size:# 从空闲块中切分self.free_blocks.pop(i)ifblock_sizealigned_size:self.free_blocks.append((offsetaligned_size,block_size-aligned_size))self.allocated[offset]{size:aligned_size,tag:tag,}self.total_allocatedaligned_sizereturnoffset# 没有足够大的空闲块尝试合并self._merge_free_blocks()# 再试一次fori,(offset,block_size)inenumerate(self.free_blocks):ifblock_sizealigned_size:self.free_blocks.pop(i)ifblock_sizealigned_size:self.free_blocks.append((offsetaligned_size,block_size-aligned_size))self.allocated[offset]{size:aligned_size,tag:tag,}self.total_allocatedaligned_sizereturnoffsetraiseMemoryError(f显存不足: 需要{aligned_size}bytes, 池中最大空闲块 f{max(b[1]forbinself.free_blocks)ifself.free_blockselse0}bytes)deffree(self,offset:int):释放内存ifoffsetnotinself.allocated:returninfoself.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(keylambdax:x[0])merged[self.free_blocks[0]]foroffset,sizeinself.free_blocks[1:]:prev_offset,prev_sizemerged[-1]ifprev_offsetprev_sizeoffset:merged[-1](prev_offset,prev_sizesize)else:merged.append((offset,size))self.free_blocksmergeddefstats(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_free0:return0.0max_blockmax(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优化器状态Adam1400GB → 700GBFP32 主权重 两个 FP16 动量总计2100GB → 1050GB显存节省了但精度不能丢。所以需要 FP32 主权重 FP16 计算的混合精度策略。5.3 显存监控与预警classMemoryMonitor:显存使用监控 持续监控显存使用超过阈值时告警。 支持预测根据增长趋势预测何时会 OOM。 def__init__(self,warning_ratio0.85,critical_ratio0.95):self.warning_ratiowarning_ratio self.critical_ratiocritical_ratio self.history[]defcheck(self,current_usage_gb:float,total_gb:float64.0):usage_ratiocurrent_usage_gb/total_gb self.history.append({usage_gb:current_usage_gb,ratio:usage_ratio,})ifusage_ratioself.critical_ratio:returnCRITICAL,f显存使用{usage_ratio:.1%}即将 OOMelifusage_ratioself.warning_ratio:returnWARNING,f显存使用{usage_ratio:.1%}建议减少 batch sizeelse:returnOK,f显存使用{usage_ratio:.1%}defpredict_oom_time(self,growth_rate_gb_per_sec:float,total_gb:float64.0)-float:预测 OOM 时间ifnotself.history:returnfloat(inf)currentself.history[-1][usage_gb]remainingtotal_gb-currentifgrowth_rate_gb_per_sec0:returnfloat(inf)returnremaining/growth_rate_gb_per_sec六、常见问题问题原因解决方案推理 OOMbatch size 太大或序列太长减小 batch / 用 KV Cache 分页训练 OOM模型太大放不进单卡用 ZeRO 分片 流水线并行碎片化严重频繁分配释放不同大小的块用内存池 分桶分配显存泄漏张量被引用但未释放检查引用计数用弱引用相关仓库CANN- 昇腾计算架构 https://gitee.com/ascend/cannDeepSpeed- ZeRO 显存优化 https://github.com/microsoft/DeepSpeedPyTorch Memory- 显存管理 API https://pytorch.org/docs/stable/torch_cuda_memory.htmlCANN Profiling- 显存分析工具 https://gitee.com/ascend/cann
CANN 显存管理与内存优化：NPU 存储体系的深度剖析
一、NPU 存储架构全景1.1 三级存储体系理解 NPU 的存储架构是做好内存优化的前提。昇腾 NPU 有三级存储每一级的容量、带宽、延迟差异巨大HBMHigh Bandwidth Memory是 NPU 的主存类似于 GPU 的显存。Ascend 910B 配备 64GB HBM带宽约 1.6TB/s。所有模型参数、中间结果、输入输出数据最终都存在这里。HBM 的容量决定了你能跑多大的模型。L2 Cache是片上缓存容量在几 MB 到十几 MB 之间带宽远高于 HBM。L2 Cache 主要用于向量计算单元Vector Core的临时数据存放。当向量算子需要反复读写同一块数据时把数据放在 L2 里比反复访问 HBM 快得多。L1 CacheSRAM是片上静态随机存储器容量最小通常 128KB 左右但延迟最低、带宽最高。L1 Cache 直接服务 Cube 核矩阵计算单元矩阵乘法的中间累加结果就存在这里。三级存储之间的数据搬运通过 DMADirect Memory Access引擎完成。DMA 搬运不占用计算单元可以和计算并行执行——这就是流水线优化的物理基础。1.2 数据搬运的代价很多人低估了数据搬运的开销。以一次典型的卷积计算为例从 HBM 读取输入数据到 SRAM假设数据量是 1MBHBM 带宽 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_id0,pool_ratio0.8):self.device_iddevice_id self.pool_sizeself._get_device_memory()*pool_ratio self.allocated{}self.free_blocks[(0,self.pool_size)]# (offset, size)self.total_allocated0def_get_device_memory(self):获取设备总显存模拟return64*1024*1024*1024# 64GBdefmalloc(self,size:int,tag:str)-int:从池中分配内存 返回: 虚拟偏移地址实际地址 base offset # 对齐到 32 字节aligned_size(size31)//32*32# First-fit 查找fori,(offset,block_size)inenumerate(self.free_blocks):ifblock_sizealigned_size:# 从空闲块中切分self.free_blocks.pop(i)ifblock_sizealigned_size:self.free_blocks.append((offsetaligned_size,block_size-aligned_size))self.allocated[offset]{size:aligned_size,tag:tag,}self.total_allocatedaligned_sizereturnoffset# 没有足够大的空闲块尝试合并self._merge_free_blocks()# 再试一次fori,(offset,block_size)inenumerate(self.free_blocks):ifblock_sizealigned_size:self.free_blocks.pop(i)ifblock_sizealigned_size:self.free_blocks.append((offsetaligned_size,block_size-aligned_size))self.allocated[offset]{size:aligned_size,tag:tag,}self.total_allocatedaligned_sizereturnoffsetraiseMemoryError(f显存不足: 需要{aligned_size}bytes, 池中最大空闲块 f{max(b[1]forbinself.free_blocks)ifself.free_blockselse0}bytes)deffree(self,offset:int):释放内存ifoffsetnotinself.allocated:returninfoself.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(keylambdax:x[0])merged[self.free_blocks[0]]foroffset,sizeinself.free_blocks[1:]:prev_offset,prev_sizemerged[-1]ifprev_offsetprev_sizeoffset:merged[-1](prev_offset,prev_sizesize)else:merged.append((offset,size))self.free_blocksmergeddefstats(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_free0:return0.0max_blockmax(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优化器状态Adam1400GB → 700GBFP32 主权重 两个 FP16 动量总计2100GB → 1050GB显存节省了但精度不能丢。所以需要 FP32 主权重 FP16 计算的混合精度策略。5.3 显存监控与预警classMemoryMonitor:显存使用监控 持续监控显存使用超过阈值时告警。 支持预测根据增长趋势预测何时会 OOM。 def__init__(self,warning_ratio0.85,critical_ratio0.95):self.warning_ratiowarning_ratio self.critical_ratiocritical_ratio self.history[]defcheck(self,current_usage_gb:float,total_gb:float64.0):usage_ratiocurrent_usage_gb/total_gb self.history.append({usage_gb:current_usage_gb,ratio:usage_ratio,})ifusage_ratioself.critical_ratio:returnCRITICAL,f显存使用{usage_ratio:.1%}即将 OOMelifusage_ratioself.warning_ratio:returnWARNING,f显存使用{usage_ratio:.1%}建议减少 batch sizeelse:returnOK,f显存使用{usage_ratio:.1%}defpredict_oom_time(self,growth_rate_gb_per_sec:float,total_gb:float64.0)-float:预测 OOM 时间ifnotself.history:returnfloat(inf)currentself.history[-1][usage_gb]remainingtotal_gb-currentifgrowth_rate_gb_per_sec0:returnfloat(inf)returnremaining/growth_rate_gb_per_sec六、常见问题问题原因解决方案推理 OOMbatch size 太大或序列太长减小 batch / 用 KV Cache 分页训练 OOM模型太大放不进单卡用 ZeRO 分片 流水线并行碎片化严重频繁分配释放不同大小的块用内存池 分桶分配显存泄漏张量被引用但未释放检查引用计数用弱引用相关仓库CANN- 昇腾计算架构 https://gitee.com/ascend/cannDeepSpeed- ZeRO 显存优化 https://github.com/microsoft/DeepSpeedPyTorch Memory- 显存管理 API https://pytorch.org/docs/stable/torch_cuda_memory.htmlCANN Profiling- 显存分析工具 https://gitee.com/ascend/cann
