[autoparallel] add rotor C version

colossalai/fx/passes/algorithms/build_c_ext.py

@@ -0,0 +1,15 @@
+from setuptools import setup, Extension
+import os
+
+this_dir = os.path.dirname(os.path.abspath(__file__))
+ext_modules = [Extension(
+    'dynamic_programs_C_version',
+    sources=[os.path.join(this_dir, 'dynamic_programs.c')],
+)]
+
+setup(
+    name='rotor c extension',
+    version='0.1',
+    description='rotor c extension for faster dp computing',
+    ext_modules=ext_modules,
+)

colossalai/fx/passes/algorithms/ckpt_solver_rotor.py

@@ -5,9 +5,7 @@
 import math
 from .linearize import linearize
 from .operation import ForwardCheck, ForwardEnable, ForwardNograd, Backward, Loss, Chain, Sequence, Function
-from colossalai.fx.passes.meta_info_prop import MetaInfoProp
 from colossalai.fx.codegen.activation_checkpoint_codegen import _find_nested_ckpt_regions
-from colossalai import META_COMPATIBILITY
 
 
 # this is the python compute table code from rotor
@@ -323,34 +321,72 @@ def solver_rotor(gm: ColoGraphModule,
                  mem_limit: int,
                  mem_slots: int = 500,
                  cnode: List[str] = None,
-                 eps: float = 0.0) -> ColoGraphModule:
+                 eps: float = 0.0,
+                 force_python: bool = False) -> ColoGraphModule:
     """solver that automatically find activation checkpoint in rotor's manner
 
     Args:
-        gm (ColoGraphModule): ColoGraphModule generated by tracing model.
+        gm (ColoGraphModule): ColoGraphModule generated by tracing model and MetaInfoProp.
         data (torch.Tensor): input data.
         mem_limit (int): memory budget in Byte.
         mem_slots (int, optional): number of slots for discretizing memory budget. Defaults to 500.
         cnode (List[Node], optional): common node list for linearize. Defaults to None.
         eps (float): epsilon for memory decay. Defaults to 0.0
+        force_python (bool): force to use python version of dynamic programs
 
     Returns:
         ColoGraphModule: annotated ColoGraphModuled with __sequence__ attribute
     """
 
+    # try to import C version solver if force_python is not set
+    if not force_python:
+        try:
+            from .dynamic_programs_C_version import persistent_compute_table
+            CVERSION = True
+
+        # build module if module not found
+        except ModuleNotFoundError:
+            import subprocess
+            import os
+            print("dynamic_programs_C_version hasn't been built! Building library...")
+            this_dir = os.path.dirname(os.path.abspath(__file__))
+            result = subprocess.Popen(
+                f'python {os.path.join(this_dir, "build_c_ext.py")} build_ext --build-lib={this_dir}',
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                shell=True)
+            if result.wait() == 0:
+                print("dynamic_programs_C_version has been built!")
+                from .dynamic_programs_C_version import persistent_compute_table
+                CVERSION = True
+            else:
+                print("dynamic_programs_C_version built failed! Using python version!")
+                CVERSION = False
+
+    # check if metainfoprop is done
+    if any(len(node.meta) == 0 for node in gm.graph.nodes):
+        raise RuntimeError(
+            "Nodes meta information hasn't been prepared! Please run MetaInfoProp before calling solver!")
+
+    # linearize the graph
     node_list = linearize(gm, cnode)
-    mem_unit = mem_limit * (1.0 - eps) // mem_slots
-    if META_COMPATIBILITY:
-        from colossalai.fx.profiler import MetaTensor
-        data = MetaTensor(data, fake_device=next(gm.parameters()).device)
-    MetaInfoProp(gm).run(data)
 
+    # construct chain
+    mem_unit = mem_limit * (1.0 - eps) // mem_slots
     chain: Chain = _construct_chain(node_list, data)
     chain._discretize(mem_unit)
-    opt_table = _compute_table(chain, mem_slots)
+
+    # use C version if possible
+    if CVERSION and not force_python:
+        opt_table = persistent_compute_table(chain, mem_slots)
+    else:
+        opt_table = _compute_table(chain, mem_slots)
+
+    # found sequence
     sequence = _rec(chain, 0, chain.length, mem_slots - chain.cweight[0], opt_table)
     _annotate_from_sequence(sequence, node_list)
 
     # set __sequence__ attribute to GraphModule
     setattr(gm, "__sequence__", sequence)
+    gm.recompile()
     return gm