Repository: pytorch
Issue: TorchBench - moco - RuntimeError: Tensors must be CUDA and dense
Labels: triaged,bug

### 🐛 Describe the bug

Run TB model moco
```
benchmarks/dynamo/torchbench.py -d cuda --inductor --training --float32 --accuracy --no-skip --only moco
```

### Error logs

```
cuda train moco                               TorchDynamo optimized model failed to run because of following error
ERROR:common:
from user code:
   File "/scratch/ybliang/work/repos/torchbenchmark/torchbenchmark/models/moco/moco/builder.py", line 172, in <graph break in concat_all_gather>
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
Set torch._dynamo.config.verbose=True for more information
You can suppress this exception and fall back to eager by setting:
    torchdynamo.config.suppress_errors = True
==========
Traceback (most recent call last):
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/variables/tensor.py", line 131, in _get_fake_value
    return wrap_fake_exception(
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/utils.py", line 725, in wrap_fake_exception
    return fn()
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/variables/tensor.py", line 132, in <lambda>
    lambda: _run_node(tx.output, node, args, kwargs, nnmodule)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/variables/tensor.py", line 51, in _run_node
    return node.target(*args, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/distributed/distributed_c10d.py", line 1318, in wrapper
    return func(*args, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/distributed/distributed_c10d.py", line 2313, in all_gather
    work = default_pg.allgather([tensor_list], [tensor])
  File "/scratch/ybliang/work/repos/pytorch/torch/_subclasses/fake_tensor.py", line 894, in __torch_dispatch__
    r = func(*args, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/_ops.py", line 257, in __call__
    return self._op(*args, **kwargs or {})
RuntimeError: Tensors must be CUDA and dense
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
  File "/scratch/ybliang/work/repos/pytorch/benchmarks/dynamo/common.py", line 1095, in check_accuracy
    new_result = optimized_model_iter_fn(model, example_inputs)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/benchmarks/dynamo/common.py", line 999, in run_n_iterations
    self.model_iter_fn(mod, inputs, collect_outputs=False)
  File "/scratch/ybliang/work/repos/pytorch/benchmarks/dynamo/torchbench.py", line 332, in forward_and_backward_pass
    cloned_inputs = clone_inputs(inputs)
  File "/scratch/ybliang/work/repos/pytorch/benchmarks/dynamo/torchbench.py", line 335, in <graph break in forward_and_backward_pass>
    pred = mod(*cloned_inputs)
  File "/scratch/ybliang/work/repos/pytorch/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/nn/parallel/distributed.py", line 1093, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/nn/parallel/distributed.py", line 1047, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])
  File "/scratch/ybliang/work/repos/pytorch/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/scratch/ybliang/work/repos/torchbenchmark/torchbenchmark/models/moco/moco/builder.py", line 130, in forward
    self._momentum_update_key_encoder()  # update the key encoder
  File "/scratch/ybliang/work/repos/torchbenchmark/torchbenchmark/models/moco/moco/builder.py", line 133, in <graph break in forward>
    im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
  File "/scratch/ybliang/work/repos/pytorch/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/scratch/ybliang/work/repos/torchbenchmark/torchbenchmark/models/moco/moco/builder.py", line 76, in _batch_shuffle_ddp
    x_gather = concat_all_gather(x)
  File "/scratch/ybliang/work/repos/pytorch/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/scratch/ybliang/work/repos/torchbenchmark/torchbenchmark/models/moco/moco/builder.py", line 171, in concat_all_gather
    for _ in range(torch.distributed.get_world_size())]
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/eval_frame.py", line 236, in catch_errors
    return callback(frame, cache_size)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/convert_frame.py", line 466, in _convert_frame
    result = inner_convert(frame, cache_size)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/convert_frame.py", line 118, in _fn
    return fn(*args, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/utils.py", line 92, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/convert_frame.py", line 339, in _convert_frame_assert
    return _compile(
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/convert_frame.py", line 394, in _compile
    out_code = transform_code_object(code, transform)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/bytecode_transformation.py", line 341, in transform_code_object
    transformations(instructions, code_options)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/convert_frame.py", line 382, in transform
    tracer.run()
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/symbolic_convert.py", line 1452, in run
    super().run()
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/symbolic_convert.py", line 352, in run
    and self.step()
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/symbolic_convert.py", line 322, in step
    getattr(self, inst.opname)(inst)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/symbolic_convert.py", line 174, in wrapper
    return inner_fn(self, inst)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/symbolic_convert.py", line 807, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/symbolic_convert.py", line 264, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/variables/torch.py", line 381, in call_function
    tensor_variable = TensorVariable.create(
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/variables/tensor.py", line 201, in create
    example_value = _get_fake_value(proxy.node, tx)
  File "/scratch/ybliang/work/repos/pytorch/torch/_dynamo/variables/tensor.py", line 145, in _get_fake_value
    raise TorchRuntimeError() from e
torch._dynamo.exc.TorchRuntimeError: 
from user code:
   File "/scratch/ybliang/work/repos/torchbenchmark/torchbenchmark/models/moco/moco/builder.py", line 172, in <graph break in concat_all_gather>
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
Set torch._dynamo.config.verbose=True for more information
You can suppress this exception and fall back to eager by setting:
    torchdynamo.config.suppress_errors = True
==========
FAIL
```

Look at ```"/scratch/ybliang/work/repos/torchbenchmark/torchbenchmark/models/moco/moco/builder.py", line 172```, the failure op is ```c10d.allgather_.default```. It seems ```moco``` model explicitly calling these ops rather than calling them through DDP. We have to think about how to support such case.

### Minified repro

_No response_