Cannot use compiled model together with the ddp strategy
Labels: oncall: distributed

### 🐛 Describe the bug

Cannot use compiled model together with the `ddp` strategy, and the error disappears when change the strategy to `auto`.

CMD:
```
python code.py
```

code (`code.py`):
```python
from typing import *
import pytorch_lightning as pl
import torch
import torch.nn as nn
from jsonargparse import lazy_instance
from packaging.version import Version
from torch import Tensor
from torch.nn import MultiheadAttention
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, Dataset


def default_collate_func(batches: List[Tuple[Tensor, Tensor, Dict[str, Any]]]) -> List[Any]:
    mini_batch = []
    for x in zip(*batches):
        if isinstance(x[0], Tensor):
            x = torch.stack(x)
        mini_batch.append(x)
    return mini_batch


class SmsWsjPlusDataset(Dataset):

    def __init__(self,) -> None:
        super().__init__()

    def __getitem__(self, index: int):
        return torch.randn((129, 251, 12)), torch.randn((2, 129, 251)), {'index': index, 'sample_rate': 8000}

    def __len__(self):
        return 100000


class SmsWsjPlusDataModule(LightningDataModule):

    def __init__(self,):
        super().__init__()

    def construct_dataloader(self):
        ds = SmsWsjPlusDataset()
        return DataLoader(ds, batch_size=1, collate_fn=default_collate_func, num_workers=2)

    def train_dataloader(self) -> DataLoader:
        return self.construct_dataloader()

    def val_dataloader(self) -> DataLoader:
        return self.construct_dataloader()


class SpatialNet(nn.Module):

    def __init__(self,):
        super().__init__()
        self.encoder = nn.Conv1d(in_channels=12, out_channels=96, kernel_size=5, stride=1, padding="same")

        layers = []
        for l in range(2):
            layer = nn.ModuleList([nn.LayerNorm(96), MultiheadAttention(embed_dim=96, num_heads=4, batch_first=True)])
            layers.append(layer)
        self.layers = nn.ModuleList(layers)
        self.decoder = nn.Linear(in_features=96, out_features=4)

    def forward(self, x: Tensor, return_attn_score: bool = False) -> Tensor:
        # x: [Batch, Freq, Time, Feature]
        B, F, T, H0 = x.shape
        x = self.encoder(x.reshape(B * F, T, H0).permute(0, 2, 1)).permute(0, 2, 1)
        H = x.shape[2]

        attns = [] if return_attn_score else None
        x = x.reshape(B, F, T, H)
        for m in self.layers:
            x = x.reshape(B * F, T, H)
            x = m[0](x)
            x, attn = m[1].forward(x, x, x)
            x = x.reshape(B, F, T, H)
            if return_attn_score:
                attns.append(attn)

        y = self.decoder(x)
        if return_attn_score:
            return y.contiguous(), attns
        else:
            return y.contiguous()


class TrainModule(pl.LightningModule):

    def __init__(self, arch: nn.Module = lazy_instance(SpatialNet), compile: bool = False):
        super().__init__()

        if compile != False:
            assert Version(torch.__version__) >= Version('2.0.0'), f'compile only works for torch>=2.0: current version: {torch.__version__}'
            self.arch = torch.compile(arch)
        else:
            self.arch = arch

    def forward(self, x: Tensor) -> Tensor:
        B, F, T, H = x.shape
        out = self.arch(x)
        out = torch.view_as_complex(out.float().reshape(B, F, T, -1, 2))  # [B,F,T,Spk]
        out = out.permute(0, 3, 1, 2)  # [B,Spk,F,T]
        return out

    def training_step(self, batch, batch_idx):
        x, ys, paras = batch
        yr_hat = self.forward(x)
        loss = (ys - yr_hat).abs().mean()
        self.log('train/loss', loss, batch_size=ys[0].shape[0], prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, ys, paras = batch
        yr_hat = self.forward(x)
        loss = (ys - yr_hat).abs().mean()
        self.log('val/loss', loss, sync_dist=True, batch_size=ys.shape[0])

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())


if __name__ == '__main__':
    arch = SpatialNet()
    model = TrainModule(arch=arch, compile=True)
    datamodule = SmsWsjPlusDataModule()

    from pytorch_lightning import Trainer
    trainer = Trainer(strategy='ddp', devices="0,") # strategy='auto' is OK
    trainer.fit(model=model, datamodule=datamodule)

```

