torch.compile FakeTensor tracing fails with foreach ops with multiple devices
Labels: module: optimizer, triaged, module: custom-operators, release notes: foreach_frontend, oncall: pt2, module: fakeTensor

### 🐛 Describe the bug

Foreach ops support multiple devices in eager, but once we try to compile them they no longer support this behavior.

ex.
```py
import torch


@torch.compile()
def test_foreach_add(a0, a1, b0, b1):
    return torch._foreach_add([a0, a1], [b0, b1])


print(
    test_foreach_add(
        torch.ones(10, 10, device="cuda"),
        torch.ones(20, 20, device="cpu"),
        torch.zeros(10, 10, device="cuda"),
        torch.zeros(20, 20, device="cpu"),
    )
)```

removing compile fixes the issue.

### Error logs

`--> python foreach_ex.py
Traceback (most recent call last):
  File "/scratch/mlazos/test/foreach_ex.py", line 10, in <module>
    test_foreach_add(
  File "/scratch/mlazos/pytorch/torch/_dynamo/eval_frame.py", line 286, in _fn
    return fn(*args, **kwargs)
  File "/scratch/mlazos/pytorch/torch/_dynamo/eval_frame.py", line 439, in catch_errors
    return callback(frame, cache_size, hooks, frame_state)
  File "/scratch/mlazos/pytorch/torch/_dynamo/convert_frame.py", line 522, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/scratch/mlazos/pytorch/torch/_dynamo/convert_frame.py", line 125, in _fn
    return fn(*args, **kwargs)
  File "/scratch/mlazos/pytorch/torch/_dynamo/convert_frame.py", line 358, in _convert_frame_assert
    return _compile(
  File "/scratch/mlazos/pytorch/torch/_dynamo/utils.py", line 177, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/mlazos/pytorch/torch/_dynamo/convert_frame.py", line 428, in _compile
    out_code = transform_code_object(code, transform)
  File "/scratch/mlazos/pytorch/torch/_dynamo/bytecode_transformation.py", line 1000, in transform_code_object
    transformations(instructions, code_options)
  File "/scratch/mlazos/pytorch/torch/_dynamo/convert_frame.py", line 413, in transform
    tracer.run()
  File "/scratch/mlazos/pytorch/torch/_dynamo/symbolic_convert.py", line 2009, in run
    super().run()
  File "/scratch/mlazos/pytorch/torch/_dynamo/symbolic_convert.py", line 703, in run
    and self.step()
  File "/scratch/mlazos/pytorch/torch/_dynamo/symbolic_convert.py", line 663, in step
    getattr(self, inst.opname)(inst)
  File "/scratch/mlazos/pytorch/torch/_dynamo/symbolic_convert.py", line 385, in wrapper
    return inner_fn(self, inst)
  File "/scratch/mlazos/pytorch/torch/_dynamo/symbolic_convert.py", line 1094, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/scratch/mlazos/pytorch/torch/_dynamo/symbolic_convert.py", line 554, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/scratch/mlazos/pytorch/torch/_dynamo/variables/torch.py", line 565, in call_function
    tensor_variable = wrap_fx_proxy(
  File "/scratch/mlazos/pytorch/torch/_dynamo/variables/builder.py", line 1003, in wrap_fx_proxy
    return wrap_fx_proxy_cls(
  File "/scratch/mlazos/pytorch/torch/_dynamo/variables/builder.py", line 1038, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx)
  File "/scratch/mlazos/pytorch/torch/_dynamo/utils.py", line 1288, in get_fake_value
    raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
  File "/scratch/mlazos/pytorch/torch/_dynamo/utils.py", line 1256, in get_fake_value
    return wrap_fake_exception(
  File "/scratch/mlazos/pytorch/torch/_dynamo/utils.py", line 850, in wrap_fake_exception
    return fn()
  File "/scratch/mlazos/pytorch/torch/_dynamo/utils.py", line 1257, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
  File "/scratch/mlazos/pytorch/torch/_dynamo/utils.py", line 1322, in run_node
    raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e
  File "/scratch/mlazos/pytorch/torch/_dynamo/utils.py", line 1309, in run_node
    return node.target(*args, **kwargs)
  File "/scratch/mlazos/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/scratch/mlazos/pytorch/torch/_subclasses/fake_tensor.py", line 1105, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/scratch/mlazos/pytorch/torch/_subclasses/fake_tensor.py", line 1321, in dispatch
    return self.wrap_meta_outputs_with_default_device_logic(r, func, args, kwargs)
  File "/scratch/mlazos/pytorch/torch/_subclasses/fake_tensor.py", line 1385, in wrap_meta_outputs_with_default_device_logic
    return tree_map(partial(wrap), r)
  File "/scratch/mlazos/pytorch/torch/utils/_pytree.py", line 196, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/scratch/mlazos/pytorch/torch/utils/_pytree.py", line 196, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/scratch/mlazos/pytorch/torch/_subclasses/fake_tensor.py", line 1406, in wrap
    ) = FakeTensor._find_common_device(func, args, kwargs)
  File "/scratch/mlazos/pytorch/torch/_subclasses/fake_tensor.py", line 1041, in _find_common_device
    tree_map(merge_devices, args)
  File "/scratch/mlazos/pytorch/torch/utils/_pytree.py", line 196, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/scratch/mlazos/pytorch/torch/utils/_pytree.py", line 196, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/scratch/mlazos/pytorch/torch/_subclasses/fake_tensor.py", line 1037, in merge_devices
    raise RuntimeError(
torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in method _foreach_add of type object at 0x7fea51669688>(*([FakeTensor(..., device='cuda:0', size=(10, 10)), FakeTensor(..., size=(20, 20))], [FakeTensor(..., device='cuda:0', size=(10, 10)), FakeTensor(..., size=(20, 20))]), **{}):
Unhandled FakeTensor Device Propagation for aten._foreach_add.List, found two different devices cuda:0, cpu

from user code:
   File "/scratch/mlazos/test/foreach_ex.py", line 6, in test_foreach_add
    return torch._foreach_add([a0, a1], [b0, b1])

Set torch._dynamo.config.verbose=True or TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

### Minified repro

_No response_

### Versions

0a7ea9627f087afef5c59b2500e018abb6c3c1b5

cc @vincentqb @jbschlosser @albanD @janeyx99 @ezyang @msaroufim @wconstab @ngimel @bdhirsh @anijain2305