#!/bin/bash
# Copyright (c) 2024 BAAI. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License")
#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
# args
DATA_DIR=$1
GPUS_PER_NODE=$2
NNODES=$3
NODE_RANK=$4
MASTER_ADDR=$5
MASTER_PORT=$6
Compile_path=$7
MBS=$8
ITERS=$9
TP=${10}
PP=${11}
is_share_memory=${12}
let GBS=$NNODES*16

compile() {
    Target_path=${Compile_path}/Megatron-LM
    if [ -d "${Target_path}" ]; then
        echo "Compile path exists, skipping unzip"
        return
    fi
    unzip Megatron-LM.zip -d ${Compile_path}
}
if [ "$is_share_memory" = 1 ] && [ "$NODE_RANK" -eq 0 ]; then
  compile
elif [ "$is_share_memory" = 0 ]; then
  compile
fi

MEGAPATH=${Compile_path}/Megatron-LM
# Runs the "LLaMA3-8B" parameter model
export PYTHONPATH=$PYTHONPATH:$MEGAPATH
# algorithm args

MODEL_ARGS=" \
    --num-layers 32 \
    --hidden-size 4096 \
    --ffn-hidden-size 14336 \
    --num-attention-heads 32 \
    --group-query-attention \
    --num-query-groups 8 \
    --seq-length 8192 \
    --max-position-embeddings 8192 \
    --swiglu \
    --normalization RMSNorm \
    --norm-epsilon 1.0e-5 \
    --global-batch-size $GBS \
    --attention-softmax-in-fp32"

OPT_ARGS=" \
    --lr 1.0e-5 \
    --train-iters $ITERS \
    --lr-decay-iters $ITERS \
    --lr-decay-style cosine \
    --min-lr 1.0e-6 \
    --weight-decay 0.1 \
    --clip-grad 1.0 \
    --adam-beta1 0.9 \
    --adam-beta2 0.95 \
    --init-method-std 0.008"
    
ALGO_ARGS="$MODEL_ARGS $OPT_ARGS"

# data args

DATA_ARGS=" \
    --data-path $DATA_DIR/wudao_llama3bpe_content_document \
    --tokenizer-type Llama3Tokenizer \
    --tokenizer-model $DATA_DIR \
    --split 100,0,0
"

# training args

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

TRAINING_ARGS=" \
    --micro-batch-size $MBS \
    --tensor-model-parallel-size $TP \
    --pipeline-model-parallel-size $PP \
    --sequence-parallel \
    --bf16
"

# vendor args

VENDOR_ARGS=" \
    --transformer-impl transformer_engine \
    --use-distributed-optimizer \
    --use-mcore-models \
    --use-flash-attn
"

OUTPUT_ARGS=" --log-interval 1"

run_cmd="torchrun $DISTRIBUTED_ARGS $MEGAPATH/pretrain_gpt.py \
    $ALGO_ARGS \
    $DATA_ARGS \
    $TRAINING_ARGS \
    $VENDOR_ARGS \
    $OUTPUT_ARGS"

sleep 5
echo ${run_cmd}
eval ${run_cmd}
