[fx] supported data-dependent control flow in model tracing

colossalai/fx/__init__.py

@@ -0,0 +1 @@
+from .tracer import ColoTracer

colossalai/fx/proxy.py

@@ -37,6 +37,12 @@ def has_meta_tensor(self):
     def _assert_has_meta(self):
         assert self.has_meta_tensor, f'Meta tensor is not set for {self.node.name}'
 
+    @property
+    def device(self):
+        # Hack so we can track when devices are used. During meta-tensor propagation,
+        # replace these values with a constant 'meta'
+        return MetaDeviceAttribute(self, "device")
+
     @property
     def dtype(self):
         self._assert_has_meta()
@@ -72,3 +78,27 @@ def __getattr__(self, k):
 
     def __setitem__(self, indices, values):
         return self.tracer.create_proxy("call_function", operator.setitem, (self, indices, values), {})
+
+
+class ColoAttribute(ColoProxy):
+
+    def __init__(self, root, attr: str):
+        # this class is copied from torch.fx.Attribute
+        # but inherits ColoProxy
+        self.root = root
+        self.attr = attr
+        self.tracer = root.tracer
+        self._node = None
+
+    @property
+    def node(self):
+        if self._node is None:
+            self._node = self.tracer.create_proxy("call_function", getattr, (self.root, self.attr), {}).node
+        return self._node
+
+    def __call__(self, *args, **kwargs):
+        return self.tracer.create_proxy("call_method", self.attr, (self.root,) + args, kwargs)
+
+
+class MetaDeviceAttribute(ColoAttribute):
+    pass

colossalai/fx/tracer/__init__.py

@@ -0,0 +1 @@
+from .tracer import ColoTracer

colossalai/fx/tracer/_tracer_utils.py

@@ -0,0 +1,31 @@
+from typing import List, Union, Any
+from ..proxy import ColoProxy, MetaDeviceAttribute
+
+__all__ = ['is_element_in_list', 'extract_meta']
+
+
+def is_element_in_list(elements: Union[List[Any], Any], list_: List[Any]):
+    if isinstance(elements, (tuple, list, set)):
+        for ele in elements:
+            if ele not in list_:
+                return False, ele
+    else:
+        if elements not in list_:
+            return False, elements
+
+    return True, None
+
+
+def extract_meta(*args, **kwargs):
+
+    def _convert(val):
+        if isinstance(val, MetaDeviceAttribute):
+            return 'meta'
+        elif isinstance(val, ColoProxy):
+            assert val.meta_tensor is not None
+            return val.meta_tensor
+        return val
+
+    new_args = [_convert(val) for val in args]
+    new_kwargs = {k: _convert(v) for k, v in kwargs.items()}
+    return new_args, new_kwargs

colossalai/fx/tracer/meta_patch/__init__.py

@@ -0,0 +1,4 @@
+from sys import meta_path
+from .registry import *
+from .patched_function import *
+from .patched_module import *

colossalai/fx/tracer/meta_patch/patched_module.py

@@ -0,0 +1,7 @@
+import torch
+from .registry import meta_patched_module
+
+
+@meta_patched_module.register(torch.nn.Linear)
+def torch_nn_linear(self, input):
+    return torch.empty(input.shape[:-1] + (self.out_features,), device="meta")

colossalai/fx/tracer/meta_patch/registry.py

@@ -0,0 +1,25 @@
+class PatchRegistry:
+
+    def __init__(self, name):
+        self.name = name
+        self.store = {}
+
+    def register(self, source):
+
+        def wrapper(func):
+            self.store[source] = func
+            return func
+
+        return wrapper
+
+    def get(self, source):
+        assert source in self.store
+        target = self.store[source]
+        return target
+
+    def has(self, source):
+        return source in self.store
+
+
+meta_patched_function = PatchRegistry(name='patched_functions_for_meta_execution')
+meta_patched_module = PatchRegistry(name='patched_modules_for_meta_execution')