Repository: pytorch
Issue: [Inductor] StableDiffusion unet with `cudagraphs` backend raises fake tensor mismatch error
Labels: triaged,module: cuda graphs,oncall: pt2,module: inductor

### 🐛 Describe the bug

Fake tensor fake mode mismatch

repro
```python
def create_sd2():
    from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
    model_id = "stabilityai/stable-diffusion-2"

    # Create a model instance
    # Use the Euler scheduler here instead
    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
    pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
    return pipe

model = create_sd2().to('cuda)
model.unet = torch.compile(model.unet, fullgraph=True, backend="cudagraphs")
model("my prompt")
```

Issue seems to be in the `partition_cudagraphs` pass.

### Error logs

```
Exception occurred during processing of request from ('127.0.0.1', 50382)
Traceback (most recent call last):
  File "/usr/lib/python3.10/socketserver.py", line 316, in _handle_request_noblock
    self.process_request(request, client_address)
  File "/usr/lib/python3.10/socketserver.py", line 347, in process_request
    self.finish_request(request, client_address)
  File "/usr/lib/python3.10/socketserver.py", line 360, in finish_request
    self.RequestHandlerClass(request, client_address, self)
  File "/usr/lib/python3.10/socketserver.py", line 747, in __init__
    self.handle()
  File "/usr/lib/python3.10/http/server.py", line 433, in handle
    self.handle_one_request()
  File "/usr/lib/python3.10/http/server.py", line 421, in handle_one_request
    method()
  File "/home/jonch/Desktop/sdpa.py", line 2228, in do_POST
    output = get_model_instance()(post_data)
  File "/home/jonch/Desktop/sdpa.py", line 2253, in compiled
    return model(prompt, compile=True)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/jonch/.local/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 903, in __call__
    encoder_hidden_states=prompt_embeds,
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/eval_frame.py", line 489, in _fn
    return fn(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/eval_frame.py", line 655, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/convert_frame.py", line 721, in _convert_frame
    result = inner_convert(frame, cache_entry, hooks, frame_state)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/convert_frame.py", line 383, in _convert_frame_assert
    compiled_product = _compile(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/convert_frame.py", line 645, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/utils.py", line 244, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/convert_frame.py", line 562, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/bytecode_transformation.py", line 1033, in transform_code_object
    transformations(instructions, code_options)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/convert_frame.py", line 151, in _fn
    return fn(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/convert_frame.py", line 527, in transform
    tracer.run()
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/symbolic_convert.py", line 2123, in run
    super().run()
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/symbolic_convert.py", line 818, in run
    and self.step()
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/symbolic_convert.py", line 781, in step
    getattr(self, inst.opname)(inst)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/symbolic_convert.py", line 2238, in RETURN_VALUE
    self.output.compile_subgraph(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/output_graph.py", line 938, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/output_graph.py", line 1080, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/utils.py", line 244, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/output_graph.py", line 1152, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/output_graph.py", line 1133, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/repro/after_dynamo.py", line 117, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/__init__.py", line 1696, in __call__
    return self.compiler_fn(model_, inputs_, **self.kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/backends/common.py", line 55, in compiler_fn
    cg = aot_module_simplified(gm, example_inputs, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_functorch/aot_autograd.py", line 5021, in aot_module_simplified
    compiled_fn = create_aot_dispatcher_function(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/utils.py", line 244, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_functorch/aot_autograd.py", line 4561, in create_aot_dispatcher_function
    compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config, fw_metadata=fw_metadata)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_functorch/aot_autograd.py", line 2881, in aot_wrapper_dedupe
    return compiler_fn(flat_fn, leaf_flat_args, aot_config, fw_metadata=fw_metadata)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_functorch/aot_autograd.py", line 3067, in aot_wrapper_synthetic_base
    return compiler_fn(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_functorch/aot_autograd.py", line 2179, in aot_dispatch_base
    compiled_fw = compiler(fw_module, updated_flat_args)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_dynamo/backends/cudagraphs.py", line 135, in cudagraphs
    model = partition_cudagraphs(model, inputs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/fx/passes/backends/cudagraphs.py", line 49, in partition_cudagraphs
    FakeTensorProp(gm).propagate(*inputs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/fx/passes/fake_tensor_prop.py", line 53, in propagate
    fake_args = [
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/fx/passes/fake_tensor_prop.py", line 54, in <listcomp>
    self._mode.from_tensor(a) if isinstance(a, torch.Tensor) else a
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 1871, in from_tensor
    return self.fake_tensor_converter(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 389, in __call__
    return self.from_real_tensor(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 346, in from_real_tensor
    out = self.meta_converter(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/meta_utils.py", line 710, in __call__
    r = self.meta_tensor(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/meta_utils.py", line 410, in meta_tensor
    r = _view_from_base(base, t)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/meta_utils.py", line 404, in _view_from_base
    return base.as_strided(sizes, strides, storage_offset)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 1385, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 1522, in dispatch
    (flat_args, flat_arg_fake_tensors) = self.validate_and_convert_non_fake_tensors(
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 1760, in validate_and_convert_non_fake_tensors
    validated_args = [validate(a) for a in flat_args]
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 1760, in <listcomp>
    validated_args = [validate(a) for a in flat_args]
  File "/home/jonch/Desktop/Programming/mlsys/pytorch/torch/_subclasses/fake_tensor.py", line 1748, in validate
    raise AssertionError("Mixing fake modes NYI")
torch._dynamo.exc.BackendCompilerFailed: backend='cudagraphs' raised:
AssertionError: Mixing fake modes NYI
```

### Minified repro

_No response_

### Versions

main

cc @mcarilli @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler