github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-10 13:30:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

[auto-parallel] add auto-offload feature (#3154)

Browse Source 
* add auto-offload feature

* polish code

* fix syn offload runtime pass bug

* add offload example

* fix offload testing bug

* fix example testing bug
This commit is contained in:
Zihao
2023-03-21 14:17:41 +08:00
committed by GitHub
parent 258b43317c
commit 18dbe76cae
18 changed files with 2833 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
examples/language/gpt/experiments/auto_offload/run.sh Normal file
View File

@@ -0,0 +1,8 @@
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