#!/bin/bash
#SBATCH --time=3-00:00:00
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:h200:2
#SBATCH --mem=128G
#SBATCH --partition=h200ea
#SBATCH --output=../../logs/%x-%j.out
#SBATCH --error=../../logs/%x-%j.err
#SBATCH --account=h200ea


# Optional arguments [MODEL_NAME] [NUM_GPUS] [NUM_NODES]
NUM_GPUS=${1:-4}                # default 4 GPUs per node
NUM_NODES=${2:-1}               # default 1 node

# Clear first 2 args to avoid hydra issues
shift 2

# CPUs and memory automatically determined 
CPUS_PER_TASK=8                             # 8 CPUs per GPU
TOTAL_CPUS=$((NUM_GPUS * CPUS_PER_TASK))
TOTAL_MEM=$((NUM_GPUS * 32))                # 32gb ram per GPU

# Extract model name from hydra overrides
MODEL_NAME="transformer"
for arg in "$@"; do
  if [[ "$arg" == model=* ]]; then
    MODEL_NAME="${arg#model=}"
    break
  fi
done

JOB_NAME="${MODEL_NAME}_train_medium"

DATE=$(date +%F)

echo "Launching training job: $JOB_NAME"
echo "Nodes: ${NUM_NODES}, GPUs per node: ${NUM_GPUS}"
echo "Resources per node: ${TOTAL_CPUS} CPUs, ${TOTAL_MEM} GB memory"

SBATCH_ARGS="--job-name=${JOB_NAME} \
             --nodes=${NUM_NODES} \
             --gres=gpu:h200:${NUM_GPUS} \
             --ntasks-per-node=${NUM_GPUS} \
             --cpus-per-task=${CPUS_PER_TASK} \
             --mem=${TOTAL_MEM}G \
             --output=../../logs/${MODEL_NAME}/${DATE}/%x-%j.out \
             --error=../../logs/${MODEL_NAME}/${DATE}/%x-%j.err"

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3b.sh \
    # size=1_3b \
    # training=1_3b \
    # dataset=tinygsm_resume \
    # "$@"

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=transformer_1_3B

# bert_pretrain
# cp /work/jf381/code/lm-research/main_1_2.py /work/jf381/code/lm-research/main.py
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_fst_353M_bert.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm \
#     "$@" \
#     model=fst_353M_bert
bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M.sh \
    size=medium \
    training=medium \
    dataset=slimpajama_6b \
    "$@" \
    training.lr=2e-6 \
    model=transformer_353M

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_fst_353M_bert1.sh \
#     size=medium \
#     training=medium \
#     dataset=slimpajama_6b \
#     "$@" \
#     training.lr=2e-5 \
#     model=fst_353M_bert

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_fst_353M_bert2.sh \
#     size=medium \
#     training=medium \
#     dataset=slimpajama_6b \
#     "$@" \
#     training.lr=2e-4 \
#     model=fst_353M_bert


# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M_bert.sh \
#     size=medium \
#     training=medium \
#     dataset=slimpajama_6b \
#     training.lr=2e-6 \
#     "$@" \
#     model=transformer_353M_bert

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M_bert1.sh \
#     size=medium \
#     training=medium \
#     dataset=slimpajama_6b \
#     training.lr=2e-5 \
#     "$@" \
#     model=transformer_353M_bert
  
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M_bert2.sh \
#     size=medium \
#     training=medium \
#     dataset=slimpajama_6b \
#     training.lr=2e-4 \
#     "$@" \
#     model=transformer_353M_bert

# slimpajama_6b
# bert_pretrain
# bash /work/jf381/code/lm-research/scripts/training/train_medium_bash_resume_transformer_bert.sh;
# bash /work/jf381/code/lm-research/scripts/training/train_medium_bash_resume_fst_bert.sh;
# ork/jf381/code/lm-research/main_1_2.py /work/jf381/code/lm-research/main.py
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M_bert.sh \
#     size=medium \
#     training=medium \
#     dataset=slimpajama_6b \
#     "$@" \
#     model=transformer_353M_bert

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3B.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=transformer_1_3B

# cp /work/jf381/code/lm-research/main_1_7.py /work/jf381/code/lm-research/main.py
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3B.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_1_3B


# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_353M

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_353M.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_1_3B

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=transformer_353M
# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3b.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@"

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_generate.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@"
  ## left_small 02-
  ## 04 big-
  ## 03 fst big

# # cat << 'EOF' > 1_3b_transformer.txt
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3b_generate.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@"
# EOF