ColossalAI/examples/language/opt/run_demo.sh

45 lines
822 B
Bash

set -xe
pip install -r requirements.txt
# model name or path
MODEL="facebook/opt-350m"
# path for saving model
OUTPUT_PATH="./output_model.bin"
# plugin(training strategy)
# can only be one of "torch_ddp"/"torch_ddp_fp16"/"low_level_zero"/"gemini"
PLUGIN="hybrid_parallel"
# number of gpus to use
GPUNUM=4
# batch size per gpu
BS=16
# learning rate
LR="5e-5"
# number of epoch
EPOCH=10
# weight decay
WEIGHT_DECAY=0.01
# ratio of warmup steps
WARMUP_RATIO=0.1
# run the script for demo
colossalai run \
--nproc_per_node ${GPUNUM} \
--master_port 29505 \
opt_train_demo.py \
--model_name_or_path ${MODEL} \
--output_path ${OUTPUT_PATH} \
--plugin ${PLUGIN} \
--batch_size ${BS} \
--num_epoch ${EPOCH} \
--learning_rate ${LR} \
--weight_decay ${WEIGHT_DECAY} \
--warmup_ratio ${WARMUP_RATIO}