Repository: pytorch
Issue:  Forward arguments are not updated in DDP
Labels: oncall: distributed,triaged

### 🐛 Describe the bug

I am using the example from [Initialize DDP with torch.distributed.run/torchrun](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#initialize-ddp-with-torch-distributed-run-torchrun). I add a dict state_info as an additional input to the forward function, which will track the state of each forward call. However, the dict is not updated in DDP. The codes are:
```
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x, state_info):
        state_info['tmp'] = 10 # additional input that tracks the state
        return self.net2(self.relu(self.net1(x)))

def demo_basic():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    print(f"Start running basic DDP example on rank {rank}.")

    # create model and move it to GPU with id rank
    device_id = rank % torch.cuda.device_count()
    model = ToyModel().to(device_id)
    ddp_model = DDP(model, device_ids=[device_id])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    state_info = {'tmp': 0}
    print(state_info) ####### print function 1
    outputs = ddp_model(torch.randn(20, 10), state_info)
    labels = torch.randn(20, 5).to(device_id)
    loss_fn(outputs, labels).backward()
    optimizer.step()
    print(state_info) ####### print function 2, the dict is not updated here

if __name__ == "__main__":
    demo_basic()
```

For the second print function, state_info is not changed, i.e. not changed to 10. When DDP is not used, state_info is changed. It is said that this is not a feature of DDP. Reference: [Forward arguments are not updated in DDP](https://discuss.pytorch.org/t/forward-arguments-are-not-updated-in-ddp/169748/2?u=yzxhd).

### Versions

Collecting environment information...                                                                                                                                                                                           
PyTorch version: 1.9.1                                                                                                                                                                                                          
Is debug build: False                                                                                                                                                                                                           
CUDA used to build PyTorch: 10.2                                                                                                                                                                                                
ROCM used to build PyTorch: N/A                                                                                                                                                                                                 
                                                                                                                                                                                                                                
OS: Ubuntu 20.04.5 LTS (x86_64)                                                                                                                                                                                                 
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0                                                                                                                                                                              
Clang version: Could not collect                                                                                                                                                                                                
CMake version: version 3.21.3
Libc version: glibc-2.31

Python version: 3.8.11 (default, Aug  3 2021, 15:09:35)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-136-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 10.2.89
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 2080 Ti
GPU 1: NVIDIA GeForce GTX 1080 Ti

Nvidia driver version: 470.161.03
cuDNN version: Probably one of the following:
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.2
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] neural-renderer-pytorch==1.1.3
[pip3] numpy==1.23.5
[pip3] pytorch3d==0.6.0
[pip3] torch==1.9.1
[pip3] torch-cluster==1.5.9
[pip3] torch-geometric==2.0.2
[pip3] torch-scatter==2.0.8
[pip3] torch-sparse==0.6.12
[pip3] torchaudio==0.9.0a0+a85b239
[pip3] torchvision==0.10.1
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            11_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            11_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            11_linux64_mkl    conda-forge
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0            py38h7f8727e_0  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.2            py38h51133e4_0  
[conda] neural-renderer-pytorch   1.1.3                    pypi_0    pypi
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] pytorch                   1.9.1           py3.8_cuda10.2_cudnn7.6.5_0    pytorch
[conda] pytorch3d                 0.6.0                    pypi_0    pypi
[conda] torch-cluster             1.5.9                    pypi_0    pypi
[conda] torch-geometric           2.0.2                    pypi_0    pypi
[conda] torch-scatter             2.0.8                    pypi_0    pypi
[conda] torch-sparse              0.6.12                   pypi_0    pypi
[conda] torchaudio                0.9.1                      py38    pytorch
[conda] torchvision               0.10.1               py38_cu102    pytorch


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