#! /bin/bash

####################################################
#
# Usage:
#   bash train_qwen.sh [OPTIONS]
#
# OPTIONS 分为四类配置：
#
# 1. Torchrun 配置（分布式启动参数）
#    --master_addr ADDR        Master 节点地址（默认: localhost）
#    --master_port PORT        Master 节点监听端口（默认: 19981）
#    --node_num N              分布式作业的总节点数（默认: 1）
#    --node_rank R             当前节点的索引（从 0 开始， 默认: 0）
#    --nproc_per_node P        每个节点使用的 GPU 数量（默认: 8）
#
# 2. 模型选择
#    --model_size SIZE         模型参数规模（如 8 表示 8B，具体含义由训练脚本定义）（默认: 8）
#
# 3. 并行配置（用于控制模型并行策略）
#    --tp TP                   Tensor Parallel（张量并行）（默认: 2）
#    --pp PP                   Pipeline Parallel（流水线并行）（默认: 1）
#    --cp CP                   Context Parallel（上下文并行）（默认: 1）
#    --ep EP                   Expert Parallel（专家并行，仅 MoE 模型使用）（默认: 0）
#    --etp ETP                 Expert Tensor Parallel（专家张量并行，仅 MoE 模型使用）（默认: 0）
#
# 4. 其他训练配置
#    --mbs MBS                 Megatron 的 Micro Batch Size（默认: 1）
#    --mbn MBN                 Megatron 的 Micro Batch Num（默认: 64）
#    --seq_len LEN             训练时的序列长度（默认: 32768）
#    --train_iters IT          总训练迭代次数（默认: 30）
#    --precision PREC          训练数值精度（支持 bf16、fp16 等）（默认: bf16）
#    --checkpoint_path PATH    检查点存储路径（默认: 不使用检查点）
#    --data_source PATH        训练数据路径，支持 "mock" 以使用模拟数据（默认: 使用预设数据路径）
#
# 示例：
#    - 单机8卡训练 Qwen3-8B 30轮：
#        bash train_qwen.sh --model_size 8 --tp 2 --pp 1 --cp 1 --mbs 1 --seq_len 32768 --precision bf16
#    - 多机训练 Qwen3-8B 30轮（在每台节点运行，--node_rank要变）：
#        bash train_qwen.sh --master_addr 192.168.1.10 --node_num 2 --node_rank ${RANK} \
#          --model_size 8 --tp 2 --pp 1 --mbs 1 --seq_len 32768 --precision fp16
#
# 注意：
#   - 所有选项均为可选，未指定时使用默认值。
#
# supported model size:
#	  - qwen3: {8, 32, A3B}
#	  - qwen2.5: {72}
####################################################

cd /workspace/Megatron-LM
source /opt/TrainPerfBench/run_cmd.sh
source /opt/TrainPerfBench/parse_args.sh
source /opt/TrainPerfBench/parse_torchrun_args.sh
export CUDA_DEVICE_MAX_CONNECTIONS=1

# Parse args
parse_torchrun_args "$@"
parse_qwen_args "$@"

# Data settings
DEFAULT_DATA_PATH="/data/datasets/redpajama-qwen/redpajama-sample_text_document"
if [[ "${data_source}" == "mock" ]]; then
  DATASET_ARGS="--mock-data"
elif [[ -z "${data_source}" ]]; then
  DATASET_ARGS="--data-path ${DEFAULT_DATA_PATH}"
else
  DATASET_ARGS="--data-path ${data_source}"
fi
DATASET_ARGS+=" --num-workers 16"
TOKENIZER_PATH="/data/datasets/redpajama-qwen/tokenizer"
TOKENIZER_ARGS="--tokenizer-type HuggingFaceTokenizer --tokenizer-model $TOKENIZER_PATH"

# Checkpoint settings
if [[ -z "${checkpoint_path}" ]]; then
  CHECKPOINT_ARGS=""
else
  CHECKPOINT_ARGS=" \
  --load $checkpoint_path \
	--save $checkpoint_path \
	--no-load-optim \
	--no-load-rng \
  "
fi

# Model settings
if [[ "${model_size}" == "8" ]]; then
  MODEL_NAME="qwen3-8B"
  NUM_LAYERS=36
  HIDDEN_SIZE=4096
  FFN_HIDDEN_SIZE=12288
  NUM_ATTN_HEADS=32
  NUM_KV_HEADS=8
  ROPE_THETA=1000000
  RMS_NORM_EPS=1e-6
  GQA_ARGS=" \
    --group-query-attention \
    --num-query-groups ${NUM_KV_HEADS} \
    --kv-channels 128 \
    --qk-layernorm \
  "
  OTHER_ARGS=" \
    --disable-bias-linear \
    --untie-embeddings-and-output-weights \
  "
elif [[ "${model_size}" == "32" ]]; then
  MODEL_NAME="qwen3-32B"
  NUM_LAYERS=64
  HIDDEN_SIZE=5120
  FFN_HIDDEN_SIZE=25600
  NUM_ATTN_HEADS=64
  NUM_KV_HEADS=8
  ROPE_THETA=1000000
  RMS_NORM_EPS=1e-6
  GQA_ARGS=" \
    --group-query-attention \
    --num-query-groups ${NUM_KV_HEADS} \
    --kv-channels 128 \
    --qk-layernorm \
  "
  OTHER_ARGS=" \
    --disable-bias-linear \
    --untie-embeddings-and-output-weights \
  "
elif [[ "${model_size}" == "72" ]]; then
  MODEL_NAME="qwen2.5-72B"
  NUM_LAYERS=80
  HIDDEN_SIZE=8192
  FFN_HIDDEN_SIZE=29568
  NUM_ATTN_HEADS=64
  NUM_KV_HEADS=8
  GQA_ARGS=" \
    --group-query-attention \
    --num-query-groups ${NUM_KV_HEADS} \
  "
  OTHER_ARGS=" \
    --untie-embeddings-and-output-weights \
  "
