github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-17 23:18:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autoparallel] Hook all meta information on ResNet nodes for auto activation checkpoint (#2248)

Browse Source 
* [autoparallel] hook node meta on graph nodes for checkpoint solver

* [autoparallel] polish code

* [autoparallel] restore some node handlers

* colossalai/auto_parallel/passes/meta_info_prop.py

* [autoparallel] remove some unused import

* [autoparallel] hook bwd_mem_out
This commit is contained in:
Boyuan Yao
2023-01-02 16:25:18 +08:00
committed by GitHub
parent c8c79102f0
commit ab38aebace
6 changed files with 132 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/auto_parallel/tensor_shard/node_handler/node_handler.py
View File

@@ -138,8 +138,7 @@ class NodeHandler(ABC):
return None
if self.node.op == 'call_module':
submod = self.node.graph.owning_module.get_submodule(self.node.target)
target = type(submod)
target = self.node.graph.owning_module.get_submodule(self.node.target)
elif self.node.op == 'call_function':
target = self.node.target
elif self.node.op == 'call_method':