mirror of https://github.com/hpcaitech/ColossalAI
9 lines
414 B
Bash
9 lines
414 B
Bash
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
|