add available

colossalai/autochunk/autochunk_codegen.py

@@ -1,6 +1,11 @@
 from typing import Any, Dict, Iterable, List, Tuple
 
 import torch
+
+import colossalai
+from colossalai.fx.codegen.activation_checkpoint_codegen import CODEGEN_AVAILABLE
+
+if CODEGEN_AVAILABLE:
     from torch.fx.graph import (
         CodeGen,
         PythonCode,
@@ -15,7 +20,6 @@ from torch.fx.graph import (
     )
     from torch.fx.node import Argument, Node, _get_qualified_name, _type_repr, map_arg
 
-import colossalai
 from .search_chunk import SearchChunk
 from .utils import delete_free_var_from_last_use, find_idx_by_name, get_node_shape
 
@@ -302,6 +306,8 @@ def emit_code_with_chunk(
         node_idx += 1
 
 
+if CODEGEN_AVAILABLE:
+
     class AutoChunkCodeGen(CodeGen):
         def __init__(self, meta_graph, max_memory=None, print_mem=False):
             super().__init__()
@@ -349,7 +355,9 @@ class AutoChunkCodeGen(CodeGen):
                 return global_name
 
             # set _custom_builtins here so that we needn't import colossalai in forward
-        _custom_builtins["colossalai"] = _CustomBuiltin("import colossalai", colossalai)
+            _custom_builtins["colossalai"] = _CustomBuiltin(
+                "import colossalai", colossalai
+            )
 
             # Pre-fill the globals table with registered builtins.
             for name, (_, obj) in _custom_builtins.items():
@@ -449,7 +457,9 @@ class AutoChunkCodeGen(CodeGen):
                 )
                 if node.op == "placeholder":
                     assert isinstance(node.target, str)
-                maybe_default_arg = "" if not node.args else f" = {repr(node.args[0])}"
+                    maybe_default_arg = (
+                        "" if not node.args else f" = {repr(node.args[0])}"
+                    )
                     free_vars.append(
                         f"{node.target}{maybe_type_annotation}{maybe_default_arg}"
                     )
@@ -553,7 +563,9 @@ class AutoChunkCodeGen(CodeGen):
 
             if len(wrapped_fns) > 0:
                 wrap_name = add_global("wrap", torch.fx.wrap)
-            wrap_stmts = "\n".join([f'{wrap_name}("{name}")' for name in wrapped_fns])
+                wrap_stmts = "\n".join(
+                    [f'{wrap_name}("{name}")' for name in wrapped_fns]
+                )
             else:
                 wrap_stmts = ""