[dynamo] Improve support for inspect.signature().parameters (#115047)

jansel · pytorchmergebot · commit a70c85ce90c0 · 2023-12-04T19:08:36.000Z
Pull Request resolved: #115047 Approved by: https://github.com/oulgen ghstack dependencies: #114830
diff --git a/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/aot_eager_torchbench_inference.csv
@@ -118,7 +118,7 @@ hf_T5,pass,0
 
 
 
-hf_T5_generate,pass,18
+hf_T5_generate,pass,19
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_aot_eager_torchbench_inference.csv
@@ -118,7 +118,7 @@ hf_T5,pass,0
 
 
 
-hf_T5_generate,fail_to_run,9
+hf_T5_generate,fail_to_run,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv
@@ -118,7 +118,7 @@ hf_T5,pass,0
 
 
 
-hf_T5_generate,fail_to_run,9
+hf_T5_generate,fail_to_run,7
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamo_eager_torchbench_inference.csv
@@ -118,7 +118,7 @@ hf_T5,pass,0
 
 
 
-hf_T5_generate,pass,18
+hf_T5_generate,pass,19
 
 
 
diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv
@@ -118,7 +118,7 @@ hf_T5,pass,0
 
 
 
-hf_T5_generate,pass,18
+hf_T5_generate,pass,19
 
 
 
diff --git a/test/dynamo/test_repros.py b/test/dynamo/test_repros.py
@@ -490,6 +490,20 @@ def apply_chunking_to_forward(forward_fn, *input_tensors):
     return forward_fn(*input_tensors)
 
 
+def _validate_model_kwargs(fn, model_kwargs):
+    # simplified from transformers.generation.utils._validate_model_kwargs
+    unused_model_args = []
+    model_args = set(inspect.signature(fn).parameters)
+    for key, value in model_kwargs.items():
+        if value is not None and key not in model_args:
+            unused_model_args.append(key)
+    if unused_model_args:
+        raise ValueError(
+            f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the"
+            " generate arguments will also show up in this list)"
+        )
+
+
 class FakeMamlInner(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -2417,6 +2431,24 @@ def f(x, y):
             self.assertEqual(f(x, x), opt_f(x, x))
             self.assertEqual(f(x, y), opt_f(x, y))
 
+    def test_validate_model_kwargs(self):
+        cnt = CompileCounter()
+
+        def f1(a, b):
+            return torch.sin(a) + torch.cos(b)
+
+        @torch.compile(backend=cnt, fullgraph=True)
+        def f2(**kwargs):
+            _validate_model_kwargs(f1, kwargs)
+            return f1(**kwargs)
+
+        x = torch.randn(10)
+        y = torch.randn(10)
+
+        self.assertEqual(f2(a=x, b=y), f1(x, y))
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertEqual(cnt.op_count, 3)
+
     def test_swin_base_tensor_attr(self):
         class Foo(torch.nn.Module):
             def __init__(self):
diff --git a/torch/_dynamo/variables/base.py b/torch/_dynamo/variables/base.py
@@ -299,8 +299,8 @@ def has_unpack_var_sequence(self, tx):
         except NotImplementedError:
             return False
 
-    def num_parameters(self):
-        unimplemented(f"num_parameters: {self}")
+    def inspect_parameter_names(self) -> List[str]:
+        unimplemented(f"inspect_parameter_names: {self}")
 
     def call_hasattr(self, tx, name: str) -> "VariableTracker":
         unimplemented(f"hasattr: {repr(self)}")
diff --git a/torch/_dynamo/variables/functions.py b/torch/_dynamo/variables/functions.py
@@ -82,8 +82,8 @@ def call_function(
             self, list(self.self_args()) + list(args), kwargs
         )
 
-    def num_parameters(self):
-        return len(inspect.signature(self.get_function()).parameters)
+    def inspect_parameter_names(self):
+        return list(inspect.signature(self.get_function()).parameters)
 
     def closure_vars(self, tx):
         return {}
@@ -293,8 +293,8 @@ def call_function(
                 )
         return super().call_function(tx, args, kwargs)
 
-    def num_parameters(self):
-        return super().num_parameters() - 1
+    def inspect_parameter_names(self):
+        return super().inspect_parameter_names()[1:]
 
 
 class WrappedUserMethodVariable(UserMethodVariable):
diff --git a/torch/_dynamo/variables/misc.py b/torch/_dynamo/variables/misc.py
@@ -276,10 +276,27 @@ def create(callable, **kwargs):
             unimplemented(f"inspect.signature with {kwargs}")
         return InspectSignatureVariable(callable)
 
-    def __init__(self, inspected, **kwargs):
+    def __init__(self, inspected: VariableTracker, **kwargs):
         super().__init__(**kwargs)
         self.inspected = inspected
 
+    def var_getattr(self, tx, name: str) -> "VariableTracker":
+        if name == "parameters":
+            return variables.ConstDictVariable(
+                {
+                    name: InspectParameterVariable()
+                    for name in self.inspected.inspect_parameter_names()
+                },
+                user_cls=dict,
+            )
+        return super().var_getattr(tx, name)
+
+
+class InspectParameterVariable(VariableTracker):
+    """This is not implemented, if used will graph break."""
+
+    pass
+
 
 def produce_trampoline_autograd_fwd(fn_cls):
     def trampoline_autograd_fwd(*args, **kwargs):
@@ -642,23 +659,6 @@ def call_function(
     ) -> "VariableTracker":
         return self.obj.call_method(tx, self.name, args, kwargs)
 
-    def call_method(
-        self,
-        tx,
-        name,
-        args: "List[VariableTracker]",
-        kwargs: "Dict[str, VariableTracker]",
-    ) -> "VariableTracker":
-        if (
-            name == "__len__"
-            and isinstance(self.obj, InspectSignatureVariable)
-            and self.name == "parameters"
-        ):
-            return variables.ConstantVariable.create(
-                self.obj.inspected.num_parameters(),
-            )
-        return super().call_method(tx, name, args, kwargs)
-
 
 class MethodWrapperVariable(VariableTracker):
     def __init__(self, method_wrapper, **kwargs):
