1
3
DCU技术达人
新人贴
海光 DCU 大模型推理环境全流程脚本集
2026-09-30 12:14:02
36
# 海光DCU大模型推理环境全流程脚本集|纯代码可直接复用 整理一套从环境校验、依赖安装、服务部署到压测验证的全链路可运行脚本,全程无额外依赖,复制到服务器按顺序执行即可完成整套推理环境搭建,适合快速部署生产级推理服务。 ## 1. DTK环境一键校验脚本 执行前先确认DTK环境是否正常、硬件是否可识别,避免后续安装全做无用功。 保存为 `check_dtk_env.sh`: ```bash #!/bin/bash echo "=== DCU 环境基线校验 ===" # 1. 检查rocm-smi工具 echo -n "1. rocm-smi 工具: " if command -v rocm-smi &>/dev/null; then echo "正常" else echo "未找到,请先加载DTK环境变量" exit 1 fi # 2. 检查HIP编译器 echo -n "2. hipcc 编译器: " if command -v hipcc &>/dev/null; then hipcc --version 2>/dev/null | grep "HIP version" else echo "未找到" exit 1 fi # 3. 设备识别检查 echo -n "3. DCU 设备识别: " card_count=$(rocm-smi --showhw 2>/dev/null | grep -c "Device ID") if [ "$card_count" -gt 0 ]; then echo "检测到 $card_count 张卡" else echo "未识别到设备,请检查驱动与权限" exit 1 fi # 4. 输出核心硬件信息 echo -e "\n=== 硬件明细 ===" rocm-smi --showhw 2>/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`: ```bash #!/bin/bash echo "=== 安装 DCU 推理依赖 ===" # 自动识别DTK对应ROCm版本 DTK_VER=$(hipcc --version 2>/dev/null | grep -oP "HIP version: \K[0-9]+" | head -1) if [ "$DTK_VER" -ge 26 ]; then ROCM_VER="6.0" TORCH_VER="2.4.0" elif [ "$DTK_VER" -ge 25 ]; then ROCM_VER="5.7" TORCH_VER="2.3.1" else echo "DTK版本过低,建议升级至25.04以上" exit 1 fi echo "匹配 ROCm 版本: $ROCM_VER,PyTorch 版本: $TORCH_VER" # 安装PyTorch pip install torch==$TORCH_VER torchvision==0.19.0 \ --index-url https://download.pytorch.org/whl/rocm$ROCM_VER # 安装推理常用依赖 pip install transformers accelerate sentencepiece \ modelscope vllm==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()}") EOF echo "✅ 依赖安装完成" ``` 使用方式:`bash install_deps.sh`,自动适配DTK版本,无需手动对应。 ## 3. vLLM推理服务一键启停脚本 带进程守护、日志管理的服务化脚本,支持启动、停止、状态查询,无需手动管理进程。 保存为 `vllm_service.sh`: ```bash #!/bin/bash # 配置项 - 根据实际环境修改 MODEL_PATH="./models/Qwen2-7B-Instruct" MODEL_NAME="Qwen2-7B-Instruct" PORT=8000 TP_SIZE=1 DTK_ENV="/opt/dtk-26.04/env.sh" LOG_FILE="./vllm_service.log" PID_FILE="./vllm.pid" # 加载DTK环境 source $DTK_ENV export TRITON_HIP_LLD_PATH=/opt/dtk-26.04/llvm/bin/ld.lld start_service() { if [ -f "$PID_FILE" ] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then echo "服务已在运行,PID: $(cat "$PID_FILE")" exit 1 fi echo "启动 vLLM 推理服务..." nohup numactl --cpunodebind=0 --membind=0 \ vllm serve "$MODEL_PATH" \ --served-model-name "$MODEL_NAME" \ --tensor-parallel-size "$TP_SIZE" \ --dtype bfloat16 \ --max-model-len 8192 \ --gpu-memory-utilization 0.9 \ --enforce-eager \ --max-num-seqs 128 \ --host 0.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" ]; then echo "服务未运行" exit 1 fi pid=$(cat "$PID_FILE") if kill -0 "$pid" 2>/dev/null; then echo "停止服务,PID: $pid" kill "$pid" sleep 3 kill -9 "$pid" 2>/dev/null fi rm -f "$PID_FILE" echo "服务已停止" } status_service() { if [ -f "$PID_FILE" ] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then echo "✅ 服务运行中,PID: $(cat "$PID_FILE")" echo "端口: $PORT" else echo "❌ 服务未运行" fi } case "$1" in start) start_service ;; stop) stop_service ;; restart) stop_service; sleep 2; start_service ;; status) status_service ;; *) echo "用法: $0 start|stop|restart|status"; exit 1 ;; esac ``` 使用方式: ```bash bash vllm_service.sh start # 启动服务 bash vllm_service.sh status # 查看状态 bash vllm_service.sh stop # 停止服务 ``` ## 4. 推理接口压测脚本 服务启动后,批量并发压测,输出吞吐、平均延迟等核心指标,量化验证服务性能。 保存为 `benchmark.py`: ```python import time import requests import threading from statistics import mean # 配置 API_URL = "http://localhost:8000/v1/chat/completions" MODEL_NAME = "Qwen2-7B-Instruct" CONCURRENT = 16 # 并发数 TOTAL_REQUESTS = 64 MAX_TOKENS = 256 PROMPT = "请详细介绍海光DCU的主要应用场景和技术特点" # 全局统计 success = 0 total_tokens = 0 latencies = [] lock = threading.Lock() def worker(): global success, 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() - start if resp.status_code == 200: data = resp.json() tokens = data["usage"]["completion_tokens"] with lock: success += 1 total_tokens += tokens latencies.append(cost) except Exception as e: print(f"请求失败: {e}") if __name__ == "__main__": print(f"开始压测: 并发={CONCURRENT}, 总请求={TOTAL_REQUESTS}") start_time = time.perf_counter() threads = [] for _ in range(TOTAL_REQUESTS): t = threading.Thread(target=worker) threads.append(t) t.start() # 控制并发数 while threading.active_count() > CONCURRENT + 1: time.sleep(0.01) for t in threads: t.join() total_time = time.perf_counter() - start_time print("\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`: ```bash #!/bin/bash echo "=== DCU 资源清理工具 ===" # 列出所有占用DCU的进程 echo "当前占用DCU设备的进程:" echo "PID 用户 进程名" echo "-------------------------" pids="" for dev in /dev/dri/renderD*; do dev_pids=$(fuser "$dev" 2>/dev/null) for pid in $dev_pids; do if ! echo "$pids" | grep -q "$pid"; then pids="$pids $pid" user=$(ps -p "$pid" -o user= 2>/dev/null) comm=$(ps -p "$pid" -o comm= 2>/dev/null) printf "%-6s %-8s %s\n" "$pid" "$user" "$comm" fi done done if [ -z "$pids" ]; then echo "无占用进程,显存空闲" exit 0 fi read -p "是否终止以上所有进程并释放显存? (y/n): " confirm if [ "$confirm" = "y" ]; then for pid in $pids; do kill -9 "$pid" 2>/dev/null && echo "已终止进程 $pid" done # 清理PyTorch缓存 python3 -c "import torch; torch.cuda.empty_cache()" 2>/dev/null echo "✅ 清理完成" else echo "已取消操作" fi ``` 使用方式:`bash clean_dcu.sh`,先列进程再确认,避免误杀正常任务。 --- 以上5个脚本覆盖了DCU大模型推理环境从部署到运维的全流程高频场景,全部基于原生命令与官方工具实现,无额外依赖,存到服务器固定目录即可长期复用。
全部评论(0)
只看楼主
暂无数据
--
楼主
Level.3
0 帖子
楼主热帖
暂无数据