海光 DCU 大模型推理环境全流程脚本集
2026/8/4 20:54:34 网站建设 项目流程

海光DCU大模型推理环境全流程脚本集|纯代码可直接复用

整理一套从环境校验、依赖安装、服务部署到压测验证的全链路可运行脚本,全程无额外依赖,复制到服务器按顺序执行即可完成整套推理环境搭建,适合快速部署生产级推理服务。

1. DTK环境一键校验脚本

执行前先确认DTK环境是否正常、硬件是否可识别,避免后续安装全做无用功。
保存为check_dtk_env.sh

#!/bin/bashecho"=== DCU 环境基线校验 ==="# 1. 检查rocm-smi工具echo-n"1. rocm-smi 工具: "ifcommand-vrocm-smi&>/dev/null;thenecho"正常"elseecho"未找到,请先加载DTK环境变量"exit1fi# 2. 检查HIP编译器echo-n"2. hipcc 编译器: "ifcommand-vhipcc&>/dev/null;thenhipcc--version2>/dev/null|grep"HIP version"elseecho"未找到"exit1fi# 3. 设备识别检查echo-n"3. DCU 设备识别: "card_count=$(rocm-smi--showhw2>/dev/null|grep-c"Device ID")if["$card_count"-gt0];thenecho"检测到$card_count张卡"elseecho"未识别到设备,请检查驱动与权限"exit1fi# 4. 输出核心硬件信息echo-e"\n=== 硬件明细 ==="rocm-smi--showhw2>/dev/null|grep-E"Device ID|VRAM Total|Compute Unit"echo-e"\n✅ 环境校验通过"

使用方式:bash check_dtk_env.sh,任意一项失败都会直接提示对应问题。

2. PyTorch与推理依赖一键安装

自动匹配DTK版本安装对应ROCm版PyTorch,同时安装vLLM、transformers等推理必备依赖。
保存为install_deps.sh

#!/bin/bashecho"=== 安装 DCU 推理依赖 ==="# 自动识别DTK对应ROCm版本DTK_VER=$(hipcc--version2>/dev/null|grep-oP"HIP version: \K[0-9]+"|head-1)if["$DTK_VER"-ge26];thenROCM_VER="6.0"TORCH_VER="2.4.0"elif["$DTK_VER"-ge25];thenROCM_VER="5.7"TORCH_VER="2.3.1"elseecho"DTK版本过低,建议升级至25.04以上"exit1fiecho"匹配 ROCm 版本:$ROCM_VER,PyTorch 版本:$TORCH_VER"# 安装PyTorchpipinstalltorch==$TORCH_VERtorchvision==0.19.0\--index-url https://download.pytorch.org/whl/rocm$ROCM_VER# 安装推理常用依赖pipinstalltransformers accelerate sentencepiece\modelscopevllm==0.18.1 requests# 验证安装结果echo-e"\n=== 安装结果验证 ==="python3<<'EOF' import torch import transformers import vllm print(f"PyTorch: {torch.__version__}") print(f"Transformers: {transformers.__version__}") print(f"vLLM: {vllm.__version__}") print(f"DCU可用: {torch.cuda.is_available()}") print(f"DCU数量: {torch.cuda.device_count()}") EOFecho"✅ 依赖安装完成"

使用方式:bash install_deps.sh,自动适配DTK版本,无需手动对应。

3. vLLM推理服务一键启停脚本

带进程守护、日志管理的服务化脚本,支持启动、停止、状态查询,无需手动管理进程。
保存为vllm_service.sh

#!/bin/bash# 配置项 - 根据实际环境修改MODEL_PATH="./models/Qwen2-7B-Instruct"MODEL_NAME="Qwen2-7B-Instruct"PORT=8000TP_SIZE=1DTK_ENV="/opt/dtk-26.04/env.sh"LOG_FILE="./vllm_service.log"PID_FILE="./vllm.pid"# 加载DTK环境source$DTK_ENVexportTRITON_HIP_LLD_PATH=/opt/dtk-26.04/llvm/bin/ld.lldstart_service(){if[-f"$PID_FILE"]&&kill-0"$(cat"$PID_FILE")"2>/dev/null;thenecho"服务已在运行,PID:$(cat"$PID_FILE")"exit1fiecho"启动 vLLM 推理服务..."nohupnumactl--cpunodebind=0--membind=0\vllm serve"$MODEL_PATH"\--served-model-name"$MODEL_NAME"\--tensor-parallel-size"$TP_SIZE"\--dtypebfloat16\--max-model-len8192\--gpu-memory-utilization0.9\--enforce-eager\--max-num-seqs128\--host0.0.0.0\--port"$PORT"\>>"$LOG_FILE"2>&1&echo$!>"$PID_FILE"echo"服务已启动,PID:$!"echo"日志文件:$LOG_FILE"echo"端口:$PORT"}stop_service(){if[!-f"$PID_FILE"];thenecho"服务未运行"exit1fipid=$(cat"$PID_FILE")ifkill-0"$pid"2>/dev/null;thenecho"停止服务,PID:$pid"kill"$pid"sleep3kill-9"$pid"2>/dev/nullfirm-f"$PID_FILE"echo"服务已停止"}status_service(){if[-f"$PID_FILE"]&&kill-0"$(cat"$PID_FILE")"2>/dev/null;thenecho"✅ 服务运行中,PID:$(cat"$PID_FILE")"echo"端口:$PORT"elseecho"❌ 服务未运行"fi}case"$1"instart)start_service;;stop)stop_service;;restart)stop_service;sleep2;start_service;;status)status_service;;*)echo"用法:$0start|stop|restart|status";exit1;;esac

