#!/usr/bin/env bash set_n_least_used_CUDA_VISIBLE_DEVICES() { local n=${1:-"9999"} echo "GPU Memory Usage:" local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv | tail -n +2 | nl -v 0 | tee /dev/tty | sort -g -k 2 | awk '{print $1}' | head -n $n) export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') echo "Now CUDA_VISIBLE_DEVICES is set to:" echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" } set_n_least_used_CUDA_VISIBLE_DEVICES 4 set -xu NUM_RETRY=3 BASE_DIR=$(dirname $(dirname $(realpath $BASH_SOURCE))) EXAMPLES_DIR=$BASE_DIR/examples CONFIG_DIR=$BASE_DIR/config TEMP_DIR=$BASE_DIR/temp TEST_DIR=$BASE_DIR/tests MODEL_SAVE_PATH=$TEMP_DIR/rlhf_models MODELS_DIR=$TEMP_DIR/models_config # Skip those tests due to CI tests timeout MODELS=('llama') ADVANCED_PLUGINS=('sp_split_gather' 'sp_ring' 'sp_all_to_all' 'tp_zero2' '3d' 'gemini' 'gemini_auto' 'zero2' 'zero2_cpu') # pp is still buggy PLUGINS=('3d' 'gemini' 'gemini_auto' 'zero2' 'zero2_cpu') LORA_RANK=('0') # skip to reduce CI execution time, can pass all locally export OMP_NUM_THREADS=8 get_pretrain() { local model=$1 if [[ $model == "llama" ]]; then echo "nickypro/tinyllama-110M" elif [[ $model == "opt" ]]; then echo "facebook/opt-125m" else echo "Unknown model $model" exit 1 fi } get_tokenizer_dirs() { local model=$1 if [[ $model == "llama" ]]; then echo "hf-internal-testing/llama-tokenizer" elif [[ $model == "opt" ]]; then echo "facebook/opt-125m" else echo "Unknown model $model" exit 1 fi } get_conversation_template_config() { local model=$1 if [[ $model == "llama" ]]; then echo "$TEST_DIR/llama.json" elif [[ $model == "opt" ]]; then echo "$TEST_DIR/opt.json" else echo "Unknown model $model" exit 1 fi } random_choice() { local arr=("$@") local len=${#arr[@]} local idx=$((RANDOM % len)) echo ${arr[$idx]} } echo "[Test]: testing sft ..." SKIPPED_TESTS=( llama-3d-20 # 3d plugin doesn't support lora llama-gemini_auto-20 # gemini_auto plugin doesn't support lora llama-gemini-20 # gemini doesn't support lora ) GRAD_CKPTS=('--grad_checkpoint') for lora_rank in ${LORA_RANK[@]}; do for model in ${MODELS[@]}; do for plugin in ${ADVANCED_PLUGINS[@]}; do if [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin-$lora_rank " ]]; then echo "[Test]: Skipped $model-$plugin-$lora_rank" continue elif [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin " ]]; then echo "[Test]: Skipped $model-$plugin" continue fi pretrain=$(get_pretrain $model) tokenizer_dir=$(get_tokenizer_dirs $model) grad_ckpt=$(random_choice "${GRAD_CKPTS[@]}") tp='1' bs='2' pp='1' zero_stage='0' sp='1' sp_mode='split_gather' enable_sequence_parallelism='' if [[ $plugin == "3d" ]]; then tp='4' bs='8' fi if [[ $plugin == "tp_zero2" ]]; then tp='4' bs='8' zero_stage='2' plugin='3d' fi if [[ $plugin == "tp_pp" ]]; then tp='2' bs='8' pp='2' plugin='3d' fi if [[ $plugin == "pp" ]]; then bs='8' pp='4' plugin='3d' fi if [[ $plugin == "sp_split_gather" ]]; then enable_sequence_parallelism='--enable_sequence_parallelism' sp_mode='split_gather' tp='4' sp='1' bs='8' plugin='3d' fi if [[ $plugin == "sp_ring" ]]; then enable_sequence_parallelism='--enable_sequence_parallelism' sp_mode='ring' tp='4' sp='1' bs='8' plugin='3d' fi if [[ $plugin == "sp_all_to_all" ]]; then enable_sequence_parallelism='--enable_sequence_parallelism' sp_mode='all_to_all' tp='1' sp='4' bs='8' plugin='3d' fi grad_accu='2' # Check if the plugin is either "gemini_auto" or "gemini" and set grad_accu to '1' if [[ $plugin == "gemini_auto" ]]; then grad_accu='1' fi for i in $(seq $NUM_RETRY); do echo "[Test]: $model-$plugin-$lora_rank, attempt $i" declare -a dataset=() for split in $(seq -f "%05g" 0 0); do dataset+=("$TEMP_DIR/rlhf_data/tokenized_${model}_sft/arrow/part-$split") done colossalai run --nproc_per_node 4 --master_port 31332 $EXAMPLES_DIR/training_scripts/train_sft.py \ --pretrain $pretrain \ --tokenizer_dir $tokenizer_dir \ --dataset ${dataset[@]} \ --eval_dataset ${dataset[@]} \ --save_path $MODEL_SAVE_PATH \ --config_file $MODELS_DIR/config.jsonl \ --lora_rank $lora_rank \ --plugin $plugin \ --batch_size $bs \ --max_epochs 1 \ --accumulation_steps $grad_accu \ --tp $tp \ --pp $pp \ --zero_stage $zero_stage \ --sp $sp \ --sp_mode $sp_mode \ $enable_sequence_parallelism \ --lr 2e-5 \ $grad_ckpt \ --max_len 400 \ --use_flash_attn passed=$? if [ $passed -eq 0 ]; then rm -rf ${MODEL_SAVE_PATH:?}/* rm -rf ${MODELS_DIR:?}/* break fi done if [ $passed -ne 0 ]; then echo "[Test]: Failed $model-$plugin-$lora_rank" exit 1 fi done done done echo "[Test]: testing reward model ..." SKIPPED_TESTS=( llama-3d-20 # 3d plugin doesn't support lora llama-gemini_auto-20 # gemini_auto plugin doesn't support lora llama-gemini-20 # gemini doesn't support lora ) GRAD_CKPTS=('--grad_checkpoint') for lora_rank in ${LORA_RANK[@]}; do for model in ${MODELS[@]}; do for plugin in ${PLUGINS[@]}; do if [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin-$lora_rank " ]]; then echo "[Test]: Skipped $model-$plugin-$lora_rank" continue elif [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin " ]]; then echo "[Test]: Skipped $model-$plugin" continue fi pretrain=$(get_pretrain $model) tokenizer_dir=$(get_tokenizer_dirs $model) grad_ckpt=$(random_choice "${GRAD_CKPTS[@]}") tp='1' bs='2' if [[ $plugin == "3d" ]]; then tp='4' bs='8' fi grad_accu='2' # gemini_auto and gemini doesn't support gradient accumulation if [[ $plugin == "gemini_auto" ]]; then grad_accu='1' fi for i in $(seq $NUM_RETRY); do echo "[Test]: $model-$plugin-$lora_rank, attempt $i" declare -a dataset=() for split in $(seq -f "%05g" 0 0); do dataset+=("$TEMP_DIR/rlhf_data/tokenized_${model}_preference/arrow/part-$split") done colossalai run --nproc_per_node 4 --master_port 31332 $EXAMPLES_DIR/training_scripts/train_rm.py \ --pretrain $pretrain \ --tokenizer_dir $tokenizer_dir \ --dataset ${dataset[@]} \ --eval_dataset ${dataset[@]} \ --save_dir $MODEL_SAVE_PATH \ --config_file $MODELS_DIR/config.jsonl \ --lora_rank $lora_rank \ --plugin $plugin \ --batch_size $bs \ --max_epochs 1 \ --accumulation_steps $grad_accu \ --tp $tp \ --lr 2e-5 \ $grad_ckpt \ --max_len 400 \ --use_flash_attn passed=$? if [ $passed -eq 0 ]; then rm -rf ${MODEL_SAVE_PATH:?}/* rm -rf ${MODELS_DIR:?}/* break fi done if [ $passed -ne 0 ]; then echo "[Test]: Failed $model-$plugin-$lora_rank" exit 1 fi done done done echo "[Test]: testing ppo ..." SKIPPED_TESTS=( llama-3d # 3d plugin doesn't support lora llama-gemini # gemini doesn't support lora ) GRAD_CKPTS=('--grad_checkpoint') for lora_rank in ${LORA_RANK[@]}; do for model in ${MODELS[@]}; do for plugin in ${PLUGINS[@]}; do if [[ $plugin == "gemini_auto" ]]; then echo "[Test]: Skipped $model-$plugin" continue # gemini_auto plugin doesn't support generation fi if [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin-$lora_rank " ]]; then echo "[Test]: Skipped $model-$plugin-$lora_rank" continue elif [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin " ]]; then echo "[Test]: Skipped $model-$plugin" continue fi pretrain=$(get_pretrain $model) tokenizer_dir=$(get_tokenizer_dirs $model) grad_ckpt=$(random_choice "${GRAD_CKPTS[@]}") tp='1' bs='4' ebs='8' conversation_template=$(get_conversation_template_config $model) if [[ $plugin == "3d" ]]; then tp='4' bs='16' ebs='32' fi grad_accu='2' # gemini_auto and gemini doesn't support gradient accumulation if [[ $plugin == "gemini_auto" ]]; then grad_accu='1' fi # gemini_auto and gemini doesn't support generation if [[ $plugin == "gemini_auto" ]]; then # gemini-auto doesn't support generation echo "[Test]: Skipped $model-$plugin" continue fi for i in $(seq $NUM_RETRY); do echo "[Test]: $model-$plugin-$lora_rank, attempt $i" declare -a prompt_dataset=() for split in $(seq -f "%05g" 0 0); do prompt_dataset+=("$TEMP_DIR/rlhf_data/tokenized_${model}_prompt/arrow/part-$split") done declare -a ptx_dataset=() for split in $(seq -f "%05g" 0 0); do ptx_dataset+=("$TEMP_DIR/rlhf_data/tokenized_${model}_sft/arrow/part-$split") done colossalai run --nproc_per_node 4 --master_port 31332 $EXAMPLES_DIR/training_scripts/train_ppo.py \ --pretrain $pretrain \ --rm_pretrain $pretrain \ --tokenizer_dir $tokenizer_dir \ --conversation_template_config $conversation_template \ --prompt_dataset ${prompt_dataset[@]} \ --ptx_dataset ${ptx_dataset[@]} \ --ptx_batch_size 1 \ --ptx_coef 0.2 \ --save_path $MODEL_SAVE_PATH \ --lora_rank $lora_rank \ --plugin $plugin \ --num_episodes 5 \ --num_collect_steps 1 \ --num_update_steps 1 \ --experience_batch_size $ebs \ --train_batch_size $bs \ --accumulation_steps $grad_accu \ --lr 9e-6 \ --mixed_precision "bf16" \ --grad_clip 1.0 \ --tp $tp \ --lr 2e-5 \ $grad_ckpt \ --max_len 400 \ --max_seq_len 10 \ # --use_flash_attn passed=$? if [ $passed -eq 0 ]; then rm -rf ${MODEL_SAVE_PATH:?}/* rm -rf ${MODELS_DIR:?}/* break fi done if [ $passed -ne 0 ]; then echo "[Test]: Failed $model-$plugin-$lora_rank" exit 1 fi done done done echo "[Test]: testing DPO ..." SKIPPED_TESTS=( llama-3d-20 # 3d plugin doesn't support lora llama-gemini_auto-20 # gemini_auto plugin doesn't support lora llama-gemini-20 # gemini doesn't support lora ) GRAD_CKPTS=('--grad_checkpoint') for lora_rank in ${LORA_RANK[@]}; do for model in ${MODELS[@]}; do for plugin in ${PLUGINS[@]}; do if [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin-$lora_rank " ]]; then echo "[Test]: Skipped $model-$plugin-$lora_rank" continue elif [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin " ]]; then echo "[Test]: Skipped $model-$plugin" continue fi pretrain=$(get_pretrain $model) tokenizer_dir=$(get_tokenizer_dirs $model) grad_ckpt=$(random_choice "${GRAD_CKPTS[@]}") tp='1' bs='2' if [[ $plugin == "3d" ]]; then tp='4' bs='8' fi grad_accu='2' # gemini_auto and gemini doesn't support gradient accumulation if [[ $plugin == "gemini_auto" ]]; then grad_accu='1' fi # gemini_auto doesn't support generation # (need to calculate ref_model logits through forwarding in inference mode) if [[ $plugin == "gemini_auto" ]]; then echo "[Test]: Skipped $model-$plugin" continue fi for i in $(seq $NUM_RETRY); do echo "[Test]: $model-$plugin-$lora_rank, attempt $i" declare -a dataset=() for split in $(seq -f "%05g" 0 0); do dataset+=("$TEMP_DIR/rlhf_data/tokenized_${model}_preference/arrow/part-$split") done colossalai run --nproc_per_node 4 --master_port 31332 $EXAMPLES_DIR/training_scripts/train_dpo.py \ --pretrain $pretrain \ --tokenizer_dir $tokenizer_dir \ --dataset ${dataset[@]} \ --eval_dataset ${dataset[@]} \ --save_dir $MODEL_SAVE_PATH \ --config_file $MODELS_DIR/config.jsonl \ --lora_rank $lora_rank \ --plugin $plugin \ --batch_size $bs \ --max_epochs 1 \ --accumulation_steps $grad_accu \ --tp $tp \ --lr 2e-5 \ $grad_ckpt \ --max_len 400 \ --use_flash_attn passed=$? if [ $passed -eq 0 ]; then rm -rf ${MODEL_SAVE_PATH:?}/* rm -rf ${MODELS_DIR:?}/* break fi done if [ $passed -ne 0 ]; then echo "[Test]: Failed $model-$plugin-$lora_rank" exit 1 fi done done done echo "[Test]: testing ORPO ..." SKIPPED_TESTS=( llama-3d-20 # 3d plugin doesn't support lora llama-gemini_auto-20 # gemini_auto plugin doesn't support lora llama-gemini-20 # gemini doesn't support lora ) GRAD_CKPTS=('--grad_checkpoint') for lora_rank in ${LORA_RANK[@]}; do for model in ${MODELS[@]}; do for plugin in ${PLUGINS[@]}; do if [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin-$lora_rank " ]]; then echo "[Test]: Skipped $model-$plugin-$lora_rank" continue elif [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin " ]]; then echo "[Test]: Skipped $model-$plugin" continue fi pretrain=$(get_pretrain $model) tokenizer_dir=$(get_tokenizer_dirs $model) grad_ckpt=$(random_choice "${GRAD_CKPTS[@]}") tp='1' bs='2' if [[ $plugin == "3d" ]]; then tp='4' bs='8' fi grad_accu='2' # gemini_auto and gemini doesn't support gradient accumulation if [[ $plugin == "gemini_auto" ]]; then grad_accu='1' fi # gemini_auto doesn't support generation # (need to calculate ref_model logits through forwarding in inference mode) if [[ $plugin == "gemini_auto" ]]; then echo "[Test]: Skipped $model-$plugin" continue fi for i in $(seq $NUM_RETRY); do echo "[Test]: $model-$plugin-$lora_rank, attempt $i" declare -a dataset=() for split in $(seq -f "%05g" 0 0); do dataset+=("$TEMP_DIR/rlhf_data/tokenized_${model}_preference/arrow/part-$split") done colossalai run --nproc_per_node 4 --master_port 31332 $EXAMPLES_DIR/training_scripts/train_orpo.py \ --pretrain $pretrain \ --tokenizer_dir $tokenizer_dir \ --dataset ${dataset[@]} \ --eval_dataset ${dataset[@]} \ --save_dir $MODEL_SAVE_PATH \ --config_file $MODELS_DIR/config.jsonl \ --lora_rank $lora_rank \ --plugin $plugin \ --batch_size $bs \ --max_epochs 1 \ --accumulation_steps $grad_accu \ --tp $tp \ --lr 2e-5 \ $grad_ckpt \ --max_len 400 \ --use_flash_attn passed=$? if [ $passed -eq 0 ]; then rm -rf ${MODEL_SAVE_PATH:?}/* rm -rf ${MODELS_DIR:?}/* break fi done if [ $passed -ne 0 ]; then echo "[Test]: Failed $model-$plugin-$lora_rank" exit 1 fi done done done echo "[Test]: testing KTO ..." SKIPPED_TESTS=( llama-3d-20 # 3d plugin doesn't support lora llama-gemini_auto-20 # gemini_auto plugin doesn't support lora llama-gemini-20 # gemini doesn't support lora ) GRAD_CKPTS=('--grad_checkpoint') for lora_rank in ${LORA_RANK[@]}; do for model in ${MODELS[@]}; do for plugin in ${PLUGINS[@]}; do if [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin-$lora_rank " ]]; then echo "[Test]: Skipped $model-$plugin-$lora_rank" continue elif [[ " ${SKIPPED_TESTS[*]} " =~ " $model-$plugin " ]]; then echo "[Test]: Skipped $model-$plugin" continue fi pretrain=$(get_pretrain $model) tokenizer_dir=$(get_tokenizer_dirs $model) grad_ckpt=$(random_choice "${GRAD_CKPTS[@]}") tp='1' bs='2' if [[ $plugin == "3d" ]]; then tp='4' bs='8' fi grad_accu='2' # gemini_auto and gemini doesn't support gradient accumulation if [[ $plugin == "gemini_auto" ]]; then grad_accu='1' fi # gemini_auto doesn't support generation # (need to calculate ref_model logits through forwarding in inference mode) if [[ $plugin == "gemini_auto" ]]; then echo "[Test]: Skipped $model-$plugin" continue fi for i in $(seq $NUM_RETRY); do echo "[Test]: $model-$plugin-$lora_rank, attempt $i" declare -a dataset=() for split in $(seq -f "%05g" 0 0); do dataset+=("$TEMP_DIR/rlhf_data/tokenized_${model}_kto/arrow/part-$split") done colossalai run --nproc_per_node 4 --master_port 31332 $EXAMPLES_DIR/training_scripts/train_kto.py \ --pretrain $pretrain \ --tokenizer_dir $tokenizer_dir \ --dataset ${dataset[@]} \ --eval_dataset ${dataset[@]} \ --save_dir $MODEL_SAVE_PATH \ --config_file $MODELS_DIR/config.jsonl \ --lora_rank $lora_rank \ --plugin $plugin \ --batch_size $bs \ --max_epochs 1 \ --accumulation_steps $grad_accu \ --tp $tp \ --lr 2e-5 \ --desirable_weight 1.2 \ $grad_ckpt \ --max_len 400 \ --use_flash_attn passed=$? if [ $passed -eq 0 ]; then rm -rf ${MODEL_SAVE_PATH:?}/* rm -rf ${MODELS_DIR:?}/* break fi done if [ $passed -ne 0 ]; then echo "[Test]: Failed $model-$plugin-$lora_rank" exit 1 fi done done done