#!/bin/bash -l
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=32
#SBATCH --mem=450G
#SBATCH --gres=gpu:4
#SBATCH --time=48:00:00
#SBATCH --account=<todo>
#SBATCH --partition=<todo>
#SBATCH --output=logs/slurm_%x-%j.out

set -euo pipefail

module load ML-bundle/25.04

# Virtual environment
source .venv-gh200/bin/activate

# Hugging Face
export HF_TOKEN=""   # TODO: DO NOT COMMIT THIS

export HF_HOME=$SCRATCH/huggingface
export HF_HUB_CACHE=$SCRATCH/models
export HF_HUB_ENABLE_HF_TRANSFER=1

# WandB
export WANDB_API_KEY="" # TODO
export WANDB_PROJECT="" # TODO

# NCCL
export NCCL_DEBUG=INFO
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export FI_PROVIDER="cxi"
export FI_CXI_DISABLE_HOST_REGISTER=1
export NCCL_SOCKET_IFNAME=bond0

# Other
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export OMP_NUM_THREADS=4
export MKL_NUM_THREADS=1


# Mutli node training
export HEAD_NODE_IP=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MAIN_PROCESS_PORT=$((29500 + $SLURM_JOB_ID % 100))

# Training
export BASE_MODEL="models/base/bielik-11b-v2.6-instruct-siglip2-so400m-patch14-384-llava-next"
export RUN_NAME="$(basename ${BASE_MODEL})-pretrained-$(date +%Y-%m-%d_%H-%M-%S)"

export DATA_PATH="data/processed/pretrain"
export IMAGES_PATH="data/static/pretrain"

export PER_DEVICE_TRAIN_BATCH_SIZE=1
export GLOBAL_BATCH_SIZE=256
export GRADIENT_ACCUMULATION_STEPS=$((GLOBAL_BATCH_SIZE / (SLURM_GPUS_ON_NODE * SLURM_NNODES * PER_DEVICE_TRAIN_BATCH_SIZE)))


export LAUNCHER="accelerate launch \
    --num_processes=$((SLURM_NNODES * SLURM_GPUS_ON_NODE)) \
    --num_machines=$SLURM_NNODES \
    --machine_rank=$SLURM_PROCID \
    --rdzv_backend=c10d \
    --main_process_ip=$HEAD_NODE_IP \
    --main_process_port=$MAIN_PROCESS_PORT \
    --mixed_precision=bf16 \
    --dynamo_backend=no
    "
export SCRIPT="-m src.modeling.llava.train"
export SCRIPT_ARGS="\
    --deepspeed=config/deepspeed/zero2.json \
    --model_name_or_path=${BASE_MODEL} \
    --data_path=${DATA_PATH} \
    --images_path=${IMAGES_PATH} \
    --trainer_output_dir=models/checkpoints \
    --model_output_dir=models/pretrained \
    --run_name=${RUN_NAME} \
    --seed=42 \
    --train_vision_tower=False \
    --train_multimodal_projector=True \
    --train_language_model=False \
    --train_lm_head=False \
    --max_seq_length=8192 \
    --epochs=1 \
    --lr=1e-3 \
    --lr_scheduler_type=cosine \
    --weight_decay=0.00 \
    --warmup_ratio=0.03 \
    --per_device_train_batch_size=${PER_DEVICE_TRAIN_BATCH_SIZE} \
    --per_device_eval_batch_size=2 \
    --gradient_accumulation_steps=${GRADIENT_ACCUMULATION_STEPS} \
    --gradient_checkpointing=True \
    --save_steps=250 \
    --logging_steps=10 \
    --report_to=wandb \
    --eval_strategy=steps \
    --eval_steps=250 \
    --test_size=0.01 \
    "

export CMD="${LAUNCHER} ${SCRIPT} ${SCRIPT_ARGS}"
srun ${CMD}

srun python -m src.modeling.llava.merge \
    --base_model_name_or_path="${BASE_MODEL}" \
    --trainer_output_dir="models/checkpoints/${RUN_NAME}" \
    --final_model_name_or_path="models/pretrained/${RUN_NAME}"