Make backend_accuracy_fails suppress errors in same_two_models (#100324)

ezyang · pytorchmergebot · commit 8556cf208a3f · 2023-05-02T11:44:16.000Z
The basic idea is that if we're trying to match for an accuracy error, we don't want to switch to a compile/runtime error, because that's probably us breaking things in a different way. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #100324 Approved by: https://github.com/voznesenskym
diff --git a/torch/_dynamo/debug_utils.py b/torch/_dynamo/debug_utils.py
@@ -309,7 +309,7 @@ def run_fwd_maybe_bwd(gm, args, only_fwd=False):
     return collect_results(gm, out, None, args)
 
 
-def same_two_models(gm, opt_gm, example_inputs, only_fwd=False):
+def same_two_models(gm, opt_gm, example_inputs, only_fwd=False, *, require_fp64=False):
     """
     Check two models have same accuracy.
     """
@@ -336,6 +336,8 @@ def same_two_models(gm, opt_gm, example_inputs, only_fwd=False):
         )
         fp64_ref = run_fwd_maybe_bwd(fp64_model, fp64_examples, only_fwd)
     except Exception:
+        if require_fp64:
+            raise RuntimeError("Could not generate fp64 outputs")
         log.warning("Could not generate fp64 outputs")
         fp64_ref = None
 
@@ -393,11 +395,16 @@ def cast_to_fp64(model, inputs):
     return cast_to(torch.float64, model, inputs)
 
 
-def backend_accuracy_fails(gm, example_inputs, compiler_fn, only_fwd=False):
+def backend_accuracy_fails(
+    gm, example_inputs, compiler_fn, only_fwd=False, *, require_fp64=False
+):
     try:
         compiled_gm = compiler_fn(
             copy.deepcopy(gm), clone_inputs_retaining_gradness(example_inputs)
         )
+        return not same_two_models(
+            gm, compiled_gm, example_inputs, only_fwd, require_fp64=require_fp64
+        )
     except Exception as e:
         # This means that the the minified graph is bad/exposes a different problem.
         # As we are checking accuracy here, lets log the exception and return False.
@@ -409,5 +416,3 @@ def backend_accuracy_fails(gm, example_inputs, compiler_fn, only_fwd=False):
             )
         )
         return False
-
-    return not same_two_models(gm, compiled_gm, example_inputs, only_fwd)