The error reported:
```
Traceback (most recent call last):
  File "/mnt/home/quancs/projects/NBSS/SharedTrainer2.py", line 128, in <module>
    trainer.fit(model=model, datamodule=datamodule)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 545, in fit
    call._call_and_handle_interrupt(
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 43, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 102, in launch
    return function(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 581, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 990, in _run
    results = self._run_stage()
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1034, in _run_stage
    self._run_sanity_check()
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1063, in _run_sanity_check
    val_loop.run()
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 134, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 391, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_args)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 309, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 402, in validation_step
    return self._forward_redirection(self.model, self.lightning_module, "validation_step", *args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 628, in __call__
    wrapper_output = wrapper_module(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1519, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1355, in _run_ddp_forward
    return self.module(*inputs, **kwargs)  # type: ignore[index]
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 621, in wrapped_forward
    out = method(*_args, **_kwargs)
  File "/mnt/home/quancs/projects/NBSS/SharedTrainer2.py", line 113, in validation_step
    yr_hat = self.forward(x)
  File "/mnt/home/quancs/projects/NBSS/SharedTrainer2.py", line 99, in forward
    out = self.arch(x)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 487, in catch_errors
    return hijacked_callback(frame, cache_entry, hooks, frame_state)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 641, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 133, in _fn
    return fn(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 389, in _convert_frame_assert
    return _compile(
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 569, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 189, in time_wrapper
    r = func(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 491, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 458, in transform
    tracer.run()
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 2074, in run
    super().run()
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 2162, in RETURN_VALUE
    self.output.compile_subgraph(
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 833, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 957, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 189, in time_wrapper
    r = func(*args, **kwargs)
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 1024, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 1009, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/_dynamo/backends/distributed.py", line 268, in compile_fn
    if maybe_param.requires_grad and not self._ignore_parameter(
  File "/mnt/home/quancs/miniconda3/envs/torch211/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
torch._dynamo.exc.BackendCompilerFailed: backend='compile_fn' raised:
AttributeError: 'MultiheadAttention' object has no attribute 'requires_grad'

Set TORCH_LOGS="+dynamo" and 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
```



### Versions

Collecting environment information...
PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 9.3.1 20200408 (Red Hat 9.3.1-2)
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.17

Python version: 3.9.0 (default, Nov 15 2020, 14:28:56)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 530.30.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:          x86_64
CPU op-mode(s):        32-bit, 64-bit
Byte Order:            Little Endian
CPU(s):                128
On-line CPU(s) list:   0-127
Thread(s) per core:    1
Core(s) per socket:    64
Socket(s):             2
NUMA node(s):          8
Vendor ID:             AuthenticAMD
CPU family:            23
Model:                 49
Model name:            AMD EPYC 7742 64-Core Processor
Stepping:              0
CPU MHz:               2250.000
CPU max MHz:           2250.0000
CPU min MHz:           1500.0000
BogoMIPS:              4491.67
Virtualization:        AMD-V
L1d cache:             32K
L1i cache:             32K
L2 cache:              512K
L3 cache:              16384K
NUMA node0 CPU(s):     0-15
NUMA node1 CPU(s):     16-31
NUMA node2 CPU(s):     32-47
NUMA node3 CPU(s):     48-63
NUMA node4 CPU(s):     64-79
NUMA node5 CPU(s):     80-95
NUMA node6 CPU(s):     96-111
NUMA node7 CPU(s):     112-127
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc art rep_good nopl nonstop_tsc extd_apicid aperfmperf eagerfpu pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_l2 cpb cat_l3 cdp_l3 hw_pstate sme retpoline_amd ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip overflow_recov succor smca

Versions of relevant libraries:
[pip3] mypy==1.6.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.3
[pip3] pytorch-lightning==2.1.0
[pip3] torch==2.1.0
[pip3] torchaudio==2.1.0
[pip3] torcheval==0.0.6
[pip3] torchmetrics==1.2.0
[pip3] torchtnt==0.2.1
[pip3] torchvision==0.16.0
[pip3] triton==2.1.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.24.3           py39h14f4228_0  
[conda] numpy-base                1.24.3           py39h31eccc5_0  
[conda] pytorch                   2.1.0           py3.9_cuda12.1_cudnn8.9.2_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-lightning         2.1.0                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.1.0                py39_cu121    pytorch
[conda] torcheval                 0.0.6                    pypi_0    pypi
[conda] torchmetrics              1.2.0                    pypi_0    pypi
[conda] torchtnt                  0.2.1                    pypi_0    pypi
[conda] torchtriton               2.1.0                      py39    pytorch
[conda] torchvision               0.16.0               py39_cu121    pytorch

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu