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

colossalai/fx/graph_module.py

@@ -22,7 +22,12 @@ 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,
+                     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)