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

[autoparallel] modify comm nodes' memory cost in construct chain (#2263)

Browse Source 
* [autoparallel] align the data_ptr with the old version of auto activation checkpoint pipeline

* [autoparallel] using fwd_time and bwd_time instead of fwd_flop and bwd_flop

* [autoparallel] specifycomm nodes' memory cost in construct chain
This commit is contained in:
Boyuan Yao
2023-01-03 11:38:48 +08:00
committed by GitHub
parent 1ea99b869e
commit 5c2ef9fc76
3 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/auto_parallel/passes/meta_info_prop.py
View File

@@ -151,6 +151,7 @@ class MetaInfoProp:
# fetch other memory informations
memory_cost = meta_info.memory_cost
graph_info.fwd_mem_tmp = memory_cost.fwd.temp
graph_info.fwd_mem_out = memory_cost.fwd.activation
graph_info.bwd_mem_tmp = memory_cost.bwd.temp
graph_info.bwd_mem_out = memory_cost.bwd.activation