#!/bin/bash
#SBATCH -A nvr_elm_llm
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 00:30:00
#SBATCH -N 1
#SBATCH -J train_debug_gpu
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_debug_gpu_error.log
#SBATCH -o workdir/train_debug_gpu_out.log

set -e

echo "=========================================="
echo "DFM_Qwen2VL Training Debug (GPU, 1 iter)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "GPUs: 1"
echo "CPUs: 16"
echo "Memory: 1024GB"
echo ""

# Activate environment
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

echo "Python: $(python3 --version)"
echo "PyTorch: $(python3 -c 'import torch; print(torch.__version__)')"
echo "CUDA available: $(python3 - <<'PY' 2>/dev/null || true
import torch
print(torch.cuda.is_available())
PY
)"
echo "nvidia-smi:"
nvidia-smi || true
echo ""

# Set paths
export PYTHONPATH="$PWD:$PWD/chengyue_script:$PWD/src:$PYTHONPATH"
export OMP_NUM_THREADS=16
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:64,garbage_collection_threshold:0.9
export TORCHDYNAMO_DISABLE=1
export TORCH_COMPILE=0

# 配置
exp_id=train_debug_gpu_1iter_${SLURM_JOB_ID:-$(date +%Y%m%d_%H%M%S)}
model_name_or_path=/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_DFM
tokenizer_name=Qwen/Qwen2.5-VL-3B-Instruct
dataset_path=/home/chengyuew/chengyuew/workspace/LMFlow/data/nemotron_no_thinking_mini_1k
output_dir=output_models/${exp_id}
log_dir=log/${exp_id}

echo "=========================================="
echo "配置:"
echo "  模型: ${model_name_or_path}"
echo "  Tokenizer: ${tokenizer_name}"
echo "  数据: ${dataset_path}"
echo "  输出: ${output_dir}"
echo "  最大步数: 1 (仅1个iteration)"
echo "  序列长度: 2048"
echo "  per_device_train_batch_size: 8"
echo "  精度: bf16"
echo "  bd_size: 128"
echo "  learn_padding: 1"
echo "  conplemenrary_mask: 1"
echo "=========================================="
echo ""

# 创建目录
mkdir -p ${output_dir} ${log_dir}

# 检查路径
echo "检查路径..."
if [ ! -d "${model_name_or_path}" ]; then
    echo "❌ 模型路径不存在: ${model_name_or_path}"
    exit 1
fi
echo "✅ 模型路径: ${model_name_or_path}"

if [ ! -d "${dataset_path}" ]; then
    echo "❌ 数据路径不存在: ${dataset_path}"
    exit 1
fi
echo "✅ 数据路径: ${dataset_path}"
echo ""

echo "开始训练 debug (1 iteration)..."
echo "=========================================="
echo ""

# 运行训练 - 只1个iteration（多GPU）
deepspeed --num_gpus 8 examples/finetune.py \
    --model_name_or_path ${model_name_or_path} \
    --tokenizer_name ${tokenizer_name} \
    --conversation_template qwen2_5_no_reasoning \
    --trust_remote_code 1 \
    --dataset_path ${dataset_path} \
    --output_dir ${output_dir} \
    --max_steps 1 \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio 0.03 \
    --learning_rate 2e-5 \
    --disable_group_texts 1 \
    --block_size 2048 \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 1 \
    --gradient_checkpointing \
    --run_name ${exp_id} \
    --overwrite_output_dir \
    --logging_steps 20 \
    --do_train \
    --validation_split_percentage 0 \
    --ddp_timeout 72000 \
    --save_steps 100 \
    --save_total_limit 30 \
    --dataloader_num_workers 8 \
    --preprocessing_num_workers 32 \
    --bd_size 128 \
    --mdm 1 \
    --learn_padding 1 \
    --conplemenrary_mask 1 \
    --bf16 \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    2>&1 | tee ${log_dir}/train_debug.log

exit_code=${PIPESTATUS[0]}

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ 训练 Debug 成功!"
else
    echo "❌ 训练 Debug 失败 (exit code: ${exit_code})"
    echo "请查看日志:\n  输出: ${log_dir}/train_debug.log\n  错误: workdir/train_debug_gpu_error.log"
fi
echo "=========================================="

exit $exit_code


