# 海光DCU显存排查入门：3个零成本优化技巧｜附可运行代码 刚迁移DCU的开发者常会遇到两个典型问题：同参数模型的显存占用和CUDA环境有差异，或者跑任务时突发OOM却找不到具体原因。其实DCU显存管理逻辑与CUDA同源，排查思路基本通用，只是几个专属细节容易被忽略。本文从实操角度讲解快速定位显存问题的方法，以及三个零成本即可落地的优化技巧，所有代码均可直接运行。 ## 一、两步快速定位显存占用 ### 1. 命令行全局查看 对应NVIDIA环境的 `nvidia-smi`，DCU环境使用 `rocm-smi`即可查看整机显存状态，包括单卡已用/总显存、使用率、温度等核心信息： ```bash # 静态查看显存占用与算力使用率 rocm-smi --showmemuse --showuse # 实时刷新监控（每秒更新一次） watch -n 1 rocm-smi --showmemuse --showtemp ``` 注意：DCU采用HBM高带宽显存，驱动与运行时会预留少量空间做系统调度，因此可用显存略低于标称值属于正常现象。 ### 2. Python代码级精准统计 如果需要细分模型权重、中间张量、推理缓存的显存占比，直接使用PyTorch原生接口即可。DCU适配的PyTorch通过兼容层完整保留了 `torch.cuda`系列API，代码无需额外修改： ```python import torch from transformers import AutoModelForCausalLM, AutoTokenizer def print_mem_info(tag: str = ""): allocated = torch.cuda.memory_allocated() / 1024**3 reserved = torch.cuda.memory_reserved() / 1024**3 max_alloc = torch.cuda.max_memory_allocated() / 1024**3 print(f"[{tag}] 已分配: {allocated:.2f} GB | 预留缓存: {reserved:.2f} GB | 峰值: {max_alloc:.2f} GB") model_path = "Qwen/Qwen2-7B-Instruct" tokenizer = AutoTokenizer.from_pretrained(model_path) # 加载模型前基准 print_mem_info("加载模型前") model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.float16, device_map="auto" ) print_mem_info("加载模型后") # 首次推理后统计含KV缓存的总占用 inputs = tokenizer("测试输入文本", return_tensors="pt").to("cuda") _ = model.generate(**inputs, max_new_tokens=100) print_mem_info("首次推理后") ``` 运行后可清晰区分模型权重、运行时缓存各自的占比，快速定位OOM的触发点。 ## 二、3个零成本显存优化技巧 均无需修改模型结构，一行配置即可生效，适合入门阶段快速解决显存紧张问题。 ### 技巧1：推理间隙显式释放缓存 DCU的显存缓存回收机制与CUDA存在细微差异，长时运行的推理服务容易出现闲置缓存堆积。任务结束后加一行清空命令，即可快速释放冗余显存： ```python # 批量推理间隙、服务空闲时调用 torch.cuda.empty_cache() ``` 该操作仅释放PyTorch预留的空闲缓存，不会影响正在计算的张量，适合定时巡检或批量任务间隔调用。 ### 技巧2：半精度/BF16加载模型 这是性价比最高的优化，显存占用直接减半。DCU K100及以上型号原生支持FP16与BF16精度，加载模型时指定数据类型即可： ```python # 推理场景优先FP16，显存减半且精度损失可忽略 model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.float16, device_map="auto" ) # 微调场景推荐BF16，数值稳定性更好，避免梯度溢出 # torch_dtype=torch.bfloat16 ``` 以7B参数模型为例，FP32加载约占28GB显存，切换FP16后仅需约14GB，单张32GB显存的DCU即可轻松承载。 ### 技巧3：微调场景开启梯度检查点 小样本微调时若显存不足以承载优化器状态与中间激活值，开启梯度检查点可通过“计算换显存”的方式，降低30%-40%的峰值显存占用，训练速度损失通常不到10%： ```python # transformers 原生支持，一行配置生效 model.gradient_checkpointing_enable() ``` 原理是反向传播时重新计算中间激活值，而非全程保存在显存中，适合显存紧张的单机微调场景。 ## 三、新手避坑提醒 1. 不要打满显存：DCU运行时需要预留5%-10%的显存用于调度与多卡通信，推理框架如vLLM的 `gpu-memory-utilization`建议设为0.9左右，避免高并发触发OOM。 2. 多卡场景优先做NUMA绑核：跨NUMA组通信会额外占用显存缓存，完成绑核优化后既能提升性能，也能减少这部分冗余开销。 3. 优先使用官方适配权重：第三方自行转换的权重可能存在显存排布不合理的问题，容易出现隐性显存浪费。 整体来看，DCU的显存使用体验与CUDA高度对齐，入门阶段掌握基础排查方法和这三个轻量化技巧，就能解决80%的日常显存问题。更深层的优化可结合算子融合、量化等手段逐步推进，适合生产环境再深入落地。