#!/bin/bash

#SBATCH --job-name=search
#SBATCH --output=logs/%x/%A/%j/stdout.log

#SBATCH --partition=flame
#SBATCH --time=07-00:00:00
#SBATCH --qos=flame-t1b_g1_qos
#SBATCH --array=0-79%8

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --mem=1536G
#SBATCH --cpus-per-task=208
#SBATCH --gres=gpu:8

source scripts/config.sh
source scripts/secret.sh

configs=(
    "  256  1368  176  6  [0]*1+[1]*5  32 8   2377 1e18" # 0
    "  768  4104  528  6  [0]*1+[1]*5  16 8    543 1e18" # 1
    "  512  2736  352  6  [0]*1+[1]*5  32 8    967 1e18" # 2
    "  512  2736  352  9  [0]*1+[1]*8  32 8    808 1e18" # 3
    "  256  1368  176  9  [0]*1+[1]*8  32 8   2121 1e18" # 4
    "  768  4104  528  9  [0]*1+[1]*8  16 8    436 1e18" # 5
    " 1536  8208 1056 12 [0]*1+[1]*11   4 8    111 1e18" # 6
    " 1024  5472  704 12 [0]*1+[1]*11   8 8    224 1e18" # 7
    "  768  4104  528 12 [0]*1+[1]*11  16 8    363 1e18" # 8
    "  512  2736  352 12 [0]*1+[1]*11  16 8    695 1e18" # 9
    " 1024  5472  704 15 [0]*1+[1]*14   8 8    190 1e18" # 10
    " 2048 10944 1408 15 [0]*1+[1]*14   1 8     55 1e18" # 11
    " 1536  8208 1056 15 [0]*1+[1]*14   4 8     92 1e18" # 12
    " 1536  8208 1056 18 [0]*1+[1]*17   4 8     79 1e18" # 13
    " 1024  5472  704 18 [0]*1+[1]*17   8 8    165 1e18" # 14
    " 2048 10944 1408 18 [0]*1+[1]*17   1 8     46 1e18" # 15
    "  256  1368  176  6  [0]*1+[1]*5  32 8   7130 3e18" # 16
    "  768  4104  528  6  [0]*1+[1]*5  16 8   1629 3e18" # 17
    "  512  2736  352  6  [0]*1+[1]*5  32 8   2900 3e18" # 18
    "  512  2736  352  9  [0]*1+[1]*8  32 8   2424 3e18" # 19
    "  256  1368  176  9  [0]*1+[1]*8  32 8   6363 3e18" # 20
    "  768  4104  528  9  [0]*1+[1]*8  16 8   1306 3e18" # 21
    " 1536  8208 1056 12 [0]*1+[1]*11   4 8    331 3e18" # 22
    " 1024  5472  704 12 [0]*1+[1]*11   8 8    672 3e18" # 23
    "  768  4104  528 12 [0]*1+[1]*11  16 8   1089 3e18" # 24
    "  512  2736  352 12 [0]*1+[1]*11  16 8   2083 3e18" # 25
    " 1024  5472  704 15 [0]*1+[1]*14   8 8    569 3e18" # 26
    " 2048 10944 1408 15 [0]*1+[1]*14   1 8    163 3e18" # 27
    " 1536  8208 1056 15 [0]*1+[1]*14   4 8    276 3e18" # 28
    " 1536  8208 1056 18 [0]*1+[1]*17   4 8    236 3e18" # 29
    " 1024  5472  704 18 [0]*1+[1]*17   8 8    493 3e18" # 30
    " 2048 10944 1408 18 [0]*1+[1]*17   1 8    138 3e18" # 31
    "  256  1368  176  6  [0]*1+[1]*5  32 8  14260 6e18" # 32
    "  768  4104  528  6  [0]*1+[1]*5  16 8   3258 6e18" # 33
    "  512  2736  352  6  [0]*1+[1]*5  32 8   5799 6e18" # 34
    "  512  2736  352  9  [0]*1+[1]*8  32 8   4848 6e18" # 35
    "  256  1368  176  9  [0]*1+[1]*8  32 8  12726 6e18" # 36
    "  768  4104  528  9  [0]*1+[1]*8  16 8   2611 6e18" # 37
    " 1536  8208 1056 12 [0]*1+[1]*11   4 8    662 6e18" # 38
    " 1024  5472  704 12 [0]*1+[1]*11   8 8   1344 6e18" # 39
    "  768  4104  528 12 [0]*1+[1]*11  16 8   2178 6e18" # 40
    "  512  2736  352 12 [0]*1+[1]*11  16 8   4165 6e18" # 41
    " 1024  5472  704 15 [0]*1+[1]*14   8 8   1137 6e18" # 42
    " 2048 10944 1408 15 [0]*1+[1]*14   1 8    325 6e18" # 43
    " 1536  8208 1056 15 [0]*1+[1]*14   4 8    551 6e18" # 44
    " 1536  8208 1056 18 [0]*1+[1]*17   4 8    472 6e18" # 45
    " 1024  5472  704 18 [0]*1+[1]*17   8 8    986 6e18" # 46
    " 2048 10944 1408 18 [0]*1+[1]*17   1 8    276 6e18" # 47
    "  256  1368  176  6  [0]*1+[1]*5  32 8  23767 1e19" # 48
    "  768  4104  528  6  [0]*1+[1]*5  16 8   5429 1e19" # 49
    "  512  2736  352  6  [0]*1+[1]*5  32 8   9664 1e19" # 50
    "  512  2736  352  9  [0]*1+[1]*8  32 8   8080 1e19" # 51
    "  256  1368  176  9  [0]*1+[1]*8  32 8  21210 1e19" # 52
    "  768  4104  528  9  [0]*1+[1]*8  16 8   4351 1e19" # 53
    " 1536  8208 1056 12 [0]*1+[1]*11   4 8   1102 1e19" # 54
    " 1024  5472  704 12 [0]*1+[1]*11   8 8   2240 1e19" # 55
    "  768  4104  528 12 [0]*1+[1]*11  16 8   3630 1e19" # 56
    "  512  2736  352 12 [0]*1+[1]*11  16 8   6942 1e19" # 57
    " 1024  5472  704 15 [0]*1+[1]*14   8 8   1895 1e19" # 58
    " 2048 10944 1408 15 [0]*1+[1]*14   1 8    541 1e19" # 59
    " 1536  8208 1056 15 [0]*1+[1]*14   4 8    918 1e19" # 60
    " 1536  8208 1056 18 [0]*1+[1]*17   4 8    786 1e19" # 61
    " 1024  5472  704 18 [0]*1+[1]*17   8 8   1643 1e19" # 62
    " 2048 10944 1408 18 [0]*1+[1]*17   1 8    460 1e19" # 63
    "  256  1368  176  6  [0]*1+[1]*5  32 8  71300 3e19" # 64
    "  768  4104  528  6  [0]*1+[1]*5  16 8  16286 3e19" # 65
    "  512  2736  352  6  [0]*1+[1]*5  32 8  28992 3e19" # 66
    "  512  2736  352  9  [0]*1+[1]*8  32 8  24239 3e19" # 67
    "  256  1368  176  9  [0]*1+[1]*8  32 8  63628 3e19" # 68
    "  768  4104  528  9  [0]*1+[1]*8  16 8  13052 3e19" # 69
    " 1536  8208 1056 12 [0]*1+[1]*11   4 8   3306 3e19" # 70
    " 1024  5472  704 12 [0]*1+[1]*11   8 8   6718 3e19" # 71
    "  768  4104  528 12 [0]*1+[1]*11  16 8  10889 3e19" # 72
    "  512  2736  352 12 [0]*1+[1]*11  16 8  20825 3e19" # 73
    " 1024  5472  704 15 [0]*1+[1]*14   8 8   5685 3e19" # 74
    " 2048 10944 1408 15 [0]*1+[1]*14   1 8   1621 3e19" # 75
    " 1536  8208 1056 15 [0]*1+[1]*14   4 8   2752 3e19" # 76
    " 1536  8208 1056 18 [0]*1+[1]*17   4 8   2358 3e19" # 77
    " 1024  5472  704 18 [0]*1+[1]*17   8 8   4928 3e19" # 78
    " 2048 10944 1408 18 [0]*1+[1]*17   1 8   1379 3e19" # 79
)

