#!/bin/bash

# ======================== 参数解析与检查 ========================
# 默认参数

export DISABLE_MEMORY_REPORT=1

DATA_FORMAT="bf16"  # Llama3目前不支持切换，保留参数
DATASET_DIR=""      # 必须参数

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case "$1" in
        --data_format)
            DATA_FORMAT="$2"
            shift 2
            ;;
        --dataset_dir)
            DATASET_DIR="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# 检查必要参数
if [[ -z "$DATASET_DIR" ]]; then
    echo "Error: --dataset_dir is required"
    exit 1
fi

# ======================== 基础配置 ========================
# 生成唯一实验ID
CURRENT_TIME=$(date "+%Y%m%d_%H%M%S")
WORK_HOME="$PWD"
PATCH_HOME=/home/megatron-lm-musa-patch #"$PWD/../.."
HOSTFILE="./hostfile"

# 创建目录
mkdir -p "./output/$CURRENT_TIME"
mkdir -p "./logs" "./tboard" "./wandb" "./output_log/$CURRENT_TIME"

#mv /home/XLC_2025_llama3/flagscale_ckpt_tp8pp4/ ./checkpoints/tp8_pp4_dp1_mbs1_numbs512_gbs512_gpus32_bf16/


# ======================== 训练并行配置 ========================
# Llama3 8B并行参数
TP_SIZE=8
PP_SIZE=4
EP_SIZE=1
GPUS_PER_NODE=8
MICRO_BATCH_SIZE=1
NUM_MICROBATCHES=512

# 多机配置
NUM_NODES=$(grep -v '^$' "$HOSTFILE" | grep -v '^#' | wc -l)
WORLD_SIZE=$((GPUS_PER_NODE * NUM_NODES))
DP_SIZE=$((WORLD_SIZE / (TP_SIZE * PP_SIZE)))
GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * DP_SIZE))

# 实验名称
EXPNAME="tp${TP_SIZE}_pp${PP_SIZE}_dp${DP_SIZE}_mbs${MICRO_BATCH_SIZE}_numbs${NUM_MICROBATCHES}_gbs${GLOBAL_BATCH_SIZE}_gpus${WORLD_SIZE}_${DATA_FORMAT}"

# ======================== 环境变量配置 ========================
# 获取本机IP和节点信息
# export NODE_ADDR=$(ip a | grep -E 'inet (?!127.0.0.1)' | grep -v inet6 | grep -v -E 'docker0|cni0|lo|veth' | awk '{print $2}' | cut -d '/' -f1 | head -n1)
# export NODE_ADDR=$(ip a | grep -E 'inet (?!127.0.0.1)' | grep -v inet6 | grep -v -E 'docker0|cni0|lo|veth' | awk '{print $2}' | cut -d '/' -f1 | head -n1)
export NODE_ADDR=$(ip -4 addr show | grep -A1 'bond0:' | grep -oP '(?<=inet\s)\d+(\.\d+){3}')
export GPUS_PER_NODE=$GPUS_PER_NODE
export NUM_NODES=$NUM_NODES
export MASTER_ADDR=$(grep -v '^$' "$HOSTFILE" | grep -v '^#' | head -n1 | awk '{print $1}')
export NODE_RANK=$(awk '/^$/ || /^#/ {next} {ranks[$1]=(NR-1)} END {print ranks["'"$NODE_ADDR"'"] + 0}' "$HOSTFILE")
# export NODE_RANK=${NODE_RANK:-0}  # 兜底默认值
export MASTER_PORT=14388
export RDZV_ID=$CURRENT_TIME

# 系统环境变量
export OMP_NUM_THREADS=4
export MUSA_VISIBLE_DEVICES='0,1,2,3,4,5,6,7'
export MUSA_KERNEL_TIMEOUT=3200000
export ACCELERATOR_BACKEND="musa"
export MCCL_PROTOS=2
export MCCL_CHECK_POINTERS=0
export CUDA_DEVICE_MAX_CONNECTIONS=1

# Python路径配置
MEGATRON_PATH=${PATCH_HOME}/../Megatron-LM
export PYTHONPATH=${MEGATRON_PATH}:${PATCH_HOME}:$PYTHONPATH

# 编译Megatron（如果未编译）
if [ ! -d "${MEGATRON_PATH}/build" ]; then
    echo "Compiling Megatron-LM..."
    cd "${MEGATRON_PATH}"
    python setup.py build_ext --inplace
    cd -
fi

# ======================== 路径配置 ========================
#CHECKPOINT_PATH="$WORK_HOME/checkpoints/$EXPNAME"
CHECKPOINT_PATH=/home/XLC_2025_llama3/flagscale_ckpt_tp8pp4/
DATA_PATH="${DATASET_DIR}/dataset/dedup-md5-pile-pile-cc_text_document"
LOG_PATH="$WORK_HOME/logs/$EXPNAME"
TB_PATH="$WORK_HOME/tboard/$EXPNAME"
WB_PATH="$WORK_HOME/wandb/$EXPNAME"
TOKENIZED_MODEL="${DATASET_DIR}/tokenizer"

# 创建目录
mkdir -p $CHECKPOINT_PATH $LOG_PATH $TB_PATH $WB_PATH
cp "$0" $LOG_PATH/  # 备份当前脚本

