Repository: pytorch
Issue: DDP grads not synced when static_graph=True and module output is a dict subclass?
Labels: oncall: distributed,module: pytree

### 🐛 Describe the bug

I am using something very similar to [`transformers.utils.ModelOutput`](https://github.com/huggingface/transformers/blob/d533465150532b0c5de167b574e59f64c68b1154/src/transformers/utils/generic.py#L237) for my model outputs, and I noticed that when I use DDP with `static_graph=True`, the grads do not appear to be synchronized after `backward()`, but they are if I use a plain `dict` instead. Is there some intricacy of supported DDP outputs that I am missing?

command:
```
CUDA_VISIBLE_DEVICES=0,1 torchrun \
--nproc_per_node=2 \
--nnodes=1 \
--node_rank=0 \
--rdzv_id=462 \
--rdzv_backend=c10d \
a.py
```

**a.py**:
```python
from typing import Dict

import torch
import torch.distributed as dist
from torch import hub, nn


class ModelOutput(dict):
    """Like `transformers.utils.ModelOutput`.

    See: https://github.com/huggingface/transformers/blob/d533465150532b0c5de167b574e59f64c68b1154/src/transformers/utils/generic.py#L237.
    """


class MobileNetBinaryClassifier(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        self.model = hub.load("pytorch/vision:v0.6.0", "mobilenet_v2", pretrained=True, verbose=False)
        self.model.classifier[1] = nn.Linear(1280, num_classes)
        for p in self.model.parameters():
            p.requires_grad = True

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        output = self.model(x)
        # return dict(output=output)  # grads are synced
        return ModelOutput(output=output)  # grads are not synced?


def setup():
    dist.init_process_group(backend="nccl")


def cleanup():
    dist.destroy_process_group()


def demo_basic():
    setup()

    rank = dist.get_rank() if dist.is_initialized() else 0

    model = MobileNetBinaryClassifier().to(rank)
    ddp_model = nn.parallel.DistributedDataParallel(model, device_ids=[rank], static_graph=True)
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()

    data = torch.randn((16, 3, 224, 224), device=torch.device(rank))
    outputs = ddp_model(data)
    outputs["output"].sum().backward()

    print(f"rank{rank}: {ddp_model.module.model.classifier[1].weight.grad[0,:5]}")

    cleanup()


if __name__ == "__main__":
    demo_basic()

```

`return dict(output=output)` (expected):
```
rank0: tensor([ 8.5775,  9.8441, 10.3479,  8.0700,  7.9398], device='cuda:0')
rank1: tensor([ 8.5775,  9.8441, 10.3479,  8.0700,  7.9398], device='cuda:1')
```

`return ModelOutput(output=output)`
```
rank1: tensor([9.3800, 8.2608, 9.2448, 9.0277, 9.6185], device='cuda:1')
rank0: tensor([8.1090, 9.8833, 9.2154, 6.4282, 8.2903], device='cuda:0')
```

### Versions

```console
$ python collect_env.py
Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 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.27.0
Libc version: glibc-2.31

Python version: 3.10.6 (main, Oct 24 2022, 16:07:47) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-76-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000

Nvidia driver version: 525.105.17
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
Address sizes:                   43 bits physical, 48 bits virtual
CPU(s):                          64
On-line CPU(s) list:             0-63
Thread(s) per core:              2
Core(s) per socket:              32
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD Ryzen Threadripper 3970X 32-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         2746.204
CPU max MHz:                     3700.0000
CPU min MHz:                     2200.0000
BogoMIPS:                        7400.62
Virtualization:                  AMD-V
L1d cache:                       1 MiB
L1i cache:                       1 MiB
L2 cache:                        16 MiB
L3 cache:                        128 MiB
NUMA node0 CPU(s):               0-63
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Mitigation; untrained return thunk; SMT enabled with STIBP protection
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, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
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 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 hw_pstate ssbd mba ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 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 arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es

Versions of relevant libraries:
[pip3] facenet-pytorch==2.5.2
[pip3] flytekitplugins-kfpytorch==1.6.2
[pip3] mypy==1.4.1
[pip3] mypy-boto3-s3==1.26.0.post1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.25.2
[pip3] onnx2pytorch==0.4.1
[pip3] torch==2.0.1+cu117
[pip3] torchaudio==2.0.2
[pip3] torchdata==0.6.1
[pip3] torchjpeg==0.9.29
[pip3] torchmetrics==1.0.2
[pip3] torchsnapshot==0.1.0
[pip3] torchvision==0.15.2+cu117
[pip3] triton==2.0.0
[conda] facenet-pytorch           2.5.2                    pypi_0    pypi
[conda] flytekitplugins-kfpytorch 1.6.2                    pypi_0    pypi
[conda] numpy                     1.25.2                   pypi_0    pypi
[conda] onnx2pytorch              0.4.1                    pypi_0    pypi
[conda] torch                     2.0.1+cu117              pypi_0    pypi
[conda] torchaudio                2.0.2                    pypi_0    pypi
[conda] torchdata                 0.6.1                    pypi_0    pypi
[conda] torchjpeg                 0.9.29                   pypi_0    pypi
[conda] torchmetrics              1.0.2                    pypi_0    pypi
[conda] torchsnapshot             0.1.0                    pypi_0    pypi
[conda] torchvision               0.15.2+cu117             pypi_0    pypi
[conda] triton                    2.0.0                    pypi_0    pypi
```

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