Repository: pytorch
Issue: [inline-inbuilt-nn-modules] Torch compile with DDP errors on parameterized modules
Labels: high priority,oncall: distributed,triaged,module: ddp,module: nn.utils.parametrize,bug,oncall: pt2,module: dynamo,pt2d-triage-nov2024

### 🐛 Describe the bug

On Pytorch 2.1.0, when there is a parametrized module in a model, torch.compile on the DDP-wrapped module results in an error like:
```
[2023-11-08 17:28:50,112]   File "~/.local/lib/python3.10/site-packages/torch/_dynamo/backends/distributed.py", line 268, in compile_fn
[2023-11-08 17:28:50,112]     if maybe_param.requires_grad and not self._ignore_parameter(
[2023-11-08 17:28:50,112]   File "~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__
[2023-11-08 17:28:50,112]     raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
[2023-11-08 17:28:50,112] torch._dynamo.exc.BackendCompilerFailed: backend='compile_fn' raised:
[2023-11-08 17:28:50,112] AttributeError: 'ParametrizedConv2d' object has no attribute 'requires_grad'
```

The following snippet should reproduce the error:
```Python
from types import SimpleNamespace

import torch
import torch.nn as nn
from torch.nn.parallel.distributed import DistributedDataParallel


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 3, 1, bias=False)
        nn.utils.parametrizations.spectral_norm(self.conv, eps=1e-12)

    def forward(self, input):
        output = self.conv(input)
        return output.sum()


if __name__ == "__main__":
    input = torch.rand((4, 3, 12, 12)).to(device="cuda")
    original_model = Model().to("cuda")

    fake_ddp = SimpleNamespace(bucket_bytes_cap=int(100 * 1024 * 1024))
    DistributedDataParallel._get_active_ddp_module = lambda: fake_ddp
    model = torch.compile(original_model, mode="default")
    out = model(input)
    out.backward()
```

The error happens because Pytorch 2.1.0 generates the following FX graph:
```
ipdb> print(gm.graph)
graph():
    %l_input_ : torch.Tensor [num_users=1] = placeholder[target=L_input_]
    %l__self___conv : [num_users=1] = get_attr[target=L__self___conv]
    %forward : [num_users=1] = call_method[target=forward](args = (%l__self___conv, %l_input_), kwargs = {})
    %sum_1 : [num_users=1] = call_method[target=sum](args = (%forward,), kwargs = {})
    return (sum_1,)
```
The `get_attr`'s target is the `ParametrizedConv2d` object which becomes `maybe_param` in [this line](https://github.com/pytorch/pytorch/blob/1488bafb274fcc82c8aac429bad61738bc3f950e/torch/_dynamo/backends/distributed.py#L285).

While in Pytorch 2.0.1 we get the following graph which gets the weight from the ParametrizedConv2d and uses it with the functional version of Conv2d.
```
ipdb> print(gm.graph)
graph():
    %input_1 : torch.Tensor [#users=1] = placeholder[target=input]
    %weight_weight : [#users=1] = call_module[target=weight_weight](args = (), kwargs = {})
    %conv2d : [#users=1] = call_function[target=torch.conv2d](args = (%input_1, %weight_weight, None, (1, 1), (0, 0), (1, 1), 1), kwargs = {})
    %sum_1 : [#users=1] = call_method[target=sum](args = (%conv2d,), kwargs = {})
    return (sum_1,) 
```

Not quite sure why the FX graph is different between the 2 versions, but the fix may be as simple as adding:
```
if torch.nn.utils.parametrize.is_parametrized(maybe_param):
    maybe_param = maybe_param.parametrizations.weight.original
```
Happy to put up a fix if this sounds good.

### Versions

Ran on Ubuntu 20.04 LTS, Nvidia A100, using Pytorch 2.1.0 built from source.

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @amjames @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @XilunWu @bdhirsh @anijain2305 @aakhundov @kiukchung @lucasllc @tianyu-l