# ======================== 训练参数配置 ========================
# 分布式参数
DISTRIBUTED_ARGS=(
    --nproc_per_node $GPUS_PER_NODE
    --nnodes $NUM_NODES
    --node_rank $NODE_RANK
    --master_addr $MASTER_ADDR
    --master_port $MASTER_PORT
    --log_dir "$WORK_HOME/output_log/$RDZV_ID/$EXPNAME"
    --redirects ${LOG_REDIRECTS_LEVEL:-0}
)

# 模型参数（Llama3 8B）
MODEL_ARGS=(
    --num-layers 80
    --hidden-size 8192
    --ffn-hidden-size 28672
    --num-attention-heads 64
    --group-query-attention
    --num-query-groups 8
    --seq-length 8192
    --vocab-size 128256 #102400
    --max-position-embeddings 8192
    --norm-epsilon 1e-5
    --attention-dropout 0.0
    --hidden-dropout 0.0
    --disable-bias-linear
    --position-embedding-type rope
    --swiglu
    --normalization RMSNorm
    --untie-embeddings-and-output-weights
    --make-vocab-size-divisible-by 64
    --no-position-embedding
    --use-rotary-position-embeddings
    --rotary-base 500000
)

# 训练参数
TRAINING_ARGS=(
    --seed 42
    --micro-batch-size $MICRO_BATCH_SIZE
    --global-batch-size $GLOBAL_BATCH_SIZE
    # --train-samples 24414062
    # --train-iters 47684
    --train-samples 256000
    # --train-samples 6160066
    --init-method-std 0.02
    --use-mcore-models
    --no-bias-dropout-fusion
    --no-bias-swiglu-fusion
    # --use-distributed-optimizer
    --use-flash-attn
    --sequence-parallel
    # --recompute-granularity full
    # --recompute-method block
    # --recompute-num-layers 0
    --distributed-backend nccl
    # --no-load-optim
    # --no-load-rng
    --finetune
)

# 正则化参数
REGULARIZATION_ARGS=(
    --weight-decay 0.01
    --adam-beta1 0.9
    --adam-beta2 0.95
    --clip-grad 1.0
    --start-weight-decay 0
    --end-weight-decay 5e-07
    --weight-decay-incr-style cosine
)

# 学习率参数
WARMUP_STEPS=2000
WARMUP_SAMPLES=$((WARMUP_STEPS * GLOBAL_BATCH_SIZE))
WARMUP_SAMPLES=25600
LEARNING_RATE_ARGS=(
    --lr 5e-6
    --lr-decay-style cosine
    # --lr-warmup-samples ${WARMUP_SAMPLES}
    --lr-warmup-samples ${WARMUP_SAMPLES}
    # --lr-warmup-iters 2000 
    --min-lr 0
    --initial-loss-scale 4294967296
    --min-loss-scale 1.0
)

# 模型并行参数
MODEL_PARALLEL_ARGS=(
    --tensor-model-parallel-size $TP_SIZE
    --pipeline-model-parallel-size $PP_SIZE
    # --decoder-last-pipeline-num-layers 14
)

# 混合精度参数
MIXED_PRECISION_ARGS=(
    --bf16
    --attention-softmax-in-fp32
    --no-masked-softmax-fusion
    --accumulate-allreduce-grads-in-fp32
)

# 数据参数
DATA_ARGS=(
    --data-path "$DATA_PATH"
    --tokenizer-type HuggingFaceTokenizer
    --tokenizer-model "$TOKENIZED_MODEL"
    --split 1
)

# 日志与评估参数
EVAL_AND_LOGGING_ARGS=(
    --log-interval 1
    --log-throughput
    --save-interval 200000
    --eval-interval 1000
    --save "$CHECKPOINT_PATH"
    --load "$CHECKPOINT_PATH"
    --eval-iters 0
    --tensorboard-dir "$TB_PATH"
)

# Transformer Engine参数
TRANSFORMER_ENGINE_ARGS=(
    --transformer-impl transformer_engine
)

# ======================== 执行训练 ========================
echo -e "\n=== Training Configuration ==="
echo "Dataset dir: $DATASET_DIR"
echo "Data format: $DATA_FORMAT (Llama3 currently only supports BF16)"
echo "Hostfile: $(cat $HOSTFILE)"
echo "Global batch size: $GLOBAL_BATCH_SIZE"
echo "NODE_RANK: $NODE_RANK"
echo "MASTER_ADDR: $MASTER_ADDR"
echo "Experiment name: $EXPNAME"
echo -e "\n=== Starting Training ==="

# 拼接并执行命令
cmd="torchrun ${DISTRIBUTED_ARGS[@]} $WORK_HOME/pretrain_gpt.py \
    ${MODEL_ARGS[@]} \
    ${TRAINING_ARGS[@]} \
    ${REGULARIZATION_ARGS[@]} \
    ${LEARNING_RATE_ARGS[@]} \
    ${MODEL_PARALLEL_ARGS[@]} \
    ${MIXED_PRECISION_ARGS[@]} \
    ${DATA_ARGS[@]} \
    ${EVAL_AND_LOGGING_ARGS[@]} \
    ${TRANSFORMER_ENGINE_ARGS[@]}"

echo -e "\nCommand:\n$cmd\n"
eval $cmd