Fx trans for nn graph

examples/torch_interpretor.py

@@ -1,4 +1,5 @@
 # HF_HUB_OFFLINE=1 python3 examples/torch_interpretor.py
+import os
 import torch
 from diffusers import StableDiffusionPipeline
 from onediff.infer_compiler import torchbackend
@@ -10,13 +11,14 @@
  torch_dtype=torch.float16,
 )
 
+os.environ["with_interp"] = "0"
+os.environ["with_graph"] = "1"
 pipe.unet = torch.compile(pipe.unet, fullgraph=True, mode="reduce-overhead", backend=torchbackend)
 pipe = pipe.to("cuda")
 
 prompt = "a photo of an astronaut riding a horse on mars"
 with torch.autocast("cuda"):
- images = pipe(prompt).images
- images = pipe(prompt).images
- images = pipe(prompt).images
- for i, image in enumerate(images):
- image.save(f"{prompt}-of-{i}.png")
+ for i in range(3):
+ images = pipe(prompt).images
+ for j, image in enumerate(images):
+ image.save(f'{prompt}-of-{i}-{j}.png')

src/onediff/infer_compiler/__init__.py

@@ -1,202 +1,30 @@
 import os
 import torch
-import diffusers
-import oneflow
 import oneflow as flow
-from torch.fx.experimental.proxy_tensor import make_fx
-from torch.func import functionalize
-import importlib
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
-from .attention_1f import BasicTransformerBlock, FeedForward, GEGLU
-from .attention_processor_1f import Attention
-from .lora_1f import LoRACompatibleLinear
-
-
-_is_diffusers_quant_available = False
-try:
- import diffusers_quant
- _is_diffusers_quant_available = True
-except:
- pass
-
-
-def replace_class(cls):
- if cls.__module__.startswith("torch"):
- mod_name = cls.__module__.replace("torch", "oneflow")
- mod = importlib.import_module(mod_name)
- return getattr(mod, cls.__name__)
- elif cls == diffusers.models.attention.BasicTransformerBlock:
- return BasicTransformerBlock
- elif cls == diffusers.models.attention_processor.Attention:
- return Attention
- elif cls == diffusers.models.attention.FeedForward:
- return FeedForward
- elif cls == diffusers.models.attention.GEGLU:
- return GEGLU
- elif cls == diffusers.models.lora.LoRACompatibleLinear:
- return LoRACompatibleLinear
-
- if _is_diffusers_quant_available:
- if cls == diffusers_quant.FakeQuantModule:
- return diffusers_quant.OneFlowFakeQuantModule
- if cls == diffusers_quant.StaticQuantModule:
- return diffusers_quant.OneFlowStaticQuantModule
- if cls == diffusers_quant.DynamicQuantModule:
- return diffusers_quant.OneFlowDynamicQuantModule
-
-
-def replace_obj(obj):
- cls = type(obj)
- if cls == torch.dtype:
- return {
- "torch.float16": flow.float16,
- "torch.float32": flow.float32,
- "torch.double": flow.double,
- "torch.int8": flow.int8,
- "torch.int32": flow.int32,
- "torch.int64": flow.int64,
- "torch.uint8": flow.uint8,
- }[str(obj)]
- if cls == torch.fx.immutable_collections.immutable_list:
- return [e for e in obj]
- replacement = replace_class(cls)
- if replacement is not None:
- if cls in [torch.device]:
- return replacement(str(obj))
- elif cls == torch.nn.parameter.Parameter:
- return flow.utils.tensor.from_torch(obj.data)
- else:
- raise RuntimeError("don't know how to create oneflow obj for: " + str(cls))
- else:
- return obj
-
-
-def replace_func(func):
- if func == torch.conv2d:
- return oneflow.nn.functional.conv2d
- if func == torch._C._nn.linear:
- return oneflow.nn.functional.linear
- if func.__module__.startswith("torch"):
- mod_name = func.__module__.replace("torch", "oneflow")
- mod = importlib.import_module(mod_name)
- return getattr(mod, func.__name__)
- else:
- return func
-
-
-def map_args(args, kwargs):
- args = [replace_obj(a) for a in args]
- kwargs = dict((k, replace_obj(v)) for (k, v) in kwargs.items())
- return (args, kwargs)
-
-
-class ProxySubmodule:
- def __init__(self, submod):
- self._1f_proxy_submod = submod
- self._1f_proxy_parameters = dict()
- self._1f_proxy_children = dict()
-
- def __getattribute__(self, attribute):
- if attribute.startswith("_1f_proxy"):
- return object.__getattribute__(self, attribute)
- elif attribute in ["forward", "_conv_forward"]:
- replacement = replace_class(type(self._1f_proxy_submod))
- return lambda *args, **kwargs: getattr(replacement, attribute)(
- self, *args, **kwargs
- )
- elif (
- isinstance(
- self._1f_proxy_submod, diffusers.models.attention_processor.Attention
- )
- and attribute == "get_attention_scores"
- ):
- replacement = replace_class(type(self._1f_proxy_submod))
- return lambda *args, **kwargs: getattr(replacement, attribute)(
- self, *args, **kwargs
- )
- elif (
- isinstance(self._1f_proxy_submod, torch.nn.Linear)
- and attribute == "use_fused_matmul_bias"
- ):
- return (
- self.bias is not None
- and os.getenv("ONEFLOW_KERNEL_ENABLE_FUSED_LINEAR") == "1"
- )
- elif (
- isinstance(self._1f_proxy_submod, torch.nn.Dropout)
- and attribute == "generator"
- ):
- return flow.Generator()
- elif (
- isinstance(self._1f_proxy_submod, torch.nn.Conv2d)
- and attribute == "channel_pos"
- ):
- return "channels_first"
- else:
- a = getattr(self._1f_proxy_submod, attribute)
- if isinstance(a, torch.Tensor):
- a = flow.utils.tensor.from_torch(a.data)
- elif isinstance(a, torch.nn.parameter.Parameter):
- # TODO(oneflow): assert a.requires_grad == False
- if attribute not in self._1f_proxy_parameters:
- a = flow.utils.tensor.from_torch(a.data)
- self._1f_proxy_parameters[attribute] = a
- else:
- a = self._1f_proxy_parameters[attribute]
- elif isinstance(a, torch.nn.ModuleList):
- a = [ProxySubmodule(m) for m in a]
- elif isinstance(a, torch.nn.Module):
- if attribute not in self._1f_proxy_children:
- a = ProxySubmodule(a)
- self._1f_proxy_children[attribute] = a
- else:
- a = self._1f_proxy_children[attribute]
- assert (
- type(a).__module__.startswith("torch") == False
- and type(a).__module__.startswith("diffusers") == False
- ), "can't be a torch module at this point! But found " + str(type(a))
- return a
-
- def __call__(self, *args: Any, **kwargs: Any) -> Any:
- replacement = replace_class(type(self._1f_proxy_submod))
- if replacement is not None:
- return replacement.__call__(self, *args, **kwargs)
- else:
- raise RuntimeError(
- "can't find oneflow module for: " + str(type(self._1f_proxy_submod))
- )
-
-
-class OneFlowInterpreter(torch.fx.Interpreter):
- from torch.fx.node import Argument, Node, Target, map_arg, map_aggregate
-
- def call_function(self, target: Target, args: Tuple, kwargs: Dict) -> Any:
- if target == torch.sigmoid:
- return torch.neg(*args, **kwargs)
- args, kwargs = map_args(args, kwargs)
- target = replace_func(target)
- return super().call_function(target, args, kwargs)
-
- def call_method(self, target: Target, args: Tuple, kwargs: Dict) -> Any:
- args, kwargs = map_args(args, kwargs)
- return super().call_method(target, args, kwargs)
-
- def call_module(
- self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any]
- ) -> Any:
- submod = self.fetch_attr(target)
- submod = ProxySubmodule(submod)
- return submod(*args, **kwargs)
+from .with_fx_interpreter import OneFlowInterpreter
+from .with_fx_graph import fx_node_tranform
 
 
 def torchbackend(gm, example_inputs):
