#!/bin/bash
#SBATCH --job-name=transformer_1_3b
#SBATCH --time=72:00:00
#SBATCH --nodes=1
#SBATCH --gres=gpu:h200:1
#SBATCH --exclude=dcc-h200-gpu-04,dcc-h200-gpu-05
#SBATCH --nodelist=dcc-h200-gpu-01,dcc-h200-gpu-02,dcc-h200-gpu-03,dcc-h200-gpu-06,dcc-h200-gpu-07
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=32G
#SBATCH --partition=h200ea
#SBATCH --output=../../logs/%x-%j.out
#SBATCH --error=../../logs/%x-%j.err
#SBATCH --account=h200ea

cd ../../

if [ -f "./scripts/env/env.sh" ]; then
    source ./scripts/env/env.sh
fi

export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -n 1)
RANDOM_OFFSET=$((100 + RANDOM % 100))
export MASTER_PORT=$((29500 + RANDOM_OFFSET))

echo "MASTER_ADDR: $MASTER_ADDR"
echo "MASTER_PORT: $MASTER_PORT"
echo "SLURM_NNODES: $SLURM_NNODES"
echo "SLURM_NTASKS_PER_NODE: $SLURM_NTASKS_PER_NODE"
echo "SLURM_PROCID: $SLURM_PROCID"
echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"

# torchrun \
#   --nnodes=$SLURM_NNODES \
#   --nproc_per_node=1 \
#   --node_rank=$SLURM_PROCID \
#   --master_addr=$MASTER_ADDR \
#   --master_port=$MASTER_PORT \
#   /work/jf381/code/lm-research/main.py mode=generate +training.resume_checkpoint_path="/hpc/group/carin/shared/hf_models/checkpoints/fst_1_3B/best" checkpoint_name=fst_1_3B_resume "$@"

# torchrun \
#   --nnodes=$SLURM_NNODES \
#   --nproc_per_node=8 \
#   --node_rank=$SLURM_PROCID \
#   --master_addr=$MASTER_ADDR \
#   --master_port=$MASTER_PORT \
#   /work/jf381/code/lm-research/main.py seed=3407 +training.resume_checkpoint_path="/hpc/group/carin/shared/hf_models/checkpoints/fst_353M/best" checkpoint_name=fst_353M_resume_3407 "$@"

# torchrun \
#   --nnodes=$SLURM_NNODES \
#   --nproc_per_node=1 \
#   --node_rank=$SLURM_PROCID \
#   --master_addr=$MASTER_ADDR \
#   --master_port=$MASTER_PORT \
#   /work/jf381/code/lm-research/main.py seed=3407 mode=generate checkpoint_name=fst_353M_resume "$@"

# torchrun \
#   --nnodes=$SLURM_NNODES \
#   --nproc_per_node=1 \
#   --node_rank=$SLURM_PROCID \
#   --master_addr=$MASTER_ADDR \
#   --master_port=$MASTER_PORT \
#   /work/jf381/code/lm-research/main.py seed=3407 mode=generate checkpoint_name=fst_353M_resume_3407 "$@"

  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=8 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py  +training.resume_checkpoint_path="/work/jf381/code/lm-research/pretrain_weights/fst_353M" checkpoint_name=fst_353M_resume_new "$@"


  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=1 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py  mode=generate checkpoint_name=fst_353M_resume_new "$@"


  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=4 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py  +training.resume_checkpoint_path="/work/jf381/code/lm-research/pretrain_weights/fst_1_3B" checkpoint_name=fst_1_3B_resume_new "$@"

  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=4 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py  +training.resume_checkpoint_path="/cwork/jf381/code/lm-research_backup/pretrain_weights/fst_353M" checkpoint_name=fst_353M_resume_new_4gpu_rerun_wo_fi_mlp "$@"

