fix: torch.compile compatibility

Author: ZoroZhao

traincheck/config/config.py

@@ -249,3 +249,12 @@ def should_disable_proxy_dumping() -> bool:
     "preprocessing",
     "postprocessing",
 }
+
+COMPILE_INTERNAL_MODULE = (
+    "torch.fx",
+    # "torch._dynamo",
+    "torch._inductor",
+    "torch._subclasses",
+    "torch._higher_order_ops",
+    "torch.utils._sympy",
+)

traincheck/instrumentor/dumper.py

@@ -18,13 +18,14 @@
 
 # if torch.cuda.is_available():
 from traincheck.proxy_wrapper.hash import tensor_hash
+from traincheck.proxy_wrapper.proxy_basics import is_fake_tensor
 from traincheck.proxy_wrapper.proxy_config import (
     attribute_black_list,
     primitive_types,
     proxy_attribute,
     tensor_dump_format,
 )
-from traincheck.utils import get_timestamp_ns, typename
+from traincheck.utils import get_timestamp_ns, typename, typename_compile
 
 DEBUG = os.environ.get("ML_DAIKON_DEBUG", False)
 THREAD_DATA = threading.local()
@@ -45,12 +46,48 @@
 logger = logging.getLogger(__name__)
 
 
+def _json_default(o):
+    try:
+        if type(o).__name__ in ("SymInt", "SymFloat", "SymBool"):
+            return str(o)
+
+        if isinstance(o, torch.device):
+            return str(o)
+        if isinstance(o, torch.dtype):
+            return str(o)
+        if isinstance(o, torch.Size):
+            out = []
+            for d in o:
+                try:
+                    out.append(int(d))
+                except Exception:
+                    out.append(str(d))
+            return out
+    except Exception:
+        pass
+
+    if isinstance(o, set):
+        return list(o)
+    if isinstance(o, tuple):
+        return list(o)
+
+    try:
+        import numpy as np
+
+        if isinstance(o, (np.generic,)):
+            return o.item()
+    except Exception:
+        pass
+
+    return repr(o)
+
+
 def serialize(obj_dict: dict[str, object | str]) -> str:
     try:
-        return orjson.dumps(obj_dict).decode("utf-8")
+        return orjson.dumps(obj_dict, default=_json_default).decode("utf-8")
     except Exception:
         # if orjson fails (e.g. cannot handle ints larger than 64-bit), fallback to json
-        return json.dumps(obj_dict)
+        return json.dumps(obj_dict, default=_json_default)
 
 
 def monitor_main_thread(main_thread, stop_event):
@@ -350,12 +387,17 @@ def convert_var_to_dict(var, include_tensor_data=True, dump_config=None) -> dict
 
         attr = safe_getattr(var, attr_name)
         if attr is NOT_FOUND:
-            logger.warning(
-                f"Failed to get attribute {attr_name} of object type {type(var)}, skipping it for all following dumps for this attribute."
-            )
-            if var_type not in skip_attrs_due_to_errs:
-                skip_attrs_due_to_errs[var_type] = set()
-            skip_attrs_due_to_errs[var_type].add(attr_name)
+            if not (
+                attr_name == "data"
+                and isinstance(var, torch.Tensor)
+                and not include_tensor_data
+            ):
+                logger.warning(
+                    f"Failed to get attribute {attr_name} of object type {type(var)}, skipping it for all following dumps for this attribute."
+                )
+                if var_type not in skip_attrs_due_to_errs:
+                    skip_attrs_due_to_errs[var_type] = set()
+                skip_attrs_due_to_errs[var_type].add(attr_name)
             continue
 
         attr_name = str(attr_name)
@@ -399,7 +441,25 @@ def convert_var_to_dict(var, include_tensor_data=True, dump_config=None) -> dict
     return result
 
 
+def convert_fake_tensor_to_dict(var):
+    try:
+        shape = tuple(var.shape)
+    except Exception:
+        shape = None
+    try:
+        dtype = str(var.dtype)
+    except Exception:
+        dtype = None
+    return {
+        "fake": True,
+        "shape": shape,
+        "dtype": dtype,
+    }
+
+
 def obj_to_serializable(obj, dump_config=None) -> dict[str, object]:
