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

7f822a5c · jiaruifang · 025b482d · 5db3a5bf · 7f822a5c
Commit 7f822a5c authored Jan 18, 2023 by jiaruifang
Hide whitespace changes
Inline Side-by-side

Showing with 7 additions and 2 deletions

colossalai/fx/graph_module.py colossalai/fx/graph_module.py +7 -2

No files found.
--- a/colossalai/fx/graph_module.py
+++ b/colossalai/fx/graph_module.py
@@ -22,8 +22,13 @@ if COLOGM:
    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)
        def bind(self, ckpt_def, globals):