github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-05-10 01:17:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' of https://github.com/hpcaitech/ColossalAI into dev0116

Browse Source
This commit is contained in:
jiaruifang 2023-01-18 18:43:11 +08:00
commit 7f822a5c45
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/fx/graph_module.py
View File

@ -22,8 +22,13 @@ if COLOGM:
class ColoGraphModule(GraphModule): class ColoGraphModule(GraphModule):
def __init__(self, root: Union[torch.nn.Module, Dict[str, Any]], graph: Graph, class_name: str = 'GraphModule'): def __init__(self,
graph.set_codegen(ActivationCheckpointCodeGen()) root: Union[torch.nn.Module, Dict[str, Any]],
graph: Graph,
class_name: str = 'GraphModule',
ckpt_codegen: bool = True):
if ckpt_codegen:
graph.set_codegen(ActivationCheckpointCodeGen())
super().__init__(root, graph, class_name) super().__init__(root, graph, class_name)
def bind(self, ckpt_def, globals): def bind(self, ckpt_def, globals):