set -x

GPUS=${GPUS:-128}
GPUS_PER_NODE=${GPUS_PER_NODE:-8}
QUOTA_TYPE=${QUOTA_TYPE:-"reserved"}
NODES=$((GPUS / GPUS_PER_NODE))
CPUS_PER_TASK=${CPUS_PER_TASK:-12}
SRUN_ARGS=${SRUN_ARGS:-""}

export PYTHONPATH="${PYTHONPATH}:$(pwd)"
export MASTER_PORT=34229
export TF_CPP_MIN_LOG_LEVEL=3

OUTPUT_DIR="./shell/playground/train_states/SDLM_3B_D8"
BASE_MODEL="./shell/playground/ckpt/Qwen2.5-3B"
META_JSON="./shell/playground/data/meta/sft/sft_opc436k_scale_math_1m_smoltalk_1m_tulu_1m.json"


if [ ! -d "$OUTPUT_DIR" ]; then
  mkdir -p "$OUTPUT_DIR"
fi

# number of gpus: 128
# batch size per gpu: 2
# gradient accumulation steps: 1
# total batch size: ~256
# num of examples:
# epoch: 1
srun --gres=gpu:${GPUS_PER_NODE} \
  --nodes=${NODES} \
  --ntasks=${GPUS} \
  --ntasks-per-node=${GPUS_PER_NODE} \
  --cpus-per-task=${CPUS_PER_TASK} \
  --kill-on-bad-exit=1 \
  --quotatype=${QUOTA_TYPE} \
  ${SRUN_ARGS} \
  python -u sdlm/train/sdlm_train.py \
  --model_name_or_path ${BASE_MODEL} \
  --conv_style "Qwen-2-5" \
  --use_fast_tokenizer False \
  --output_dir ${OUTPUT_DIR} \
  --meta_path ${META_JSON} \
  --overwrite_output_dir True \
  --dataloader_num_workers 16 \
  --bf16 True \
  --num_train_epochs 1 \
  --per_device_train_batch_size 2 \
  --gradient_accumulation_steps 1 \
  --evaluation_strategy "no" \
  --save_strategy "steps" \
  --save_steps 800 \
  --save_total_limit 10 \
  --learning_rate 5e-6 \
  --weight_decay 0.01 \
  --warmup_ratio 0.03 \
  --lr_scheduler_type "constant" \
  --logging_steps 1 \
  --max_seq_length 5632 \
  --do_train True \
  --grad_checkpoint True \
  --deepspeed "zero_stage1_config.json" \
  --report_to "tensorboard" \
  --block_size 8 \
  --causal_attn False \
   --attn_implementation sdpa \
  2>&1 | tee -a "${OUTPUT_DIR}/training_log.txt"