+    if is_fake_tensor(obj):
+        return {typename_compile(obj): convert_fake_tensor_to_dict(obj)}
     if (
         type(obj) in skip_type_due_to_recursion
         and skip_type_due_to_recursion[type(obj)] > RECURSION_ERR_THRESHOLD
@@ -433,6 +493,9 @@ def var_to_serializable(obj, dump_config=None) -> dict[str, object]:
     If you want to dump the `data` attribute of a tensor, use `convert_var_to_dict` and set `include_tensor_data=True`.
     """
 
+    if is_fake_tensor(obj):
+        return {typename_compile(obj): convert_fake_tensor_to_dict(obj)}
+
     if issubclass(type(obj), dict) and type(obj) != dict:  # noqa E721
         return obj_to_serializable(obj, dump_config=dump_config)
 

traincheck/instrumentor/tracer.py

@@ -31,7 +31,7 @@
 )
 from traincheck.proxy_wrapper.proxy_basics import (
     is_proxied,
-    is_proxyparamtetr,
+    is_proxyparameter,
     unproxy_func,
 )
 from traincheck.proxy_wrapper.proxy_config import enable_C_level_observer
@@ -219,7 +219,7 @@ def global_wrapper(
 
         def find_proxy_in_args(args):
             for i, arg in enumerate(args):
-                if is_proxied(arg) or is_proxyparamtetr(arg):
+                if is_proxied(arg) or is_proxyparameter(arg):
                     proxy_in_args.append(arg)
                 elif type(arg) in [list, tuple]:
                     find_proxy_in_args(arg)
@@ -238,7 +238,7 @@ def find_proxy_in_args(args):
             if "proxy_obj_names" not in pre_record:
                 pre_record["proxy_obj_names"] = []
             for proxy in proxy_in_args:
-                if is_proxyparamtetr(proxy):
+                if is_proxyparameter(proxy):
                     pre_record["proxy_obj_names"].append(
                         [proxy.__dict__["var_name"], "Parameter"]
                     )

traincheck/proxy_wrapper/proxy_basics.py

@@ -4,18 +4,58 @@
 
 import astor
 
+from traincheck.config.config import COMPILE_INTERNAL_MODULE
+
+
+def is_compile_internal_module(obj):
+    mod = getattr(type(obj), "__module__", "") or ""
+    if any(mod.startswith(p) for p in COMPILE_INTERNAL_MODULE):
+        return True
+    name = type(obj).__name__
+    if mod.startswith("torch._dynamo") and name != "OptimizedModule":
+        return True
+    return False
+
+
+def is_fake_tensor(x) -> bool:
+    try:
+        from torch._subclasses.fake_tensor import FakeTensor
+        from torch.fx import Proxy as FxProxy
+
+        if isinstance(x, FakeTensor):
+            return True
+        if isinstance(x, FxProxy):
+            return True
+    except Exception:
+        pass
+
+    try:
+        if is_compile_internal_module(x):
+            return True
+    except Exception:
+        return True
+
+    try:
+        return x.device.type == "meta"
+    except Exception:
+        return True
+
 
 def is_proxied(obj):
     try:
+        if is_fake_tensor(obj):
+            return False
         if obj is not None and "is_traincheck_proxied_obj" in obj.__dict__:
             return True
     except Exception:
         return False
     return False
 
 
-def is_proxyparamtetr(obj):
+def is_proxyparameter(obj):
     try:
+        if is_fake_tensor(obj):
+            return False
         if obj is not None and "is_traincheck_proxyparameter" in obj.__dict__:
             return True
     except Exception:

traincheck/proxy_wrapper/proxy_observer.py

@@ -9,7 +9,7 @@
     from traincheck.proxy_wrapper.proxy import Proxy
     from traincheck.proxy_wrapper.subclass import ProxyParameter
 
-from .proxy_basics import is_proxied, unproxy_func
+from .proxy_basics import is_proxied, is_proxyparameter, unproxy_func
 
 
 def observe_proxy_var(
@@ -41,9 +41,9 @@ def wrapper(*args, **kwargs):
             # if the arg is list or tuple, check if it contains proxied object
             if type(arg) in [list, tuple]:
                 for element in arg:
-                    if is_proxied(element) or isinstance(element, ProxyParameter):
+                    if is_proxied(element) or is_proxyparameter(element):
                         proxied_vars.append(element)
-            if is_proxied(arg) or isinstance(arg, ProxyParameter):
+            if is_proxied(arg) or is_proxyparameter(arg):
                 proxied_vars.append(arg)
 
         # pre observe

traincheck/proxy_wrapper/subclass.py

@@ -9,6 +9,7 @@
 from traincheck.utils import get_timestamp_ns
 
 from .dumper import json_dumper as dumper
+from .proxy_basics import is_fake_tensor
 
 # from .proxy_registry import get_global_registry
 # from .utils import print_debug
@@ -23,22 +24,6 @@ def in_dynamo() -> bool:
         return False
 
 
-def is_fake_tensor(x: torch.Tensor) -> bool:
-    try:
-        from torch._subclasses.fake_tensor import FakeTensor  # 2.x
-
-        if isinstance(x, FakeTensor):
-            return True
-    except Exception:
-        pass
-    if getattr(x, "fake_mode", None) is not None:
-        return True
-    if getattr(x, "_is_fake", False):
-        return True
-
-    return isinstance(x, torch.Tensor) and x.device.type == "meta"
-
-
 class ProxyParameter(torch.nn.Parameter):
     loglevel = logging.INFO
     jsondumper = dumper(
@@ -59,13 +44,13 @@ def __new__(
         # TODO
         # from_copy=False,
     ):
+        if isinstance(data, ProxyParameter):
+            return data
         if in_dynamo() or is_fake_tensor(data):
             if isinstance(data, nn.Parameter):
                 return data
             return nn.Parameter(data, requires_grad=data.requires_grad)
         # TODO: verify
-        if isinstance(data, ProxyParameter):
-            return data
 
         return torch.Tensor._make_subclass(cls, data.detach(), data.requires_grad)
 

traincheck/utils.py

@@ -35,6 +35,14 @@ def safe_getattr(obj, attr, default=None):
         raise
 
 
+def typename_compile(o):
+    try:
+        mod = getattr(type(o), "__module__", "") or ""
+        return f"{mod}.{type(o).__name__}"
+    except Exception:
+        return "compile_stage"
+
+
 def typename(o, is_runtime=False):
     if isinstance(o, torch.nn.Parameter):
         return "torch.nn.Parameter"