+ with_interp = os.getenv("with_interp", "True").lower() in (
+ "true",
+ "1",
+ "t",
+ )
+ if not with_interp:
+ transformed_fn = fx_node_tranform(gm)
+
  def wrapped_forward(*args, **kwargs):
  args = [flow.utils.tensor.from_torch(a) for a in args]
- output = OneFlowInterpreter(gm, garbage_collect_values=False).run(
- *args, **kwargs
- )
+ if with_interp:
+ output = OneFlowInterpreter(gm, garbage_collect_values=False).run(
+ *args, **kwargs
+ )
+ else:
+ output = transformed_fn(*args, **kwargs)
  if isinstance(output, tuple):
  return tuple(flow.utils.tensor.to_torch(i) for i in output)
  return flow.utils.tensor.to_torch(output)
 
  return wrapped_forward
+

src/onediff/infer_compiler/obj_1f_from_torch.py

@@ -0,0 +1,145 @@
+import importlib
+import os
+import torch
+import oneflow as flow
+import diffusers
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+from .attention_1f import BasicTransformerBlock, FeedForward, GEGLU
+from .attention_processor_1f import Attention
+from .lora_1f import LoRACompatibleLinear
+
+def replace_class(cls):
+ if cls.__module__.startswith("torch"):
+ mod_name = cls.__module__.replace("torch", "oneflow")
+ mod = importlib.import_module(mod_name)
+ return getattr(mod, cls.__name__)
+ elif cls == diffusers.models.attention.BasicTransformerBlock:
+ return BasicTransformerBlock
+ elif cls == diffusers.models.attention_processor.Attention:
+ return Attention
+ elif cls == diffusers.models.attention.FeedForward:
+ return FeedForward
+ elif cls == diffusers.models.attention.GEGLU:
+ return GEGLU
+ elif cls == diffusers.models.lora.LoRACompatibleLinear:
+ return LoRACompatibleLinear
+
+
+def replace_obj(obj):
+ cls = type(obj)
+ if cls == torch.dtype:
+ return {
+ "torch.float16": flow.float16,
+ "torch.float32": flow.float32,
+ "torch.double": flow.double,
+ "torch.int8": flow.int8,
+ "torch.int32": flow.int32,
+ "torch.int64": flow.int64,
+ "torch.uint8": flow.uint8,
+ }[str(obj)]
+ if cls == torch.fx.immutable_collections.immutable_list:
+ return [e for e in obj]
+ replacement = replace_class(cls)
+ if replacement is not None:
+ if cls in [torch.device]:
+ return replacement(str(obj))
+ elif cls == torch.nn.parameter.Parameter:
+ return flow.utils.tensor.from_torch(obj.data)
+ else:
+ raise RuntimeError("don't know how to create oneflow obj for: " + str(cls))
+ else:
+ return obj
+
+
+def replace_func(func):
+ if func == torch.conv2d:
+ return flow.nn.functional.conv2d
+ if func == torch._C._nn.linear:
+ return flow.nn.functional.linear
+ if func.__module__.startswith("torch"):
+ mod_name = func.__module__.replace("torch", "oneflow")
+ mod = importlib.import_module(mod_name)
+ return getattr(mod, func.__name__)
+ else:
+ return func
+
+
+def map_args(args, kwargs):
+ args = [replace_obj(a) for a in args]
+ kwargs = dict((k, replace_obj(v)) for (k, v) in kwargs.items())
+ return (args, kwargs)
+
+
+class ProxySubmodule:
+ def __init__(self, submod):
+ self._1f_proxy_submod = submod
+ self._1f_proxy_parameters = dict()
+ self._1f_proxy_children = dict()
+
+ def __getattribute__(self, attribute):
+ if attribute.startswith("_1f_proxy"):
+ return object.__getattribute__(self, attribute)
+ elif attribute in ["forward", "_conv_forward"]:
+ replacement = replace_class(type(self._1f_proxy_submod))
+ return lambda *args, **kwargs: getattr(replacement, attribute)(
+ self, *args, **kwargs
+ )
+ elif (
+ isinstance(
+ self._1f_proxy_submod, diffusers.models.attention_processor.Attention
+ )
+ and attribute == "get_attention_scores"
+ ):
+ replacement = replace_class(type(self._1f_proxy_submod))
+ return lambda *args, **kwargs: getattr(replacement, attribute)(
+ self, *args, **kwargs
+ )
+ elif (
+ isinstance(self._1f_proxy_submod, torch.nn.Linear)
+ and attribute == "use_fused_matmul_bias"
+ ):
+ return (
+ self.bias is not None
+ and os.getenv("ONEFLOW_KERNEL_ENABLE_FUSED_LINEAR") == "1"
+ )
+ elif (
+ isinstance(self._1f_proxy_submod, torch.nn.Dropout)
+ and attribute == "generator"
+ ):
+ return flow.Generator()
+ elif (
+ isinstance(self._1f_proxy_submod, torch.nn.Conv2d)
+ and attribute == "channel_pos"
+ ):
+ return "channels_first"
+ else:
+ a = getattr(self._1f_proxy_submod, attribute)
+ if isinstance(a, torch.nn.parameter.Parameter):
+ # TODO(oneflow): assert a.requires_grad == False
+ if attribute not in self._1f_proxy_parameters:
+ a = flow.utils.tensor.from_torch(a.data)
+ self._1f_proxy_parameters[attribute] = a
+ else:
+ a = self._1f_proxy_parameters[attribute]
+ elif isinstance(a, torch.nn.ModuleList):
+ a = [ProxySubmodule(m) for m in a]
+ elif isinstance(a, torch.nn.Module):
+ if attribute not in self._1f_proxy_children:
+ a = ProxySubmodule(a)
+ self._1f_proxy_children[attribute] = a
+ else:
+ a = self._1f_proxy_children[attribute]
+ assert (
+ type(a).__module__.startswith("torch") == False
+ and type(a).__module__.startswith("diffusers") == False
+ ), "can't be a torch module at this point! But found " + str(type(a))
+ return a
+
+ def __call__(self, *args: Any, **kwargs: Any) -> Any:
+ replacement = replace_class(type(self._1f_proxy_submod))
+ if replacement is not None:
+ return replacement.__call__(self, *args, **kwargs)
+ else:
+ raise RuntimeError(
+ "can't find oneflow module for: " + str(type(self._1f_proxy_submod))
+ )