#!/bin/bash
#SBATCH -A nvr_lpr_llm
#SBATCH -p batch_block1
#SBATCH -t 04:00:00
#SBATCH -N 8
#SBATCH -J fastdllmv2_mmmathinstruct150k_8node_bd32_block_causal_no_dynamic_lr_1e-5_ds
#SBATCH --array=1-10%1
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/fastdllmv2_mmmathinstruct150k_8node_bd32_block_causal_no_dynamic_lr_1e-5_ds_error.log
#SBATCH -o workdir/fastdllmv2_mmmathinstruct150k_8node_bd32_block_causal_no_dynamic_lr_1e-5_ds_out.log
#SBATCH --exclusive

idx=$SLURM_ARRAY_TASK_ID
total=$SLURM_ARRAY_TASK_COUNT

# Prefer SLURM submit dir to avoid slurm spool path issues.
PROJECT_ROOT="${SLURM_SUBMIT_DIR:-$PWD}"
cd "${PROJECT_ROOT}" || { echo "ERROR: cannot cd to PROJECT_ROOT=${PROJECT_ROOT}" >&2; exit 1; }
if [ ! -f "scripts/cluster_env.sh" ]; then
    echo "ERROR: scripts/cluster_env.sh not found in PROJECT_ROOT=${PROJECT_ROOT}" >&2
    exit 1
fi

# Unified cluster env + conda activation (for deepspeed command availability).
source "scripts/cluster_env.sh"
lmflow_activate_conda

nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
slurm_name=${idx}-of-${total}

export LOGLEVEL=INFO

# ============================================
# Training Parameters
# ============================================
model_name_or_path=src/models/Qwen2.5-VL-3B-Instruct-mdm
tokenizer_name=Qwen/Qwen2.5-VL-3B-Instruct
dataset_path=fast_dllmv2_mmmathinstruct150k
conversation_template=qwen2_5_no_reasoning
trust_remote_code=1

export CUDA_HOME=${CUDA_HOME:-/cm/shared/apps/cuda12.4/toolkit/12.4.1}

export MDM_VL_KEEP_VISION=1
export TOKENIZERS_PARALLELISM=false

export DATASET_BACKEND="${DATASET_BACKEND:-webdataset_multi_modal}"

job_name=fastdllmv2_mmmathinstruct150k_8node_bd32_block_causal_no_dynamic_lr_1e-5_ds
output_dir=output_models/${job_name}
deepspeed_args="--master_port=23456 --num_nodes=$SLURM_NNODES --node_rank=$SLURM_PROCID"

# ============================================
# Print Node Info
# ============================================
echo "=========================================="
echo "MDM_Qwen2.5-VL Training (FastDLLMV2 + MMMathInstruct150k - DeepSpeed 8 Nodes)"
echo "=========================================="
echo "Nodes allocated: $SLURM_JOB_NODELIST"
echo "Number of nodes: $SLURM_NNODES"
echo "head_node: $head_node"
echo "head_node_ip: $head_node_ip"
echo "slurm_name: $slurm_name"
echo ""
echo "Configuration:"
echo "  Model: ${model_name_or_path}"
echo "  Tokenizer: ${tokenizer_name}"
echo "  Dataset: ${dataset_path}"
echo "  Output: ${output_dir}"
echo "  LR: 1e-5"
echo "  BD_SIZE: 32"
echo "  Block Causal: 1"
echo "=========================================="

# ============================================
# Auto-detect latest checkpoint for resume
# ============================================
latest_checkpoint=""
if [ -d "${output_dir}" ]; then
    latest_checkpoint=$(find "${output_dir}" -name "checkpoint-*" -type d | sort -V | tail -1)
    if [ -n "${latest_checkpoint}" ]; then
        echo "Found checkpoint: ${latest_checkpoint}"
    else
        echo "No checkpoint found in ${output_dir}"
        latest_checkpoint=""
    fi
