torch.compile does not know .half() has side-effects
Labels: triaged, bug, oncall: pt2, module: dynamo

### 🐛 Describe the bug

So torch does not know that calling .half() on a nn.Module will change it's type. Therefore, it will calculate gradients for the fp16 weights in fp32 which will throw an error. We encountered this in MosaicML diffusion here and made a workaround PR here: https://github.com/mosaicml/diffusion/pull/75 .

It does not seem to know that calling:
```
self.vae.half()
```
actually changes the dtype. Fixing it is pretty easy.
```
self.vae = self.vae.half()
```
makes pytorch compile it properly and the backward pass works as expected.

### Error logs

```
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 2239, in _train_batch
    self._train_microbatches(microbatches, total_loss_dict)
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 2339, in _train_microbatches
    microbatch_loss_dict = self._train_microbatch(use_grad_scaling, current_batch_size, is_final_microbatch)
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 2403, in _train_microbatch
    self.state.outputs = self.state.model(self.state.batch)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
  File "/diffusion/diffusion/models/stable_diffusion.py", line 183, in forward
    latents = self.vae.encode(inputs.half())['latent_dist'].sample().data
  File "/diffusion/diffusion/models/stable_diffusion.py", line 224, in <resume in forward>
    return self.unet(noised_latents, timesteps, conditioning,
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 839, in forward
    output = self._fsdp_wrapped_module(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/diffusers/models/unet_2d_condition.py", line 841, in forward
    emb = self.time_embedding(t_emb, timestep_cond)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/diffusers/models/embeddings.py", line 189, in forward
    def forward(self, sample, condition=None):
  File "/usr/lib/python3/dist-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 3905, in forward
    return compiled_fn(full_args)
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1482, in g
    return f(*args)
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 2527, in runtime_wrapper
    all_outs = call_func_with_args(
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1506, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1482, in g
    return f(*args)
  File "/usr/lib/python3/dist-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 3010, in forward
    fw_outs = call_func_with_args(
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1506, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/usr/lib/python3/dist-packages/torch/_dynamo/repro/after_aot.py", line 109, in deferred_for_real_inputs
    return inner_debug_fn(real_inputs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/repro/after_aot.py", line 123, in inner_debug_fn
    copy_tensor_attrs = [
  File "/usr/lib/python3/dist-packages/torch/_dynamo/repro/after_aot.py", line 124, in <listcomp>
    fake_mode.from_tensor(x) if isinstance(x, torch.Tensor) else x
  File "/usr/lib/python3/dist-packages/torch/_subclasses/fake_tensor.py", line 1721, in from_tensor
    return self.fake_tensor_converter(
  File "/usr/lib/python3/dist-packages/torch/_subclasses/fake_tensor.py", line 371, in __call__
    return self.from_real_tensor(
  File "/usr/lib/python3/dist-packages/torch/_subclasses/fake_tensor.py", line 324, in from_real_tensor
    out = self.meta_converter(
  File "/usr/lib/python3/dist-packages/torch/_subclasses/meta_utils.py", line 591, in __call__
    r = self.meta_tensor(
  File "/usr/lib/python3/dist-packages/torch/_subclasses/meta_utils.py", line 307, in meta_tensor
    base = self.meta_tensor(
  File "/usr/lib/python3/dist-packages/torch/_subclasses/meta_utils.py", line 478, in meta_tensor
    r.grad = self.meta_tensor(
RuntimeError: attempting to assign a gradient with dtype 'float' to a tensor with dtype 'c10::Half'. Please ensure that the gradient and the tensor have the same dtype
```

### Minified repro

_No response_

### Versions

PyTorch 2.1.0

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng