Repository: pytorch
Issue: activation memory budget partitioner can fail with unbacked symints
Labels: high priority,triaged,oncall: pt2,module: dynamic shapes,module: aotdispatch,module: pt2-dispatcher

internal xref: https://fb.workplace.com/groups/1075192433118967/posts/1567692087202330/?comment_id=1572673046704234&reply_comment_id=1577244289580443

Stacktrace below. Still working on a minimal repro, but a few things that become apparent from looking at the [tlparse](https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/aps-ig_v4_4t_4_brian_fix_more_log-1f26ffb451/attempt_0/version_0/rank_0/1_0_0/aot_joint_graph_7.txt?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=100&fbclid=IwZXh0bgNhZW0CMTEAAR3PXnExS6X2RbGJVZA4YDlJxOplWmMtBCqTDDkB1F4nLGu1N3aML9CZ-kM_aem_XfnDO_Q9FfE_kYaJbP3IHw):

(1) The partitioner runs some ops in the graph under the FlopCounter during mem budget partitioning, but errors inside one of the meta functions

(2) surprisingly, these meta functions didn't error previously during dynamo / AOT tracing - they only fail later when they are re-run in the partitioner.

Still working on a min repro
```
Traceback (most recent call last):
  File "/.../torch/_dynamo/convert_frame.py", line 986, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/.../torch/_dynamo/convert_frame.py", line 715, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/.../torch/_utils_internal.py", line 310, in wrapper_function
    return func(*args, **kwargs)
  File "/.../torch/_dynamo/convert_frame.py", line 750, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "/.../torch/_dynamo/bytecode_transformation.py", line 1361, in transform_code_object
    transformations(instructions, code_options)
  File "/.../torch/_dynamo/convert_frame.py", line 231, in _fn
    return fn(*args, **kwargs)
  File "/.../torch/_dynamo/convert_frame.py", line 662, in transform
    tracer.run()
  File "/.../torch/_dynamo/symbolic_convert.py", line 2870, in run
    super().run()
  File "/.../torch/_dynamo/symbolic_convert.py", line 1053, in run
    while self.step():
  File "/.../torch/_dynamo/symbolic_convert.py", line 963, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/.../torch/_dynamo/symbolic_convert.py", line 3050, in RETURN_VALUE
    self._return(inst)
  File "/.../torch/_dynamo/symbolic_convert.py", line 3035, in _return
    self.output.compile_subgraph(
  File "/.../torch/_dynamo/output_graph.py", line 1136, in compile_subgraph
    self.compile_and_call_fx_graph(
  File "/.../torch/_dynamo/output_graph.py", line 1382, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/.../torch/_dynamo/output_graph.py", line 1432, in call_user_compiler
    return self._call_user_compiler(gm)
  File "/.../torch/_dynamo/output_graph.py", line 1483, in _call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/.../torch/_dynamo/output_graph.py", line 1462, in _call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/.../torch/_dynamo/repro/after_dynamo.py", line 130, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/.../torch/__init__.py", line 2314, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
  File "/.../torch/_inductor/compile_fx.py", line 1880, in compile_fx
    return aot_autograd(
  File "/.../torch/_dynamo/backends/common.py", line 83, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
  File "/.../torch/_functorch/aot_autograd.py", line 1155, in aot_module_simplified
    compiled_fn = dispatch_and_compile()
  File "/.../torch/_functorch/aot_autograd.py", line 1131, in dispatch_and_compile
    compiled_fn, _ = create_aot_dispatcher_function(
  File "/.../torch/_functorch/aot_autograd.py", line 580, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
  File "/.../torch/_functorch/aot_autograd.py", line 830, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
  File "/.../torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 447, in aot_dispatch_autograd
    fw_module, bw_module = aot_config.partition_fn(
  File "/.../torch/_inductor/compile_fx.py", line 1797, in partition_fn
    return min_cut_rematerialization_partition(
  File "/.../torch/_functorch/partitioners.py", line 1823, in min_cut_rematerialization_partition
    saved_values = choose_saved_values_set(
  File "/.../torch/_functorch/partitioners.py", line 1573, in choose_saved_values_set
    runtimes_banned_nodes = [
  File "/.../torch/_functorch/partitioners.py", line 1574, in <listcomp>
    estimate_runtime(node) for node in all_recomputable_banned_nodes
  File "/.../torch/_functorch/partitioners.py", line 1458, in estimate_runtime
    node.target(*args, **kwargs)
  File "/.../torch/_ops.py", line 722, in __call__
    return self._op(*args, **kwargs)
  File "/.../torch/utils/flop_counter.py", line 772, in __torch_dispatch__
    out = func(*args, **kwargs)
  File "/.../torch/_ops.py", line 722, in __call__
    return self._op(*args, **kwargs)
  File "/.../torch/utils/_stats.py", line 21, in wrapper
    return fn(*args, **kwargs)
  File "/.../torch/_subclasses/fake_tensor.py", line 1276, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/.../torch/_subclasses/fake_tensor.py", line 1817, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
  File "/.../torch/_subclasses/fake_tensor.py", line 1387, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
  File "/.../torch/_subclasses/fake_tensor.py", line 2385, in _dispatch_impl
    r = func(*args, **kwargs)
  File "/.../torch/_ops.py", line 722, in __call__
    return self._op(*args, **kwargs)
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError:
Exception raised from <custom_op_meta>
```

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @chauhang @penguinwu @bobrenjc93 @yf225