Repository: pytorch
Issue: RuntimeError: expect_autograd_hooks_ INTERNAL ASSERT FAILED at "../torch/csrc/distributed/c10d/reducer.cpp"
Labels: oncall: distributed,module: ddp

### 🐛 Describe the bug

When using `DistributedDataParallel` (DDP) with `static_graph=True` and multiple backward passes on the same forward pass within a `no_sync()` context, a runtime error may occur. Specifically, if the very first forward/backward call sequence on the model is made within a `no_sync()` block and involves calling `backward(retain_graph=True)` (also occurs when calling with `retain_graph=False`) on one loss followed by a second backward call on another loss derived from the same forward pass, an internal PyTorch assertion error can be triggered. This issue does not occur if a normal forward/backward pass is performed first (outside of `no_sync()`), and it also does not happen if `no_sync()` is never used.

Run the below scripts by running `torchrun script_name.py`

## Reproduces the error
```python
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp

def func(model):
    im = torch.empty(1, 3, 224, 224, device="cuda")
    seq = torch.randint(0, 1000, (1, 128), device="cuda").long()
    loss, speculative_loss = model(im, seq)
    loss.backward(retain_graph=True)
    speculative_loss.backward()

def worker(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.lin = torch.nn.Linear(3*224*224, 10)

        def forward(self, im, seq):
            out = self.lin(im.flatten(1))
            loss = out.mean()
            return loss, loss

    model = Model().to(rank)
    model = DDP(model, device_ids=[rank], static_graph=True)

    # This scenario triggers the error
    with model.no_sync():
        func(model)

if __name__ == "__main__":
    worker(0, 1)
```

expected output
```bash
[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspaces/data-engine/temp/script_mixed_sync.py", line 39, in <module>
[rank0]:     worker(0, 1)
[rank0]:   File "/workspaces/data-engine/temp/script_mixed_sync.py", line 32, in worker
[rank0]:     func(model)
[rank0]:   File "/workspaces/data-engine/temp/script_mixed_sync.py", line 10, in func
[rank0]:     loss.backward(retain_graph=True)
[rank0]:   File "/workspaces/data-engine/jobs/extractor-training/.venv/lib/python3.10/site-packages/torch/_tensor.py", line 581, in backward
[rank0]:     torch.autograd.backward(
[rank0]:   File "/workspaces/data-engine/jobs/extractor-training/.venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 347, in backward
[rank0]:     _engine_run_backward(
[rank0]:   File "/workspaces/data-engine/jobs/extractor-training/.venv/lib/python3.10/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
[rank0]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]: RuntimeError: expect_autograd_hooks_ INTERNAL ASSERT FAILED at "../torch/csrc/distributed/c10d/reducer.cpp":1603, please report a bug to PyTorch.
```

## Demonstration of it working without `no_sync`
```python
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp

def func(model):
    im = torch.empty(1, 3, 224, 224, device="cuda")
    seq = torch.randint(0, 1000, (1, 128), device="cuda").long()
    loss, speculative_loss = model(im, seq)
    loss.backward(retain_graph=True)
    speculative_loss.backward()

def worker(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.lin = torch.nn.Linear(3*224*224, 10)

        def forward(self, im, seq):
            out = self.lin(im.flatten(1))
            loss = out.mean()
            return loss, loss

    model = Model().to(rank)
    model = DDP(model, device_ids=[rank], static_graph=True)

    # No no_sync() context - works without error
    func(model)

if __name__ == "__main__":
    worker(0, 1)
```

## Demonstration of it working when running  without `no_sync` first

```python
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp

def func(model):
    im = torch.empty(1, 3, 224, 224, device="cuda")
    seq = torch.randint(0, 1000, (1, 128), device="cuda").long()
    loss, speculative_loss = model(im, seq)
    loss.backward(retain_graph=True)
    speculative_loss.backward()

def worker(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.lin = torch.nn.Linear(3*224*224, 10)

        def forward(self, im, seq):
            out = self.lin(im.flatten(1))
            loss = out.mean()
            return loss, loss

    model = Model().to(rank)
    model = DDP(model, device_ids=[rank], static_graph=True)

    func(model)
    with model.no_sync():
        func(model)

if __name__ == "__main__":
    worker(0, 1)
```
[scripts.zip](https://github.com/user-attachments/files/18199631/scripts.zip)



### Versions

Collecting environment information...
PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 12 (bookworm) (x86_64)
GCC version: (Debian 12.2.0-14) 12.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.36

Python version: 3.10.13 (main, Mar 12 2024, 12:22:40) [GCC 12.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-1021-aws-x86_64-with-glibc2.36
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L40S
Nvidia driver version: 560.35.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.5.1
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):                               4
On-line CPU(s) list:                  0-3
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7R13 Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   2
Socket(s):                            1
Stepping:                             1
BogoMIPS:                             5299.99
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 tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save vaes vpclmulqdq rdpid
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            64 KiB (2 instances)
L1i cache:                            64 KiB (2 instances)
L2 cache:                             1 MiB (2 instances)
L3 cache:                             8 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-3
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:   Vulnerable: Safe RET, no microcode
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
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] mypy==1.13.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] torch==2.5.1
[pip3] torchaudio==2.5.1
[pip3] torchvision==0.20.1
[pip3] triton==3.1.0
[conda] Could not collect

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