[PT2][Inductor] Add runtime numeric check for the post grad pass (#136724)

mengluy0125 · pytorchmergebot · commit a00f0d5db816 · 2024-10-01T18:56:01.000Z
Summary: Similar to D51838043, we further add post grad runtime numeric check since some graph passes are performed at aten-level. Differential Revision: D63438718 Pull Request resolved: #136724 Approved by: https://github.com/Yuzhen11
diff --git a/torch/_inductor/compile_fx.py b/torch/_inductor/compile_fx.py
@@ -284,11 +284,17 @@ def _recursive_joint_graph_passes(gm):
     joint_graph_passes(gm)
 
 
-def _recursive_post_grad_passes(gm, is_inference: bool = False):
+def _recursive_post_grad_passes(
+    gm, is_inference: bool = False, example_inputs=None, fake_mode=None
+):
     for subgraph_name in _get_subgraph_names(gm):
         subgraph = getattr(gm, subgraph_name)
-        _recursive_post_grad_passes(subgraph, is_inference)
-    post_grad_passes(gm, is_inference)
+        _recursive_post_grad_passes(
+            subgraph, is_inference, example_inputs=example_inputs, fake_mode=fake_mode
+        )
+    post_grad_passes(
+        gm, is_inference, example_inputs=example_inputs, fake_mode=fake_mode
+    )
 
 
 def split_const_gm(
@@ -770,7 +776,12 @@ def log_graph_runnable():
 
         with V.set_fake_mode(fake_mode):
             # has some issues with memory in training
-            _recursive_post_grad_passes(gm, is_inference=is_inference)
+            _recursive_post_grad_passes(
+                gm,
+                is_inference=is_inference,
+                example_inputs=example_inputs,
+                fake_mode=fake_mode,
+            )
             V.debug.fx_graph_transformed(gm, example_inputs)
             post_grad_graphs_log.debug(
                 "%s",
diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py
@@ -224,6 +224,7 @@ def autotune_remote_cache_default() -> Optional[bool]:
 # https://pytorch.org/docs/stable/notes/numerical_accuracy.html#batched-computations-or-slice-computations
 fx_passes_numeric_check: Dict[str, Any] = {
     "pre_grad": False,
+    "post_grad": False,
     "precision": 1e-4,
     "num_iterations": 1,
     "requires_optimizer": True,
diff --git a/torch/_inductor/fx_passes/numeric_utils.py b/torch/_inductor/fx_passes/numeric_utils.py
@@ -1,4 +1,5 @@
 # mypy: allow-untyped-defs
+import copy
 import gc
 import logging
 import os
@@ -11,6 +12,7 @@
 import torch.optim as optim
 
 from .. import config
+from ..pattern_matcher import stable_topological_sort
 
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -129,6 +131,14 @@ def compare_gradients(model_base, model_control, precision):
     )
 
 
+def compare_buffers(model_base, model_control, precision):
+    return compare_dict_tensors(
+        dict(model_base.named_buffers()),
+        dict(model_control.named_buffers()),
+        precision,
+    )
+
+
 def run_model(
     model_base, model_control, model_input, num_iterations=10, precision=1e-4
 ):
@@ -145,6 +155,9 @@ def run_model(
 
         res = compare_forward_output(pred_base, pred_control, precision)
         logger.info("compare loss/predict. Numerical result : %s", res)
+
+        res = compare_buffers(model_base, model_control, precision)
+        logger.info("compare buffers. Numerical result : %s", res)
         # tensor may not have a grad_fn
         try:
             _ = pred_base[0].sum().backward(retain_graph=True)
@@ -210,3 +223,36 @@ def numeric_check_if_enabled(
             "Runtime numeric check failed in pre grad fx passes with error: %s", e
         )
         traceback.print_exc()
+
+
+def enable_runtime_numeric_check(
+    example_inputs, fake_mode, gm_before_fx_passes, gm, fx_passes_numeric_check
+):
+    """
+    Enable runtime numeric check for fx passes.
+    """
+    if (
+        config.pattern_matcher
+        and hasattr(config, "fx_passes_numeric_check")
+        and fx_passes_numeric_check
+        and example_inputs is not None
+    ):
+        stable_topological_sort(gm.graph)
+        gm_after_fx_passes = copy.deepcopy(gm)
+        if fake_mode is not None:
+            with fake_mode:
+                numeric_check_if_enabled(
+                    gm_before_fx_passes,  # type: ignore[possibly-undefined]
+                    gm_after_fx_passes,
+                    example_inputs,
+                    config.fx_passes_numeric_check.get("num_iterations", 1),
+                    config.fx_passes_numeric_check.get("precision", 1e-4),
+                )
+        else:
+            numeric_check_if_enabled(
+                gm_before_fx_passes,  # type: ignore[possibly-undefined]
+                gm_after_fx_passes,
+                example_inputs,
+                config.fx_passes_numeric_check.get("num_iterations", 1),
+                config.fx_passes_numeric_check.get("precision", 1e-4),
+            )
diff --git a/torch/_inductor/fx_passes/post_grad.py b/torch/_inductor/fx_passes/post_grad.py
@@ -1,5 +1,6 @@
 # mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
+import copy
 import functools
 import itertools
 import logging
@@ -49,6 +50,7 @@
 from .ddp_fusion import fuse_ddp_communication
 from .group_batch_fusion import group_batch_fusion_passes, POST_GRAD_FUSIONS
 from .micro_pipeline_tp import micro_pipeline_tp_pass
+from .numeric_utils import enable_runtime_numeric_check
 from .pre_grad import is_same_dict, save_inductor_dict
 from .reinplace import reinplace_inplaceable_ops
 from .split_cat import POST_GRAD_PATTERNS
@@ -70,7 +72,9 @@
 ]
 
 
-def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
+def post_grad_passes(
+    gm: torch.fx.GraphModule, is_inference: bool, example_inputs=None, fake_mode=None
+):
     """
     Passes that run on after grad.  This is called once on the forwards
     graph and once on the backwards graph.
@@ -97,6 +101,11 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
 
     if config.pattern_matcher:
         lazy_init()
+        gm_before_fx_passes = None
+        if hasattr(
+            config, "fx_passes_numeric_check"
+        ) and config.fx_passes_numeric_check.get("post_grad", False):
+            gm_before_fx_passes = copy.deepcopy(gm)
         optimus_scuba_log["before_recompile_post_grad"] = upload_graph(gm.graph)
         group_batch_fusion_passes(gm.graph, pre_grad=False)
         remove_noop_ops(gm.graph)
@@ -115,6 +124,11 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
                 optimus_scuba_log[
                     f"{pattern_matcher_pass.pass_name}_post_grad"
                 ] = upload_graph(gm.graph)
+        optimus_scuba_log["after_recompile_post_grad"] = upload_graph(gm.graph)
+        fx_passes_numeric_check = config.fx_passes_numeric_check.get("post_grad", False)
+        enable_runtime_numeric_check(
+            example_inputs, fake_mode, gm_before_fx_passes, gm, fx_passes_numeric_check
+        )
         if config.b2b_gemm_pass:
             B2B_GEMM_PASS.apply(gm.graph)  # type: ignore[arg-type]
 
@@ -148,7 +162,6 @@ def post_grad_passes(gm: torch.fx.GraphModule, is_inference: bool):
     comms.reinplace_fsdp_all_gather(gm.graph)
 
     gm.recompile()
-    optimus_scuba_log["after_recompile_post_grad"] = upload_graph(gm.graph)
     gm.graph.lint()
 
 
diff --git a/torch/_inductor/fx_passes/pre_grad.py b/torch/_inductor/fx_passes/pre_grad.py
@@ -27,6 +27,7 @@
 from ..utils import is_cpu_device, pass_execution_and_save
 from .group_batch_fusion import group_batch_fusion_passes, PRE_GRAD_FUSIONS
 from .misc_patterns import numpy_compat_normalization
+from .numeric_utils import enable_runtime_numeric_check
 from .split_cat import PRE_GRAD_PATTERNS
 
 
@@ -126,10 +127,11 @@ def pre_grad_passes(gm: torch.fx.GraphModule, example_inputs=None):
     """
     if config.pattern_matcher:
         lazy_init()
+        gm_before_fx_passes = None
         if hasattr(
             config, "fx_passes_numeric_check"
         ) and config.fx_passes_numeric_check.get("pre_grad", False):
-            gm_before_fx_passes = gm.__copy__()
+            gm_before_fx_passes = copy.deepcopy(gm)
         # explicitly run with predispatch atenIR based passes
         if config.is_predispatch:
 
@@ -260,6 +262,13 @@ def shape_prop(mod) -> None:
                     optimus_scuba_log[
                         f"{pattern_matcher_pass.pass_name}_pre_grad"
                     ] = upload_graph(gm.graph)
+            optimus_scuba_log["after_recompile_pre_grad"] = upload_graph(gm.graph)
+            fx_passes_numeric_check = config.fx_passes_numeric_check.get(
+                "pre_grad", False
+            )
+            enable_runtime_numeric_check(
+                example_inputs, None, gm_before_fx_passes, gm, fx_passes_numeric_check
+            )
             # TODO: move efficient_conv_bn_eval_pass to the fusions dict too.
             efficient_conv_bn_eval_pass.apply(gm.graph)  # type: ignore[arg-type]
 
@@ -276,24 +285,6 @@ def shape_prop(mod) -> None:
 
     gm.graph.lint()
     gm.recompile()
-    optimus_scuba_log["after_recompile_pre_grad"] = upload_graph(gm.graph)
-
-    if (
-        config.pattern_matcher
-        and hasattr(config, "fx_passes_numeric_check")
-        and config.fx_passes_numeric_check.get("pre_grad", False)
-        and example_inputs is not None
-    ):
-        from .numeric_utils import numeric_check_if_enabled
-
-        gm_after_fx_passes = gm.__copy__()
-        numeric_check_if_enabled(
-            gm_before_fx_passes,  # type: ignore[possibly-undefined]
-            gm_after_fx_passes,
-            example_inputs,
-            config.fx_passes_numeric_check.get("num_iterations", 1),
-            config.fx_passes_numeric_check.get("precision", 1e-4),
-        )
 
     return gm
 