elif [[ "${model_size}" == "A3B" ]]; then
  MODEL_NAME="qwen3-30B-A3B"
  NUM_LAYERS=48
  HIDDEN_SIZE=2048
  FFN_HIDDEN_SIZE=6144
  NUM_ATTN_HEADS=32
  NUM_KV_HEADS=4
  ROPE_THETA=1000000
  RMS_NORM_EPS=1e-6
  GQA_ARGS=" \
    --group-query-attention \
    --num-query-groups ${NUM_KV_HEADS} \
    --kv-channels 128 \
    --qk-layernorm \
  "
  OTHER_ARGS=" \
    --disable-bias-linear \
    --untie-embeddings-and-output-weights \
  "
  # MoE settings
  ROUTER_TOPK=8
  NUM_EXPERTS=128
  MOE_INTERMEDIATE_SIZE=768
  MOE_ARGS=" \
    --moe-grouped-gemm \
    --moe-token-dispatcher-type alltoall \
    --moe-router-topk ${ROUTER_TOPK} \
    --num-experts ${NUM_EXPERTS} \
    --expert-tensor-parallel-size ${etp} \
    --expert-model-parallel-size ${ep} \
    --moe-ffn-hidden-size ${MOE_INTERMEDIATE_SIZE} \
    --moe-router-load-balancing-type aux_loss \
    --moe-aux-loss-coeff 0.001 \
    --moe-layer-freq '([1]*48)' \
  "
else
  echo "NotImplementedError: Model size ${model_size} not supported."
  exit 2
fi

# Precision settings
if [[ "${precision}" == "fp8" ]]; then
  FP8_ARGS=" \
		--fp8-format hybrid \
		--fp8-amax-compute-algo max \
		--fp8-amax-history-len 1024 \
	"
  DT="fp8"
elif [[ "${precision}" == "bf16" ]]; then
  FP8_ARGS=""
  DT="bf16"
else
  echo "NotImplementedError: Model precision ${precision} not supported."
  exit 2
fi

# Node settings
WORLD_SIZE=$((${nproc_per_node} * ${node_num}))
DISTRIBUTED_ARGS=" \
  --nproc_per_node ${nproc_per_node} \
  --nnodes ${node_num} \
  --node_rank ${node_rank} \
  --master_addr ${master_addr} \
  --master_port ${master_port} \
"

# Parallelism and batch size
DP=$(($WORLD_SIZE / ${tp} / ${pp} / ${cp}))
GLOBAL_BATCH_SIZE=$(($DP * ${mbs} * ${mbn}))

# Communication overlap settings
COMM_ARGS="--overlap-grad-reduce --overlap-param-gather"
if [[ "${tp}" -gt 1 ]]; then
  COMM_ARGS+=" --tp-comm-overlap"
fi

# Training command
CMD="torchrun $DISTRIBUTED_ARGS \
	pretrain_gpt.py \
	--use-mcore-models \
  $COMM_ARGS \
	--tensor-model-parallel-size ${tp} \
	--pipeline-model-parallel-size ${pp} \
	--context-parallel-size ${cp} \
	--sequence-parallel \
	--num-layers $NUM_LAYERS \
	--hidden-size $HIDDEN_SIZE \
	--ffn-hidden-size $FFN_HIDDEN_SIZE \
	--num-attention-heads $NUM_ATTN_HEADS \
	$GQA_ARGS \
	--hidden-dropout 0.0 \
	--attention-dropout 0 \
	--swiglu \
	--micro-batch-size ${mbs} \
	--global-batch-size $GLOBAL_BATCH_SIZE \
	--seq-length ${seq_len} \
	--max-position-embeddings ${seq_len} \
  --position-embedding-type rope \
  --rotary-base ${ROPE_THETA} \
	--normalization RMSNorm \
  --norm-epsilon ${RMS_NORM_EPS} \
	--train-iters ${train_iters} \
	--lr-decay-iters 3200 \
  $CHECKPOINT_ARGS \
  $DATASET_ARGS \
	$TOKENIZER_ARGS \
	--split 120,4,4 \
	--distributed-backend nccl \
	--lr 0.00015 \
	--lr-decay-style cosine \
	--min-lr 1.0e-5 \
	--weight-decay 1e-2 \
	--clip-grad 1.0 \
	--lr-warmup-fraction .01 \
	--log-interval 1 \
	--log-throughput \
	--save-interval 5000 \
	--eval-interval 5000 \
	--eval-iters 0 \
	--exit-interval 5000 \
	--use-flash-attn \
	--use-distributed-optimizer \
	--bf16 \
  --transformer-impl transformer_engine \
  --cross-entropy-loss-fusion \
  --cross-entropy-fusion-impl native \
	$FP8_ARGS \
	$MOE_ARGS \
	$TIE_ARGS \
  $OTHER_ARGS \
	"

# Log settings
LOG_DIR="/workspace/logs/${MODEL_NAME}-${WORLD_SIZE}GPUS"
LOG_FILE_NAME="megatron_${MODEL_NAME}_\
seq${seq_len}_\
tp${tp}_pp${pp}_cp${cp}_dp${DP}_ep${ep}_etp${etp}_\
mb${mbs}_gb${GLOBAL_BATCH_SIZE}_\
${DT}_\
NODE${node_rank}.log"
LOG_FILE="${LOG_DIR}/${LOG_FILE_NAME}"

# Run command
run_cmd "${CMD}" "${LOG_DIR}" "${LOG_FILE}" "${node_rank}" "${node_num}"