# export CUDA_VISIBLE_DEVICES=4,5,6,7 &&  torchrun \
#   --nnodes=$SLURM_NNODES \
#   --nproc_per_node=4 \
#   --node_rank=$SLURM_PROCID \
#   --master_addr=$MASTER_ADDR \
#   --master_port=$MASTER_PORT \
#   /work/jf381/code/lm-research/main.py  +training.resume_checkpoint_path="/cwork/jf381/code/lm-research_backup/pretrain_weights/fst_353M" checkpoint_name=fst_353M_resume_new_4gpu_rerun_wo_phi_mlp "$@"

  torchrun \
  --nnodes=$SLURM_NNODES \
  --nproc_per_node=4 \
  --node_rank=$SLURM_PROCID \
  --master_addr=$MASTER_ADDR \
  --master_port=$MASTER_PORT \
  /work/jf381/code/lm-research/main.py \
  model=fst_353M \
  size=medium \
  training=medium \
  dataset=tinygsm_resume \
  checkpoint_name=fst_353M_resume_new_4gpu_rerun_wo_fi_phi_mlp "$@"


  export CUDA_VISIBLE_DEVICES=4 &&  torchrun \
  --nnodes=$SLURM_NNODES \
  --nproc_per_node=1 \
  --node_rank=$SLURM_PROCID \
  --master_addr=$MASTER_ADDR \
  --master_port=$MASTER_PORT \
  /work/jf381/code/lm-research/main.py \
  model=fst_353M \
  size=medium \
  training=medium \
  mode=generate \
  dataset=tinygsm_resume \
  checkpoint_name=fst_353M_resume_new_4gpu_rerun_wo_fi_phi_mlp "$@"

# export CUDA_VISIBLE_DEVICES=5 &&  torchrun \
#   --nnodes=$SLURM_NNODES \
#   --nproc_per_node=1 \
#   --node_rank=$SLURM_PROCID \
#   --master_addr=$MASTER_ADDR \
#   --master_port=$MASTER_PORT \
#   /work/jf381/code/lm-research/main.py model=fst_353M mode=generate checkpoint_name=fst_353M_resume_new_4gpu_rerun_wo_phi_mlp "$@"




  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=4 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py seed=34077  +training.resume_checkpoint_path="/cwork/jf381/code/lm-research_backup/pretrain_weights/fst_353M" checkpoint_name=fst_353M_resume_new_4gpu_rerun_seed34077 "$@"

  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=1 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py +training.resume_checkpoint_path="/hpc/group/carin/shared/hf_models/checkpoints/transformer_353M/best" checkpoint_name=transformer_353M_resume_new_4gpu_rerun model=transformer_353M  "$@"



  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=1 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py mode=generate checkpoint_name=fst_353M_resume_new "$@"



  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=1 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py mode=generate  checkpoint_name=fst_1_3B_resume "$@"


  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=8 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py seed=3407 checkpoint_name=fst_353M_3407 "$@"



  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=8 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py seed=42  +training.resume_checkpoint_path="/hpc/group/carin/shared/hf_models/checkpoints/fst_353M/best" checkpoint_name=fst_353M_resume_42 "$@"

  # torchrun \
  # --nnodes=$SLURM_NNODES \
  # --nproc_per_node=8 \
  # --node_rank=$SLURM_PROCID \
  # --master_addr=$MASTER_ADDR \
  # --master_port=$MASTER_PORT \
  # /work/jf381/code/lm-research/main.py seed=2000  +training.resume_checkpoint_path="/hpc/group/carin/shared/hf_models/checkpoints/fst_353M/best" checkpoint_name=fst_353M_resume_2000 "$@"


# torchrun \
#   --nnodes=$SLURM_NNODES \
#   --nproc_per_node=8 \
#   --node_rank=$SLURM_PROCID \
#   --master_addr=$MASTER_ADDR \
#   --master_port=$MASTER_PORT \
#   /work/jf381/code/lm-research/main.py seed=42  mode=generate checkpoint_name=fst_353M_resume_42 "$@"