使用方式:

bashvllm_service.sh start# 启动服务bashvllm_service.sh status# 查看状态bashvllm_service.sh stop# 停止服务

4. 推理接口压测脚本

服务启动后,批量并发压测,输出吞吐、平均延迟等核心指标,量化验证服务性能。
保存为benchmark.py

importtimeimportrequestsimportthreadingfromstatisticsimportmean# 配置API_URL="http://localhost:8000/v1/chat/completions"MODEL_NAME="Qwen2-7B-Instruct"CONCURRENT=16# 并发数TOTAL_REQUESTS=64MAX_TOKENS=256PROMPT="请详细介绍海光DCU的主要应用场景和技术特点"# 全局统计success=0total_tokens=0latencies=[]lock=threading.Lock()defworker():globalsuccess,total_tokens payload={"model":MODEL_NAME,"messages":[{"role":"user","content":PROMPT}],"max_tokens":MAX_TOKENS,"temperature":0.1,"stream":False}start=time.perf_counter()try:resp=requests.post(API_URL,json=payload,timeout=60)cost=time.perf_counter()-startifresp.status_code==200:data=resp.json()tokens=data["usage"]["completion_tokens"]withlock:success+=1total_tokens+=tokens latencies.append(cost)exceptExceptionase:print(f"请求失败:{e}")if__name__=="__main__":print(f"开始压测: 并发={CONCURRENT}, 总请求={TOTAL_REQUESTS}")start_time=time.perf_counter()threads=[]for_inrange(TOTAL_REQUESTS):t=threading.Thread(target=worker)threads.append(t)t.start()# 控制并发数whilethreading.active_count()>CONCURRENT+1:time.sleep(0.01)fortinthreads:t.join()total_time=time.perf_counter()-start_timeprint("\n=== 压测结果 ===")print(f"成功请求:{success}/{TOTAL_REQUESTS}")print(f"总生成Token:{total_tokens}")print(f"总耗时:{total_time:.2f}s")print(f"平均吞吐:{total_tokens/total_time:.2f}token/s")print(f"单请求平均延迟:{mean(latencies):.2f}s")print(f"QPS:{success/total_time:.2f}req/s")

使用方式:python3 benchmark.py,自动输出完整性能指标。

5. 显存与进程一键清理脚本

服务异常退出、显存残留时执行,快速释放所有DCU资源,无需重启节点。
保存为clean_dcu.sh

#!/bin/bashecho"=== DCU 资源清理工具 ==="# 列出所有占用DCU的进程echo"当前占用DCU设备的进程:"echo"PID 用户 进程名"echo"-------------------------"pids=""fordevin/dev/dri/renderD*;dodev_pids=$(fuser"$dev"2>/dev/null)forpidin$dev_pids;doif!echo"$pids"|grep-q"$pid";thenpids="$pids$pid"user=$(ps-p"$pid"-ouser=2>/dev/null)comm=$(ps-p"$pid"-ocomm=2>/dev/null)printf"%-6s %-8s %s\n""$pid""$user""$comm"fidonedoneif[-z"$pids"];thenecho"无占用进程,显存空闲"exit0firead-p"是否终止以上所有进程并释放显存? (y/n): "confirmif["$confirm"="y"];thenforpidin$pids;dokill-9"$pid"2>/dev/null&&echo"已终止进程$pid"done# 清理PyTorch缓存python3-c"import torch; torch.cuda.empty_cache()"2>/dev/nullecho"✅ 清理完成"elseecho"已取消操作"fi

使用方式:bash clean_dcu.sh,先列进程再确认,避免误杀正常任务。


以上5个脚本覆盖了DCU大模型推理环境从部署到运维的全流程高频场景,全部基于原生命令与官方工具实现,无额外依赖,存到服务器固定目录即可长期复用。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询