[fx] added testing for all gpt variants (#1210)

* [fx] added testing for all gpt variants

* polish code

* polish code

colossalai/fx/proxy.py

@@ -1,8 +1,7 @@
 import operator
 import torch
 from torch.fx.proxy import Proxy, Attribute
-from typing import List, Union
-from torch.utils._pytree import PyTree
+from typing import List, Union, Any
 
 __all__ = ['ColoProxy']
 
@@ -14,34 +13,33 @@ class ColoProxy(Proxy):
 
     Usage:
         proxy = tracer.create_proxy(...)
-        proxy.meta_tensor = torch.empty(4, 2, device='meta')
+        proxy.meta_data = torch.empty(4, 2, device='meta')
         print(len(proxy)) # expect output 4
 
     """
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._meta_tensor = None
+        self._meta_data = None
 
     @property
-    def meta_tensor(self):
-        return self._meta_tensor
+    def meta_data(self):
+        return self._meta_data
 
-    @meta_tensor.setter
-    def meta_tensor(self, tensor: Union[List[torch.Tensor], torch.Tensor]):
-
-        def _is_meta(item):
-            assert torch.is_tensor(item) and item.is_meta
-
-        torch.fx.node.map_aggregate(tensor, _is_meta)
-        self._meta_tensor = tensor
+    @meta_data.setter
+    def meta_data(self, data: Any):
+        self._meta_data = data
 
     @property
-    def has_meta_tensor(self):
-        return self.meta_tensor is not None
+    def has_meta_data(self):
+        return self._meta_data is not None
 
-    def _assert_has_meta(self):
-        assert self.has_meta_tensor, f'Meta tensor is not set for {self.node.name}'
+    def _assert_meta_data_is_tensor(self):
+        assert torch.is_tensor(
+            self._meta_data) and self._meta_data.is_meta, f'Meta data is not a meta tensor for {self.node.name}'
+
+    def _assert_has_meta_data(self):
+        assert self._meta_data, f'Meta data is not set for {self.node.name}'
 
     @property
     def device(self):
@@ -51,37 +49,37 @@ class ColoProxy(Proxy):
 
     @property
     def dtype(self):
-        self._assert_has_meta()
-        return self.meta_tensor.dtype
+        self._assert_meta_data_is_tensor()
+        return self.meta_data.dtype
 
     @property
     def shape(self):
-        self._assert_has_meta()
-        return self.meta_tensor.shape
+        self._assert_meta_data_is_tensor()
+        return self.meta_data.shape
 
     def dim(self):
-        self._assert_has_meta()
-        return self.meta_tensor.dim()
+        self._assert_meta_data_is_tensor()
+        return self.meta_data.dim()
 
     def size(self, dim: int = None):
-        self._assert_has_meta()
+        self._assert_meta_data_is_tensor()
         if dim:
-            return self.meta_tensor.size(dim=dim)
+            return self.meta_data.size(dim=dim)
         else:
             # size(dim=None) will trigger runtime error for meta tensor
-            return self.meta_tensor.size()
+            return self.meta_data.size()
 
     def __len__(self):
-        self._assert_has_meta()
-        return len(self.meta_tensor)
+        self._assert_has_meta_data()
+        return len(self.meta_data)
 
     def __bool__(self):
-        self._assert_has_meta()
-        return self.meta_tensor
+        self._assert_has_meta_data()
+        return self.meta_data
 
     def __getattr__(self, k):
-        if k == "metadata":
-            return self.meta_tensor
+        if k == "meta_data":
+            return self.__getattribute__(k)
         # note: not added to the graph yet, if this is a method call
         # we peephole optimize to the method invocation
         return Attribute(self, k)

colossalai/fx/tracer/_tracer_utils.py

@@ -22,8 +22,8 @@ def extract_meta(*args, **kwargs):
         if isinstance(val, MetaDeviceAttribute):
             return 'meta'
         elif isinstance(val, ColoProxy):
-            assert val.meta_tensor is not None
-            return val.meta_tensor
+            assert val.meta_data is not None
+            return val.meta_data
         return val
 
     new_args = [_convert(val) for val in args]

colossalai/fx/tracer/meta_patch/patched_function.py

@@ -60,3 +60,32 @@ def torch_matmul(input, other, *, out=None):
     if shape is None:
         return torch.tensor(0.0, device="meta")
     return torch.empty(*shape, device="meta")
+
+
+@meta_patched_function.register(torch.arange)
+def torch_arange(*args, **kwargs):
+    n = len(args)
+    step = 1
+    if n == 1:
+        start = 0
+        end = args[0]
+    elif n == 2:
+        start, end = args
+    else:
+        start, end, step = args
+    if isinstance(start, float):
+        start = int(start)
+    if isinstance(end, float):
+        start = int(end)
+    if isinstance(step, float):
+        step = int(step)
+    step = kwargs.get("step", step)
+    dtype = kwargs.get("dtype")
+    return torch.empty((end - start) // step, dtype=dtype, device="meta")
+
+
+@meta_patched_function.register(torch.where)
+def torch_where(condition, x, y):
+    # torch.where returns the broadcasted tensor of condition, x, and y,
+    # so hack it by using addition
+    return condition.to(device="meta") + x.to(device="meta") + y.to(device="meta")

colossalai/fx/tracer/tracer.py

@@ -14,7 +14,6 @@ from torch import Tensor
 from torch.fx import Tracer
 from torch.fx.graph import Graph
 from torch.fx.proxy import Proxy, ParameterProxy
-from torch.utils import _pytree
 from ..proxy import ColoProxy
 from typing import Optional, Dict, Any
 from ._tracer_utils import is_element_in_list, extract_meta
@@ -62,7 +61,7 @@ class ColoTracer(Tracer):
         proxy: ColoProxy
 
         if kind == "placeholder" and target in self.meta_args and self.meta_args[target].is_meta:
-            proxy.meta_tensor = self.meta_args[target]
+            proxy.meta_data = self.meta_args[target]
             return proxy
 
         if target in self.orig_torch_tensor_methods:
@@ -128,7 +127,7 @@ class ColoTracer(Tracer):
 
             if not isinstance(proxy, Proxy):
                 raise ValueError("Don't support composite output yet")
-            proxy.meta_tensor = meta_out
+            proxy.meta_data = meta_out
         except Exception as e:
             raise RuntimeError(f"Could not compute metadata for {kind} target {target}: {e}")
         return proxy