#! /bin/bash
set -e

CUR_DIR=$(cd `dirname $0`; pwd)
DATESTR=$(date +"%m-%d-%H-%M")

get_prop() {
    grep "${1}" ${2} | cut -d'=' -f2
}

export NVTE_BIAS_GELU_NVFUSION=0
export NVTE_BIAS_DROPOUT_FUSION=0


MODE=${1-"ONLINE"} # [ONLINE, OFFLINE]
source $CUR_DIR/prepare_once.sh $MODE

GPUS_PER_NODE=8
MASTER_ADDR=localhost
MASTER_PORT=29600
NNODES=1
NODE_RANK=0
WORLD_SIZE=$((GPUS_PER_NODE*NNODES))
echo "ws: $WORLD_SIZE"

TP_SIZE=$4
PP_SIZE=$5
DP_SIZE=$((WORLD_SIZE/TP_SIZE/PP_SIZE))

MICRO_BATCH_SIZE=$2
ACCU_STEPS=$3
GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE*DP_SIZE*ACCU_STEPS))
SEQ_LENGTH=$6

HIDDEN_SIZE=4096
FFN_HIDDEN_SIZE=13696
NUM_LAYERS=28
NUM_HEADS=32

TRAIN_STEPS=10
LR=3e-4
MIN_LR=3e-5
WEIGHT_DECAY=0.1
GRAD_CLIP=1

OP_FUSION_OPTIONS="--no-persist-layer-norm \
                   --no-bias-gelu-fusion"

TE_OPTIONS="--transformer-impl transformer_engine \
            --tp-comm-overlap \
            --attention-softmax-in-fp32"

LLAMA_OPTIONS="--attention-dropout 0 \
               --hidden-dropout 0 \
               --use-rotary-position-embeddings \
               --untie-embeddings-and-output-weights \
               --swiglu \
               --normalization RMSNorm \
               --disable-bias-linear"

MODEL_OPTIONS="--num-layers $NUM_LAYERS \
               --hidden-size $HIDDEN_SIZE \
               --ffn-hidden-size $FFN_HIDDEN_SIZE \
               --num-attention-heads $NUM_HEADS \
               --seq-length $SEQ_LENGTH \
               --max-position-embeddings $SEQ_LENGTH \
               --micro-batch-size $MICRO_BATCH_SIZE \
               --global-batch-size $GLOBAL_BATCH_SIZE"

TRAINING_OPTIONS="--lr $LR \
                  --lr-decay-style cosine \
                  --min-lr $MIN_LR \
                  --adam-beta1 0.9 \
                  --adam-beta2 0.95 \
                  --adam-eps 1e-5 \
                  --weight-decay $WEIGHT_DECAY \
                  --clip-grad $GRAD_CLIP"

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE \
                  --nnodes $NNODES \
                  --node_rank $NODE_RANK \
                  --master_addr $MASTER_ADDR \
                  --master_port $MASTER_PORT"

cmd="eightrun runonly $DISTRIBUTED_ARGS pretrain_gpt.py \
    $MODEL_OPTIONS \
    $OP_FUSION_OPTIONS \
    $TE_OPTIONS \
    $TRAINING_OPTIONS \
    $LLAMA_OPTIONS \
    --tensor-model-parallel-size $TP_SIZE \
    --pipeline-model-parallel-size $PP_SIZE \
    --sequence-parallel \
    --train-iters $TRAIN_STEPS \
    --data-path $DATA_PATH \
    --tokenizer-type Llama2Tokenizer \
    --tokenizer-model $LLAMA_TOKENIZER_MODEL \
    --split 949,50,1 \
    --distributed-backend cncl \
    --log-interval 1 \
    --save-interval 10000 \
    --eval-interval 1000 \
    --recompute-granularity selective \
    --no-load-optim \
    --no-load-rng \
    --make-vocab-size-divisible-by 100 \
    --position-embedding-type rope \
    --no-position-embedding \
    --use-distributed-optimizer \
    --use-flash-attn \
    --eval-iters 1 \
    --bf16"

echo $cmd
eval $cmd 2>&1 | tee $CUR_DIR/logs/logs-chatglm3-6B-$DATESTR.txt
