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

set -e

echo "=========================================="
echo "MDM_Qwen2VL VLM Training (ShareGPT4V, multi-node)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "Nodes: ${SLURM_NNODES:-8}"
echo "GPUs/Node: 8"
echo "CPUs/Task: 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 (first node):"
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

# Recommended NCCL envs for multi-node
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_NET_GDR_LEVEL=2

# 配置（用户给定的 ShareGPT4V 数据路径）
exp_id=train_sharegpt4v_vlm_${SLURM_JOB_ID:-$(date +%Y%m%d_%H%M%S)}
model_name_or_path=/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_MDM
tokenizer_name=Qwen/Qwen2.5-VL-3B-Instruct

dataset_json="/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/jason-filter-sharegpt4v_mix665k_cap23k_coco-ap9k_lcs3k_sam9k_div2k.json"
image_folder="/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/data"

output_dir=output_models/${exp_id}
log_dir=log/${exp_id}
mkdir -p ${output_dir} ${log_dir}

echo "=========================================="
echo "Configuration:"
echo "  Model: ${model_name_or_path}"
echo "  Tokenizer: ${tokenizer_name}"
echo "  Dataset JSON: ${dataset_json}"
echo "  Image Folder: ${image_folder}"
echo "  Output: ${output_dir}"
echo "  Num Train Epochs: 1"
echo "  Sequence Length: 2048"
echo "  per_device_train_batch_size: 4"
echo "  gradient_accumulation_steps: 2"
echo "  Effective batch size: 4 × 2 × 8 GPUs × ${SLURM_NNODES:-8} nodes"
echo "  Precision: bf16"
echo "  bd_size: 128"
echo "  learn_padding: 1"
echo "  conplemenrary_mask: 1"
echo "=========================================="
echo ""

# Preflight checks and lightweight schema adapter (only if needed)
echo "Checking dataset schema compatibility..."
ADAPTED_JSON="${output_dir}/sharegpt4v_adapted.json"
python3 - <<PY || { echo "Schema preflight failed"; exit 1; }
import json, os, sys
dataset_json = os.environ.get('dataset_json')
image_folder = os.environ.get('image_folder')
adapted_json = os.environ.get('ADAPTED_JSON')

with open(dataset_json, 'r') as f:
    data = json.load(f)

if not isinstance(data, list) or len(data) == 0:
    print("❌ JSON must be a list with at least one item", flush=True)
    sys.exit(2)

sample = data[0]
keys = set(sample.keys())

has_conversations = 'conversations' in sample or 'messages' in sample
has_image = any(k in sample for k in ['image', 'image_path', 'image_paths', 'images'])

if not has_conversations:
    print("❌ Missing conversations/messages field in dataset items", flush=True)
    sys.exit(3)

# Decide if adaptation is needed
need_adapt = False
field_conversations = 'conversations' if 'conversations' in sample else 'messages'

def normalize_role_to_from(role):
    if role in ('user', 'human'):
        return 'human'
    if role in ('assistant', 'gpt', 'bot'):
        return 'gpt'
    return role

# Check message keys
msg0 = sample[field_conversations][0] if sample[field_conversations] else {}
if 'from' not in msg0 and 'role' in msg0:
    need_adapt = True

if 'image' not in sample and any(k in sample for k in ['image_path', 'image_paths', 'images']):
    need_adapt = True

if not need_adapt:
    print("✅ Schema looks compatible (conversations/from + image)", flush=True)
    print(dataset_json)
    sys.exit(0)

print("ℹ️ Adapting dataset schema to LLaVA format...", flush=True)
out = []
for item in data:
    new_item = dict(item)
    # Normalize conversations/messages
    conv_key = 'conversations' if 'conversations' in item else ('messages' if 'messages' in item else None)
    if conv_key is None:
        continue
    conv = []
    for m in item[conv_key]:
        if 'from' in m and 'value' in m:
            conv.append(m)
        else:
            # role/content/text to from/value
            role = m.get('role')
            text = m.get('value') or m.get('text') or m.get('content') or ''
            conv.append({'from': normalize_role_to_from(role), 'value': text})
    new_item['conversations'] = conv
    # Normalize image field
    if 'image' in item:
        pass
    elif 'image_path' in item and isinstance(item['image_path'], str):
        new_item['image'] = item['image_path']
    elif 'image_paths' in item and isinstance(item['image_paths'], list) and len(item['image_paths'])>0:
        new_item['image'] = item['image_paths'][0]
    elif 'images' in item and isinstance(item['images'], list) and len(item['images'])>0:
        new_item['image'] = item['images'][0]
    else:
        # If totally missing, skip item
        continue
    out.append(new_item)

with open(adapted_json, 'w') as f:
    json.dump(out, f)

print(adapted_json)
PY

# Capture possibly adapted path from python output
schema_result=$(tail -n 1 ${log_dir}/../${exp_id}/../../${exp_id} 2>/dev/null || true)

# The python prints the dataset path on success; capture it reliably
DATASET_TO_USE=$(python3 - <<'PY'
import sys, os
log = os.environ.get('ADAPTED_JSON')
ds = os.environ.get('dataset_json')
# Prefer adapted if it exists and is non-empty
if os.path.exists(log) and os.path.getsize(log) > 2:
    print(log)
else:
    print(ds)
PY
)

echo "Using dataset file: ${DATASET_TO_USE}"

# Keep vision modules when loading base VLM (MDM naming)
export MDM_VL_KEEP_VISION=1

# Launch deepspeed across nodes via srun
echo "Starting multi-node training..."
srun --nodes=${SLURM_NNODES} --ntasks-per-node=1 --gpus-per-task=8 --cpus-per-task=16 \
    deepspeed --num_nodes ${SLURM_NNODES} --num_gpus 8 examples/finetune_qwenvl.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_TO_USE} \
        --image_folder ${image_folder} \
        --output_dir ${output_dir} \
        --num_train_epochs 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 4 \
        --gradient_accumulation_steps 2 \
        --gradient_checkpointing \
        --run_name ${exp_id} \
        --overwrite_output_dir \
        --logging_steps 50 \
        --do_train \
        --validation_split_percentage 0 \
        --ddp_timeout 72000 \
        --save_steps 2000 \
        --save_total_limit 10 \
        --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.log

exit_code=${PIPESTATUS[0]}

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ ShareGPT4V VLM Training Success!"
else
    echo "❌ ShareGPT4V VLM Training Failed (exit code: ${exit_code})"
    echo -e "Please check logs:\n  Output: ${log_dir}/train.log\n  Error: workdir/train_sharegpt4v_vlm_error.log"
fi
echo "=========================================="

exit $exit_code


