Repository: pytorch
Issue: torch.jit.trace can not trace buffer by Module.register_buffer() when use DDP Module.
Labels: oncall: distributed

### 🐛 Describe the bug

I want to use Tensorboard add_graph func to draw the network structure, but when my module have a buffer by Module.register_buffer() setting, the add_graph func raise an error and tell me "Modules that are called during a trace must be registered as submodules of the thing being traced."

But if I don't use DDP model, I don't have this error.
At the same time, if my network has no buffer, there will be no error if I use ddp.

```python
import torch
from torch import Tensor, nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch import distributed, nn
import os
import torch
import torch.multiprocessing as mp
from torch import distributed
import torch
from torch.utils.tensorboard import SummaryWriter


class Demo(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        A = torch.rand(4, 4, 4)

        self.register_buffer('A', A)

        self.conv = nn.Conv2d(3, 32, 3)

    def forward(self, x: Tensor):
        x = x.mul(self.A)
        x = self.conv(x)
        return x


def parallel_network(network: nn.Module) -> DDP:

    distributed.barrier()

    network_ = nn.SyncBatchNorm.convert_sync_batchnorm(network)
    network_ = DDP(network_,
                   device_ids=[0],
                   output_device=0)
    return network_


def main(local_rank):
    torch.cuda.set_device(local_rank)

    if distributed.is_nccl_available():
        backend = 'nccl'
    else:
        backend = 'gloo'

    distributed.init_process_group(
        backend=backend,
        init_method='env://',
        world_size=torch.cuda.device_count(),
        rank=local_rank
    )

    distributed.barrier()

    demo = Demo()
    demo = parallel_network(demo.cuda())
    writer = SummaryWriter()
    input_ = torch.rand(2, 3, 4, 4).cuda()
    writer.add_graph(demo, input_)


if __name__ == '__main__':

    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '8880'
    mp.spawn(main,  nprocs=torch.cuda.device_count())
```

Run the code and you'll get an error like this:
```
RuntimeError: Tried to trace <__torch__.torch.classes.c10d.ProcessGroup object at 000002DD0FA310E0> but it is not part of the active trace. Modules that are called during a trace must be registered as submodules of the thing being traced.
```

Specifically, the error is raised by the line.
The specific location is : torch\nn\parallel\distributed.py", line 1689.

```
        dist._broadcast_coalesced(
            self.process_group, tensors, buffer_size, authoritative_rank
        )
```

### Versions

Collecting environment information...
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 专业版
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.2 | packaged by Anaconda, Inc. | (main, Mar 27 2023, 23:35:04) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1060
Nvidia driver version: 512.78
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2808
DeviceID=CPU0
Family=198
L2CacheSize=1024
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2808
Name=Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] torch==2.0.0+cu117
[pip3] torchaudio==2.0.1+cu117
[pip3] torchvision==0.15.1+cu117
[conda] numpy                     1.24.1                   pypi_0    pypi
[conda] torch                     2.0.0+cu117              pypi_0    pypi
[conda] torchaudio                2.0.1+cu117              pypi_0    pypi
[conda] torchvision               0.15.1+cu117             pypi_0    pypi


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