Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix Check Failed Error #585

Merged
merged 8 commits into from
Jan 30, 2024
Merged

Fix Check Failed Error #585

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d94356a
Fix Check Failed Error
ccssu Jan 30, 2024
e94be30
refine
ccssu Jan 30, 2024
193bd28
refine
ccssu Jan 30, 2024
5dfc92d
refine
ccssu Jan 30, 2024
5ef8157
fix warning
ccssu Jan 30, 2024
1de7b06
fix warning
ccssu Jan 30, 2024
a9c9f4a
refine
ccssu Jan 30, 2024
274541b
Merge branch 'main' into fix_check_model_input_count_failed_error
strint Jan 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions src/onediff/infer_compiler/utils/args_tree_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import torch
import oneflow as flow
from oneflow.framework.args_tree import ArgsTree
from .log_utils import logger


def input_output_processor(func):
Expand All @@ -13,10 +14,13 @@ def input_fn(value):
return value

args_tree = ArgsTree((args, kwargs), False, tensor_type=torch.Tensor)
input_count = len(
[v for v in args_tree.iter_nodes() if isinstance(v, torch.Tensor)]
)
out = args_tree.map_leaf(input_fn)
mapped_args = out[0]
mapped_kwargs = out[1]
return mapped_args, mapped_kwargs
return mapped_args, mapped_kwargs, input_count

def process_output(output):
def output_fn(value):
Expand All @@ -29,9 +33,19 @@ def output_fn(value):
out = out_tree.map_leaf(output_fn)
return out[0]

def wrapper(cls, *args, **kwargs):
mapped_args, mapped_kwargs = process_input(*args, **kwargs)
output = func(cls, *mapped_args, **mapped_kwargs)
def wrapper(self: "DeployableModule", *args, **kwargs):
mapped_args, mapped_kwargs, input_count = process_input(*args, **kwargs)
if self._deployable_module_use_graph:
count = self._deployable_module_input_count
if count != input_count:
logger.warning(
f"Module {type(self._deployable_module_model.oneflow_module)} input tensor count changed from {count} to {input_count}, will compile again."
)
self._deployable_module_dpl_graph = None
self._load_graph_first_run = True
self._deployable_module_input_count = input_count

output = func(self, *mapped_args, **mapped_kwargs)
return process_output(output)

return wrapper
7 changes: 7 additions & 0 deletions src/onediff/infer_compiler/with_oneflow_compile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,8 @@ def __init__(
self._deployable_module_options = options
self._deployable_module_dpl_graph = None
self._is_raw_deployable_module = True
self._load_graph_first_run = True
self._deployable_module_input_count = None

@classmethod
def from_existing(cls, existing_module, use_graph=None, dynamic=None, options=None):
Expand All @@ -207,6 +209,11 @@ def from_existing(cls, existing_module, use_graph=None, dynamic=None, options=No
instance._deployable_module_dpl_graph = (
existing_module._deployable_module_dpl_graph if use_graph else None
)
instance._load_graph_first_run = existing_module._load_graph_first_run
instance._deployable_module_input_count = (
existing_module._deployable_module_input_count
)

return instance

def get_graph(self):
Expand Down
Loading