else
    echo "Output dir ${output_dir} does not exist, training from scratch"
    latest_checkpoint=""
fi

# Create directories
log_dir=log/${job_name}
mkdir -p ${output_dir} ${log_dir}

# Build resume argument
resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
fi

nvidia-smi --query-gpu=timestamp,memory.used,memory.total,utilization.gpu,pstate \
           --format=csv,noheader,nounits -l 60 &

rm -rf ${log_dir}/deepspeed-hostfile
rm -rf ${log_dir}/temp-hostfile

# ============================================
# Multi-node DeepSpeed Launch via srun
# ============================================
cmd='
cd "'"${PROJECT_ROOT}"'";
source "'"${PROJECT_ROOT}"'/scripts/cluster_env.sh";
lmflow_activate_conda;
if ! command -v deepspeed >/dev/null 2>&1; then
  echo "ERROR: deepspeed not found after conda activation (LMFLOW_CONDA_ENV='"${LMFLOW_CONDA_ENV}"')." 1>&2;
  exit 127;
fi
export PDSH_RCMD_TYPE=ssh;
export CUDA_HOME=/cm/shared/apps/cuda12.4/toolkit/12.4.1;
export MDM_VL_KEEP_VISION=1;
export TOKENIZERS_PARALLELISM=false;
export PYTHONPATH="'"${PROJECT_ROOT}"'/src:'"${PROJECT_ROOT}"'/third_party/vila:'"${PROJECT_ROOT}"'/examples:${PYTHONPATH:-}";
export DATASET_BACKEND="'"${DATASET_BACKEND}"'";
echo "SLURMD_NODENAME:= " $SLURMD_NODENAME 1>&2;
echo "head_node:= " '"${head_node}"' 1>&2;
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo $SLURM_JOB_NODELIST > '"${log_dir}"'/temp-hostfile;
  read master_node slave_node <<< $(./expand_hostfile.sh '"${log_dir}"');
fi
# Wait for all nodes
sleep 60;
./setup_ssh.sh '"${log_dir}"';
echo "MASTER_NODE:= " $master_node 1>&2;
# Only master node runs DeepSpeed
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo "Running on master node. Executing DeepSpeed..." 1>&2;
  deepspeed --master_addr=${master_node} --hostfile '"${log_dir}"'/deepspeed-hostfile '"${deepspeed_args}"' \
    examples/finetune_qwenvl.py \
    --model_name_or_path '"${model_name_or_path}"' \
    --tokenizer_name '"${tokenizer_name}"' \
    --trust_remote_code '"${trust_remote_code}"' \
    --dataset_path '"${dataset_path}"' \
    --output_dir '"${output_dir}"' \
    '"${resume_arg}"' \
    --conversation_template '"${conversation_template}"' \
    --num_train_epochs 1 \
    --learning_rate 1e-5 \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio 0.03 \
    --block_size 4096 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    --bf16 \
    --logging_steps 10 \
    --do_train \
    --ddp_timeout 72000 \
    --save_steps 50 \
    --dataloader_num_workers 1 \
    --learn_padding 1 \
    --bd_size 32 \
    --save_total_limit 20 \
    --mdm 1 \
    --always_mask_im_end 1 \
    --complementary_mask 1 \
    --max_length 4096 \
    --gradient_checkpointing 1 \
    --flexible_bd_size 0 \
    --use_block_causal_mask 1 \
    --report_to none \
    --freeze_vision_encoder 1 \
    --block_causal_no_dynamic
else
    echo "Not running on master node. Skipping DeepSpeed execution." 1>&2;
fi
'

echo $cmd
mkdir -p ${log_dir}/slurm_log/train
srun -o ${log_dir}/slurm_log/train/slurm_%A_${slurm_name}.log -e ${log_dir}/slurm_log/train/error_%A_${slurm_name}.log \
    bash -c "${cmd}"
