github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 11:02:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* [fx] added testing for all gpt variants

* polish code

* polish code
This commit is contained in:
Frank Lee
2022-07-06 14:03:13 +08:00
committed by GitHub
parent 189946c5c4
commit 2d13a45a3b
8 changed files with 136 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/fx/tracer/tracer.py
View File

@@ -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