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

* [fx] added testing for all gpt variants

* polish code

* polish code

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