#!/bin/bash
# 快速测试单个checkpoint - 验证测试流程是否正确

export CUDA_DEVICE_MAX_CONNECTIONS=1

# 设置NPU驱动库路径
export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/stub:$LD_LIBRARY_PATH

# 设置Python路径
export PYTHONPATH=/mnt/models/CODE/MindSpeed-LLM-v2.3.0:$PYTHONPATH

# TOKENIZER_PATH="/mnt/data1/models/Qwen3-0.6B"
TOKENIZER_PATH="/mnt/models/wh/data/minimind-master/model"
# CHECKPOINT应该指向模型目录，而不是iter_*子目录
# 系统会根据latest_checkpointed_iteration.txt自动加载最新的
CHECKPOINT="${1:-/mnt/models/wh/models/qwen_200m_0615}"
TASK="${2:-batch_test}"

echo "=========================================="
echo "快速测试脚本"
echo "=========================================="
echo "模型目录: ${CHECKPOINT}"
echo "测试任务: ${TASK}"
echo "=========================================="

# 显示最新的checkpoint
if [ -f "${CHECKPOINT}/latest_checkpointed_iteration.txt" ]; then
    ITER=$(cat ${CHECKPOINT}/latest_checkpointed_iteration.txt)
    echo "最新checkpoint: iter_${ITER}"
fi

MASTER_ADDR=localhost
MASTER_PORT=6000
NPUS_PER_NODE=1
SEQ_LENGTH=4096

DISTRIBUTED_ARGS="
    --nproc_per_node $NPUS_PER_NODE \
    --nnodes 1 \
    --node_rank 0 \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

LOG_FILE="/mnt/models/CODE/MindSpeed-LLM-v2.3.0/logs/test_$(basename ${CHECKPOINT})_${TASK}.log"

cd /mnt/models/wh/code/MindSpeed-LLM-v2.3.0/examples/mcore/qwen3/

torchrun $DISTRIBUTED_ARGS inference_batch_test.py \
    --use-mcore-models \
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 1 \
    --expert-model-parallel-size 1 \
    --load ${CHECKPOINT} \
    --spec mindspeed_llm.tasks.models.spec.qwen3_spec layer_spec \
    --kv-channels 64 \
    --qk-layernorm \
    --num-layers 28 \
    --hidden-size 768 \
    --use-rotary-position-embeddings \
    --num-attention-heads 12 \
    --ffn-hidden-size 1536 \
    --max-position-embeddings 4096 \
    --seq-length ${SEQ_LENGTH} \
    --make-vocab-size-divisible-by 1 \
    --padded-vocab-size 6400 \
    --rotary-base 1000000 \
    --micro-batch-size 1 \
    --disable-bias-linear \
    --swiglu \
    --tokenizer-type PretrainedFromHF \
    --tokenizer-name-or-path ${TOKENIZER_PATH} \
    --normalization RMSNorm \
    --position-embedding-type rope \
    --norm-epsilon 1e-6 \
    --hidden-dropout 0 \
    --attention-dropout 0 \
    --max-new-tokens 256 \
    --no-gradient-accumulation-fusion \
    --attention-softmax-in-fp32 \
    --no-masked-softmax-fusion \
    --group-query-attention \
    --num-query-groups 4 \
    --seed 42 \
    --bf16 \
    --task ${TASK} \
    2>&1 | tee ${LOG_FILE}

echo ""
echo "日志保存至: ${LOG_FILE}"