一个简单测试Deepseek吞吐量的脚本,国内环境可跑

一个简单测试Deepseek吞吐量的脚本,这里用DeepSeek-R1-Distill-Qwen-32B ,支持单卡4090 24G可跑,具体看你的硬件情况做调整,理论支持所有的模型,看你需要,可以修改模型名称,重点是pip使用国内的源,模型下载用阿里的ModelScope,无障碍下载,使用.
最后可以生成一个txt与html报表.
前提是你安装了python与python-venv,你可以不用venv来做,可动手改一下脚本.
 

#!/bin/bash
set -e
# 参数配置
MODEL_REPO="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"
BATCH_SIZES=(1 2 4 10 20 30 50)  # 32B模型显存需求大，batch_size调小
SEQ_LENGTHS=(256 512)
WARMUP_STEPS=3
MEASURE_STEPS=10
VENV_NAME="deepseek_test"
REPORT_FILE="benchmark_report.txt"
HTML_REPORT_FILE="benchmark_report.html"

# 创建虚拟环境
echo "创建Python虚拟环境..."
python3 -m venv $VENV_NAME
source $VENV_NAME/bin/activate

# 配置国内pip源
echo "配置国内pip源..."
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
pip config set global.trusted-host pypi.tuna.tsinghua.edu.cn

# 安装依赖
echo "安装依赖包..."
pip install torch transformers modelscope accelerate

# 测试函数
run_benchmark() {
    local batch_size=$1
    local seq_length=$2
    
    echo -e "\n测试配置: batch_size=${batch_size}, seq_length=${seq_length}"
    
    python3 - < $REPORT_FILE
echo "GPU信息:" >> $REPORT_FILE
nvidia-smi --query-gpu=name,memory.total --format=csv >> $REPORT_FILE

echo "
DeepSeek-R1-Distill-Qwen-32B性能测试报告
" > $HTML_REPORT_FILE
echo "
GPU信息：
" >> $HTML_REPORT_FILE
nvidia-smi --query-gpu=name,memory.total --format=csv | sed 's/^/
/' | sed 's/$/<\/p>/' >> $HTML_REPORT_FILE

for bs in ${BATCH_SIZES[@]}; do
    for seq in ${SEQ_LENGTHS[@]}; do
        echo -e "\n测试配置: batch_size=${bs}, seq_length=${seq}" >> $REPORT_FILE
        echo "
测试配置: batch_size=${bs}, seq_length=${seq}
" >> $HTML_REPORT_FILE
        run_benchmark $bs $seq | tee -a $REPORT_FILE | sed 's/^/
/' | sed 's/$/<\/p>/' >> $HTML_REPORT_FILE
    done
done

echo "" >> $HTML_REPORT_FILE

deactivate
echo "测试完成"