Flag guard unbacked SymInt/SymFloat support (pytorch#94987)

I believe this fixes the AllenaiLongformerBase problem in periodic.

The longer version of the problem is here is we are currently optimistically converting all item() calls into unbacked SymInt/SymFloat, but sometimes this results in a downstream error due to a data-dependent guard. Fallbacks for this case are non-existent; this will just crash the model. This is bad. So we flag guard until we get working fallbacks.

What could these fallbacks look like? One idea I have is to optimistically make data-dependent calls unbacked, but then if it results in a crash, restart Dynamo analysis with the plan of graph breaking when the item() call immediately happened.

Signed-off-by: Edward Z. Yang <ezyang@meta.com>

Pull Request resolved: pytorch#94987
Approved by: https://github.com/Skylion007, https://github.com/malfet

benchmarks/dynamo/common.py

@@ -182,7 +182,6 @@ class CI(NamedTuple):
 CI_SKIP[CI("aot_eager", training=True, dynamic=True)] = [
     *CI_SKIP[CI("aot_eager", training=True)],
     *CI_SKIP[CI("aot_eager", training=False, dynamic=True)],
-    "AllenaiLongformerBase",  # GuardOnDataDependentSymNode
 ]
 
 CI_SKIP[CI("inductor", training=False, dynamic=True)] = [

torch/_dynamo/output_graph.py

@@ -180,8 +180,11 @@ def __init__(
         super().__init__()
         self.graph = torch.fx.Graph()
         self.graphargs: List[GraphArg] = []
+        shape_env = None
+        if config.dynamic_shapes:
+            shape_env = ShapeEnv(allow_scalar_outputs=config.capture_scalar_outputs)
         fake_mode = torch._subclasses.FakeTensorMode(
-            shape_env=ShapeEnv() if config.dynamic_shapes else None,
+            shape_env=shape_env,
         )
         self.tracing_context: TracingContext = TracingContext(fake_mode)
         if config.dynamic_shapes:

torch/_subclasses/fake_tensor.py

@@ -408,6 +408,8 @@ def local_scalar_dense(fake_mode, func, arg):
     if fake_mode.shape_env is None:
         # Without symints/symfloats, cannot handle this
         raise DataDependentOutputException(func)
+    if not fake_mode.shape_env.allow_scalar_outputs:
+        raise DataDependentOutputException(func)
     if is_float_dtype(arg.dtype):
         return fake_mode.shape_env.create_unbacked_symfloat()
     elif is_integer_dtype(arg.dtype):

torch/fx/experimental/symbolic_shapes.py

@@ -1063,7 +1063,9 @@ def _print_Symbol(self, expr) -> str:
 
 
 class ShapeEnv:
-    def __init__(self):
+    def __init__(self, allow_scalar_outputs=True):
+        # Not directly used by ShapeEnv; indirectly used by FakeTensor
+        self.allow_scalar_outputs = allow_scalar_outputs
         self.guards: List[ShapeGuard] = []
         # Maps symbolic ints to their original concrete values
         # Currently populated from tensors