i=0
for config in "${configs[@]}"; do
  read -r hidden_size ffn_hidden_size moe_ffn_hidden_size num_layers moe_layer_freq micro_batch_size expert_model_parallel_size train_iters flops <<< "$config"
    if [[ $SLURM_ARRAY_TASK_ID -eq $i ]]; then
        export HIDDEN_SIZE=$hidden_size
        export FFN_HIDDEN_SIZE=$ffn_hidden_size
        export MOE_FFN_HIDDEN_SIZE=$moe_ffn_hidden_size
        export NUM_LAYERS=$num_layers
        export MOE_LAYER_FREQ=$moe_layer_freq
        export MICRO_BATCH_SIZE=$micro_batch_size
        export EXPERT_MODEL_PARALLEL_SIZE=$expert_model_parallel_size
        export PIPELINE_MODEL_PARALLEL_SIZE=1
        export TRAIN_ITERS=$train_iters
        export WANDB_RUN_GROUP="ablation-$flops"
        export WANDB_NAME="$SLURM_JOB_ID"
        export SAVE_INTERVAL=$train_iters
        export EVAL_INTERVAL=$train_iters
        export TRAIN_DATASET="$GCP_DATASET/dclm-138b/tokenized/EleutherAI/pythia-12b"
        bash scripts/training/flame-moe.sh
    fi
    ((i++))
done
