Repository: pytorch
Issue: Bug when using reparameterized model evaluating with DDP
Labels: oncall: distributed,triaged

### 🐛 Describe the bug

I usually evaluate my model at the end of each training epoch. But in DDP mode,  validation using the reparameterized model reports an error. If a single GPU is used for training and validation neither will report an error.
The error message is given below:
```triple quotes blocks
BUG INFO
[rank0]: Traceback (most recent call last):
[rank0]:   File "run.py", line 97, in <module>
[rank0]:     main()
[rank0]:   File "run.py", line 92, in main
[rank0]:     cli.train()
[rank0]:   File "/home/xuexufeng/project/psf_cv_framework/trainer/mobileone.py", line 37, in train
[rank0]:     self.val_one_epoch()
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/home/xuexufeng/project/psf_cv_framework/trainer/mobileone.py", line 96, in val_one_epoch
[rank0]:     self.val_one_batch(batch)
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/home/xuexufeng/project/psf_cv_framework/trainer/mobileone.py", line 79, in val_one_batch
[rank0]:     predict_cls: torch.Tensor = self.model_eval(batch)["cls"]
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1589, in forward
[rank0]:     inputs, kwargs = self._pre_forward(*inputs, **kwargs)
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1487, in _pre_forward
[rank0]:     self._sync_buffers()
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 2129, in _sync_buffers
[rank0]:     self._sync_module_buffers(authoritative_rank)
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 2133, in _sync_module_buffers
[rank0]:     self._default_broadcast_coalesced(authoritative_rank=authoritative_rank)
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 2155, in _default_broadcast_coalesced
[rank0]:     self._distributed_broadcast_coalesced(bufs, bucket_size, authoritative_rank)
[rank0]:   File "/home/xuexufeng/miniconda/envs/py3.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 2070, in _distributed_broadcast_coalesced
[rank0]:     dist._broadcast_coalesced(
[rank0]: RuntimeError: !tensors.empty() INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1716905971873/work/torch/csrc/distributed/c10d/reducer.cpp":2090, please report a bug to PyTorch. 
```

A sample instance as follow:

```Python
import os
import sys
import torch
import torch.nn as nn
from torch.utils import data


class MyTrainer:
    def __init__(self):
        self.rank, self.num_gpus = get_dist_info()
        self.model: nn.Module = Mobileone("s0")
        self.val_data_loader: data.DataLoader = MyDataloader()

    @staticmethod
    def reparameterize_model(model: nn.Module) -> torch.nn.Module:
        """ Refer: Mobileone https://github.com/apple/ml-mobileone
        
            Method returns a model where a multi-branched structure
            used in training is re-parameterized into a single branch
            for inference.

        :param model: MobileOne model in train mode.
        :return: MobileOne model in inference mode.
        """
        # Avoid editing original graph
        model_local = deepcopy(model)
        for module in model_local.modules():
            if hasattr(module, 'reparameterize'):
                module.reparameterize()
        return model_local

    def train(self,):
        # Assuming the training process is complete

        # Evaluating only in rank-0
        if self.rank == 0:
            self.val_one_epoch()

    @torch.no_grad()
    def val_one_epoch(self, ):
        self.model.eval()
        self.model_eval: nn.Module = self.reparameterize_model(self.model)
        for _, batch in enumerate(self.val_data_loader):
            self.val_one_batch(batch)

    def val_one_batch(self, data):
        predict_cls: torch.Tensor = self.model_eval(data)
```

### Versions

PyTorch version: 2.3.1
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.8.19 (default, Mar 20 2024, 19:58:24)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-125-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.5.119
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090
GPU 3: NVIDIA GeForce RTX 3090
GPU 4: NVIDIA GeForce RTX 3090
GPU 5: NVIDIA GeForce RTX 3090
GPU 6: NVIDIA GeForce RTX 3090
GPU 7: NVIDIA GeForce RTX 3090

Nvidia driver version: 550.54.14
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.0.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.0.0
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn.so.9
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn_adv.so.9
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn_cnn.so.9
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn_engines_precompiled.so.9
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn_engines_runtime_compiled.so.9
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn_graph.so.9
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn_heuristic.so.9
/usr/local/cuda-12.4/targets/x86_64-linux/lib/libcudnn_ops.so.9
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
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               128
On-line CPU(s) list:                  0-127
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7543 32-Core Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   32
Socket(s):                            2
Stepping:                             1
Frequency boost:                      enabled
CPU max MHz:                          3737.8899
CPU min MHz:                          1500.0000
BogoMIPS:                             5599.44
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 rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid 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_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm
Virtualization:                       AMD-V
L1d cache:                            2 MiB (64 instances)
L1i cache:                            2 MiB (64 instances)
L2 cache:                             32 MiB (64 instances)
L3 cache:                             512 MiB (16 instances)
NUMA node(s):                         8
NUMA node0 CPU(s):                    0-7,64-71
NUMA node1 CPU(s):                    8-15,72-79
NUMA node2 CPU(s):                    16-23,80-87
NUMA node3 CPU(s):                    24-31,88-95
NUMA node4 CPU(s):                    32-39,96-103
NUMA node5 CPU(s):                    40-47,104-111
NUMA node6 CPU(s):                    48-55,112-119
NUMA node7 CPU(s):                    56-63,120-127
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Mitigation; safe RET
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] onnx==1.16.1
[pip3] onnxruntime==1.18.1
[pip3] onnxsim==0.4.36
[pip3] torch==2.3.1
[pip3] torchaudio==2.3.1
[pip3] torchvision==0.18.1
[pip3] triton==2.3.1
[conda] blas                      1.0                         mkl    conda-forge
[conda] cuda-cudart               11.8.89                       0    nvidia
[conda] cuda-cupti                11.8.87                       0    nvidia
[conda] cuda-libraries            11.8.0                        0    nvidia
[conda] cuda-nvrtc                11.8.89                       0    nvidia
[conda] cuda-nvtx                 11.8.86                       0    nvidia
[conda] cuda-runtime              11.8.0                        0    nvidia
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] libcublas                 11.11.3.6                     0    nvidia
[conda] libcufft                  10.9.0.58                     0    nvidia
[conda] libcurand                 10.3.6.82                     0    nvidia
[conda] libcusolver               11.4.1.48                     0    nvidia
[conda] libcusparse               11.7.5.86                     0    nvidia
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            16_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           hc2b9512_224  
[conda] numpy                     1.24.4           py38h59b608b_0    conda-forge
[conda] pytorch                   2.3.1           py3.8_cuda11.8_cudnn8.7.0_0    pytorch
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.3.1                py38_cu118    pytorch
[conda] torchtriton               2.3.1                      py38    pytorch
[conda] torchvision               0.18.1               py38_cu118    pytorch

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o