🐛 [Bug] Compiler error while running sd_unet model

[keehyuna]

Bug Description

Compile error when sd_unet model is running with dynamo backend. No compile error when torch_compile is used.
This error happens after below commit.
chore: Fixes required for LLM models (#3002)

It seems different device types are used in fake tensor.
...
torch._dynamo.exc.TorchRuntimeError: Failed running call_function (*(FakeTensor(..., device='cuda:0', size=(1, 320), dtype=torch.float16), Parameter(FakeTensor(..., size=(1280, 320), dtype=torch.float16, requires_grad=True)), Parameter(FakeTensor(..., size=(1280,), dtype=torch.float16, requires_grad=True))), **{}):
Unhandled FakeTensor Device Propagation for aten.mm.default, found two different devices cuda:0, cpu

To Reproduce

Steps to reproduce the behavior:

1. cd tools/perf
2. python perf_run.py --model_torch sd_unet --precision fp16 --inputs="(1, 4, 64, 64);(1);(1, 1, 768)"
   --batch_size 1 --truncate --backends dynamo --report "output.csv"

Expected behavior

No error.

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Torch-TensorRT Version (e.g. 1.0.0): 10.3.0
• PyTorch Version (e.g. 1.0): 2.5.0.dev20240816+cu124
• CPU Architecture: x86_64
• OS (e.g., Linux): linux
• How you installed PyTorch (conda, pip, libtorch, source): source
• Build command you used (if compiling from source): python setup.py develop
• Are you using local sources or building from archives: local sources, latest main branch
• Python version: 3.10.12
• CUDA version: 12.4
• GPU models and configuration: NVIDIA GeForce RTX 4080, SM Capability: 8.9
• Any other relevant information:

Additional context

Error log
WARNING:root:Running <function run_dynamo at 0x7f1521bb9090> failed
Traceback (most recent call last):
File "/root/trt/TensorRT/tools/perf/perf_run.py", line 39, in wrapper_func
return func(*args, **kwargs)
File "/root/trt/TensorRT/tools/perf/perf_run.py", line 251, in run_dynamo
model = torchtrt.compile(
File "/root/trt/TensorRT/py/torch_tensorrt/_compile.py", line 266, in compile
exp_program = dynamo_trace(
File "/root/trt/TensorRT/py/torch_tensorrt/dynamo/_tracer.py", line 83, in trace
exp_program = export(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/init.py", line 173, in export
return _export(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1069, in wrapper
raise e
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1042, in wrapper
ep = fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/exported_program.py", line 96, in wrapper
return fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 2035, in _export
export_artifact = export_func( # type: ignore[operator]
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1276, in _strict_export
return _strict_export_lower_to_aten_ir(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1304, in _strict_export_lower_to_aten_ir
gm_torch_level = _export_to_torch_ir(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 552, in _export_to_torch_ir
gm_torch_level, _ = torch._dynamo.export(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1437, in inner
result_traced = opt_f(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 469, in _fn
return fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1238, in call
return self._torchdynamo_orig_callable(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 514, in call
return _compile(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 902, in _compile
guarded_code = compile_inner(code, one_graph, hooks, transform)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 653, in compile_inner
return _compile_inner(code, one_graph, hooks, transform)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_utils_internal.py", line 85, in wrapper_function
return StrobelightCompileTimeProfiler.profile_compile_time(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
return func(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 686, in _compile_inner
out_code = transform_code_object(code, transform)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1322, in transform_code_object
transformations(instructions, code_options)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 208, in _fn
return fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 622, in transform
tracer.run()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2730, in run
super().run()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 957, in run
while self.step():
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 869, in step
self.dispatch_table[inst.opcode](self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 557, in wrapper
return inner_fn(self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1564, in CALL_FUNCTION
self.call_function(fn, args, {})
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 804, in call_function
self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type]
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py", line 441, in call_function
return tx.inline_user_function_return(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in inline_user_function_return
return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 2945, in inline_call
return cls.inline_call(parent, func, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 3061, in inline_call
tracer.run()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 957, in run
while self.step():
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 869, in step
self.dispatch_table[inst.opcode](self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 557, in wrapper
return inner_fn(self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1642, in CALL_FUNCTION_EX
self.call_function(fn, argsvars.items, kwargsvars)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 804, in call_function
self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type]
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 383, in call_function
return super().call_function(tx, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 322, in call_function
return super().call_function(tx, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 106, in call_function
return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in inline_user_function_return
return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 2945, in inline_call
return cls.inline_call(parent, func, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 3061, in inline_call
tracer.run()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 957, in run
while self.step():
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 869, in step
self.dispatch_table[inst.opcode](self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 557, in wrapper
return inner_fn(self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1564, in CALL_FUNCTION
self.call_function(fn, args, {})
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 804, in call_function
self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type]
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py", line 441, in call_function
return tx.inline_user_function_return(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in inline_user_function_return
return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 2945, in inline_call
return cls.inline_call(parent, func, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 3061, in inline_call
tracer.run()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 957, in run
while self.step():
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 869, in step
self.dispatch_table[inst.opcode](self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 557, in wrapper
return inner_fn(self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1642, in CALL_FUNCTION_EX
self.call_function(fn, argsvars.items, kwargsvars)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 804, in call_function
self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type]
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 383, in call_function
return super().call_function(tx, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 322, in call_function
return super().call_function(tx, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 106, in call_function
return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in inline_user_function_return
return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 2945, in inline_call
return cls.inline_call(parent, func, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 3061, in inline_call
tracer.run()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 957, in run
while self.step():
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 869, in step
self.dispatch_table[inst.opcode](self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 557, in wrapper
return inner_fn(self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1564, in CALL_FUNCTION
self.call_function(fn, args, {})
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 804, in call_function
self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type]
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 952, in call_function
return self.obj.call_method(tx, self.name, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 157, in call_method
).call_function(tx, [self.objvar] + args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 322, in call_function
return super().call_function(tx, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 106, in call_function
return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in inline_user_function_return
return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 2945, in inline_call
return cls.inline_call(parent, func, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/dynamo/symbolic_convert.py", line 3061, in inline_call
tracer.run()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 957, in run
while self.step():
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 869, in step
self.dispatch_table[inst.opcode](self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 557, in wrapper
return inner_fn(self, inst)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1564, in CALL_FUNCTION
self.call_function(fn, args, {})
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 804, in call_function
self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type]
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/torch.py", line 817, in call_function
tensor_variable = wrap_fx_proxy(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1895, in wrap_fx_proxy
return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1982, in wrap_fx_proxy_cls
example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 2049, in get_fake_value
raise TorchRuntimeError(str(e)).with_traceback(e.traceback) from None
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1981, in get_fake_value
ret_val = wrap_fake_exception(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1466, in wrap_fake_exception
return fn()
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1982, in
lambda: run_node(tx.output, node, args, kwargs, nnmodule)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 2117, in run_node
raise RuntimeError(make_error_message(e)).with_traceback(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 2099, in run_node
return node.target(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/utils/_stats.py", line 21, in wrapper
return fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1251, in torch_dispatch
return self.dispatch(func, types, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1705, in dispatch
return self._cached_dispatch_impl(func, types, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1352, in _cached_dispatch_impl
output = self._dispatch_impl(func, types, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1953, in _dispatch_impl
return decomposition_table[func](*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_prims_common/wrappers.py", line 269, in _fn
result = fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_decomp/decompositions.py", line 83, in inner
r = f(*tree_map(increase_prec, args), **tree_map(increase_prec, kwargs))
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_decomp/decompositions.py", line 1521, in addmm
out = alpha * torch.mm(mat1, mat2)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/utils/_stats.py", line 21, in wrapper
return fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1251, in torch_dispatch
return self.dispatch(func, types, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1705, in dispatch
return self._cached_dispatch_impl(func, types, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1352, in _cached_dispatch_impl
output = self._dispatch_impl(func, types, args, kwargs)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 2031, in _dispatch_impl
self.wrap_meta_outputs_with_default_device_logic(
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 2153, in wrap_meta_outputs_with_default_device_logic
return tree_map(wrap, r)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/utils/_pytree.py", line 964, in tree_map
return treespec.unflatten(map(func, flat_args))
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/utils/_pytree.py", line 803, in unflatten
leaves = list(leaves)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 2131, in wrap
) = FakeTensor._find_common_device(func, flat_args)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 871, in _find_common_device
merge_devices(arg)
File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 866, in merge_devices
raise RuntimeError(
torch._dynamo.exc.TorchRuntimeError: Failed running call_function ((FakeTensor(..., device='cuda:0', size=(1, 320), dtype=torch.float16), Parameter(FakeTensor(..., size=(1280, 320), dtype=torch.float16, requires_grad=True)), Parameter(FakeTensor(..., size=(1280,), dtype=torch.float16, requires_grad=True))), **{}):
Unhandled FakeTensor Device Propagation for aten.mm.default, found two different devices cuda:0, cpu