mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
9 lines
414 B
9 lines
414 B
2 years ago
|
export BATCH_SIZE=${BATCH_SIZE:-64}
|
||
|
export MODEL_TYPE=${MODEL_TYPE:-"gpt2_medium"}
|
||
|
export MEMORY_BUDGET=${MEMORY_BUDGET:-16}
|
||
|
export SOLVER_TYPE=${SOLVER_TYPE:-"asyn"}
|
||
|
|
||
|
mkdir -p offload_logs
|
||
|
|
||
|
python train_gpt_offload.py --model_type=${MODEL_TYPE} --memory_budget=${MEMORY_BUDGET} --solver_type=${SOLVER_TYPE} --batch_size=${BATCH_SIZE} 2>&1 | tee ./offload_logs/${MODEL_TYPE}_bs_${BATCH_SIZE}_st_${SOLVER_TYPE}.log
|