一、NPU 显存架构1.1 显存层次┌─────────────────────────────────────────┐ │ 昇腾 NPU 显存架构 │ ├─────────────────────────────────────────┤ │ │ │ HBM (High Bandwidth Memory) │ │ ├─ 大容量: 32-64 GB │ │ ├─ 高带宽: 1.2 TB/s │ │ └─ 用于: 模型权重、激活值 │ │ │ │ DDR (Double Data Rate) │ │ ├─ 容量: 128 GB │ │ ├─ 带宽: 50 GB/s │ │ └─ 用于: 大模型权重、数据缓存 │ │ │ │ L2 Cache │ │ ├─ 容量: 8 MB │ │ ├─ 带宽: 极高 │ │ └─ 用于: 频繁访问的数据 │ │ │ └─────────────────────────────────────────┘1.2 内存分配流程用户请求分配: │ ├─ HBM 分配 (优先) │ ├─ 检查 HBM 空闲块 │ ├─ 找到合适块 → 分配 │ └─ 未找到 → 尝试 DDR │ ├─ DDR 分配 │ ├─ 检查 DDR 空闲块 │ ├─ 找到合适块 → 分配 │ └─ 未找到 → 触发回收 │ └─ 回收机制 ├─ 释放未使用的缓存 ├─ 合并碎片 └─ OOM 错误二、显存监控2.1 查询显存状态# 查看 NPU 显存使用npu-smi info# 输出示例:# NPU ID: 0# Name: Ascend910B# Memory Size: 64 GB# Memory Used: 45 GB# Memory Free: 19 GB# Utilization: 68%# 持续监控watch-n1npu-smi info2.2 Python 显存监控importtorchimporttorch.npuclassNpuMemoryMonitor:def__init__(self):self.snapshots[]defsnapshot(self,label):记录显存快照allocatedtorch.npu.memory_allocated()reservedtorch.npu.memory_reserved()max_allocatedtorch.npu.max_memory_allocated()snapshot{label:label,allocated:allocated/1024**3,# GBreserved:reserved/1024**3,max_allocated:max_allocated/1024**3,timestamp:time.time()}self.snapshots.append(snapshot)returnsnapshotdefprint_status(self):打印显存状态snapshotself.snapshot()print(f已分配:{snapshot[allocated]:.2f}GB)print(f已预留:{snapshot[reserved]:.2f}GB)print(f峰值:{snapshot[max_allocated]:.2f}GB)defreport(self):生成显存报告ifnotself.snapshots:print(无快照记录)returnprint( 显存报告:)print(f{标签:20}{已分配(GB):12}{已预留(GB):12}{峰值(GB):12})print(-*56)forsinself.snapshots:print(f{s[label]:20}{s[allocated]:12.2f}{s[reserved]:12.2f}{s[max_allocated]:12.2f})# 使用示例monitorNpuMemoryMonitor()monitor.snapshot(初始化)modelload_large_model()monitor.snapshot(模型加载)input_datatorch.randn(1,3,224,224).npu()monitor.snapshot(输入数据)outputmodel(input_data)monitor.snapshot(推理完成)monitor.report()2.3 显存泄漏检测classMemoryLeakDetector:def__init__(self):self.snapshots[]self.leak_threshold0.1# 10% 增长阈值defcheck_leak(self,num_iterations100):检查显存泄漏# 记录初始状态initial_allocatedtorch.npu.memory_allocated()self.snapshots.append((start,initial_allocated))foriinrange(num_iterations):# 执行推理input_datatorch.randn(1,3,224,224).npu()outputmodel(input_data)delinput_data,output# 每 10 次记录一次ifi%100:allocatedtorch.npu.memory_allocated()self.snapshots.append((fiter_{i},allocated))# 检查泄漏final_allocatedtorch.npu.memory_allocated()growth(final_allocated-initial_allocated)/initial_allocatedprint(f初始显存:{initial_allocated/1024**3:.2f}GB)print(f最终显存:{final_allocated/1024**3:.2f}GB)print(f增长:{growth:.2%})ifgrowthself.leak_threshold:print(⚠️ 可能存在显存泄漏!)returnTrueelse:print(✅ 显存使用正常)returnFalse# 使用示例detectorMemoryLeakDetector()detector.check_leak(100)三、显存优化技巧3.1 及时释放显存# 错误: 未及时释放defbad_inference(model,input_data):output1model(input_data)output2model(output1)# output1 仍在显存中returnoutput2# 正确: 及时释放defgood_inference(model,input_data):withtorch.no_grad():output1model(input_data)delinput_data# 释放输入output2model(output1)deloutput1# 释放中间结果returnoutput2# 使用 torch.npu.amp 自动混合精度defamp_inference(model,input_data):withtorch.npu.amp.autocast():outputmodel(input_data)returnoutput3.2 显存池化classMemoryPool:def__init__(self,max_size1024*1024*1024):# 1GBself.max_sizemax_size self.free_blocks{}# size - list of blocksself.used_blocks{}self.lockthreading.Lock()defallocate(self,size):分配显存withself.lock:# 对齐到 256 字节aligned_size(size255)//256*256# 查找空闲块ifaligned_sizeinself.free_blocksandself.free_blocks[aligned_size]:blockself.free_blocks[aligned_size].pop()self.used_blocks[id(block)](block,aligned_size)returnblock# 分配新块blocktorch.empty(aligned_size,dtypetorch.uint8,devicenpu)self.used_blocks[id(block)](block,aligned_size)returnblockdeffree(self,block):释放显存withself.lock:block_idid(block)ifblock_idinself.used_blocks:_,sizeself.used_blocks.pop(block_id)ifsizenotinself.free_blocks:self.free_blocks[size][]self.free_blocks[size].append(block)defget_stats(self):获取统计信息withself.lock:used_totalsum(sizefor_,sizeinself.used_blocks.values())free_totalsum(len(blocks)*sizeforsize,blocksinself.free_blocks.items())return{used:used_total/1024**3,free:free_total/1024**3,total:(used_totalfree_total)/1024**3,utilization:used_total/(used_totalfree_total)*100}# 使用示例poolMemoryPool()# 分配block1pool.allocate(1024*1024)# 1MBblock2pool.allocate(2048*1024)# 2MB# 释放pool.free(block1)# 查看统计print(pool.get_stats())四、大模型显存优化4.1 模型分片classModelSharding:def__init__(self,model,num_shards2):self.modelmodel self.num_shardsnum_shards self.shardsself._split_model()def_split_model(self):将模型分片parameterslist(self.model.named_parameters())shard_sizelen(parameters)//self.num_shards shards[]foriinrange(self.num_shards):starti*shard_size endstartshard_sizeifiself.num_shards-1elselen(parameters)shard_paramsparameters[start:end]shards.append(shard_params)returnshardsdefinference(self,input_data):分片推理current_inputinput_dataforshardinself.shards:# 加载当前分片到显存shard_modelself._load_shard(shard)# 推理withtorch.no_grad():current_inputshard_model(current_input)# 释放分片delshard_model torch.npu.empty_cache()returncurrent_input# 使用示例sharded_modelModelSharding(large_model,num_shards4)outputsharded_model.inference(input_data)4.2 梯度检查点classGradientCheckpointing(nn.Module):def__init__(self,model):super().__init__()self.modelmodeldefforward(self,x):梯度检查点# 只保存中间激活值不保存梯度withtorch.no_grad():forlayerinself.model.layers:xlayer(x)returnx# 使用示例modelGradientCheckpointing(large_model)outputmodel(input_data)# 训练时使用deftrain_with_checkpointing(model,data,labels):model.train()# 前向传播 (不保存梯度)outputmodel(data)lossF.cross_entropy(output,labels)# 反向传播 (重新计算中间值)loss.backward()returnloss4.3 显存碎片整理defdefragment_memory():显存碎片整理# 获取当前状态allocated_beforetorch.npu.memory_allocated()reserved_beforetorch.npu.memory_reserved()# 清空缓存torch.npu.empty_cache()# 获取整理后状态allocated_aftertorch.npu.memory_allocated()reserved_aftertorch.npu.memory_reserved()print(f整理前: 已分配{allocated_before/1024**3:.2f}GB, 已预留{reserved_before/1024**3:.2f}GB)print(f整理后: 已分配{allocated_after/1024**3:.2f}GB, 已预留{reserved_after/1024**3:.2f}GB)print(f释放:{(reserved_before-reserved_after)/1024**3:.2f}GB)# 定期执行defragment_memory()五、常见问题问题原因解决方案OOM 错误显存不足减小 batch size、使用模型分片显存泄漏未释放引用使用del、torch.npu.empty_cache()推理速度慢显存带宽不足使用 FP16、量化碎片化严重频繁分配释放使用显存池、定期碎片整理大模型加载失败显存不足模型分片、梯度检查点相关仓库torch_npu- 显存管理 API https://atomgit.com/cann/ops-nn
