ds_pretrain.sh

LR=1e-5

MASTER_PORT=$(shuf -n 1 -i 10000-65535)

HIP_VISIBLE_DEVICES=0,1,2,3 deepspeed --num_gpus=4 --master_port $MASTER_PORT main.py \
    --deepspeed deepspeed.json \
    --do_train \
    --train_file The-Lord-of-the-Rings-1.json \
    --prompt_column prompt \
    --response_column response \
    --overwrite_cache \
    --model_name_or_path THUDM/chatglm-6b \
    --output_dir ./output/pretrain \
    --overwrite_output_dir \
    --max_source_length 8 \
    --max_target_length 128 \
    --per_device_train_batch_size 16 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --predict_with_generate \
    --max_steps 2000 \
    --logging_steps 5 \
    --save_steps 1000 \
    --learning_rate $LR \
    --fp16