Repository: pytorch
Issue: Compiled and masked MHA running with DDP and `no_grad()` leads to Graph break due to unsupported builtin torch._C._distributed_c10d.PyCapsule._broadcast_coalesced
Labels: oncall: distributed,module: nn,triaged

### 🐛 Describe the bug

As titled, calling `torch.nn.MultiheadAttention` with `attn_mask` within a compiled model wrapped in DDP and `no_grad()` leads to the graph break warning. I track the issue down to this gist: https://gist.github.com/EIFY/dee8c28e1f5c7d78388020abd36a627e

```sh
cd ~/Downloads/dee8c28e1f5c7d78388020abd36a627e
torchrun mha_mask_ddp.py --multiprocessing-distributed --compiled --mask
Use GPU: 2 for training
Use GPU: 5 for training
Use GPU: 6 for training
Use GPU: 3 for training
Use GPU: 7 for training
Use GPU: 4 for training
Use GPU: 1 for training
Use GPU: 0 for training
Train-like step with gradient: # repeated total of 8 times
# (...)
Eval-like step without gradient: # repeated and interleaved total of 8 times
/home/ubuntu/.local/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py:725: UserWarning: Graph break due to unsupported builtin torch._C._distributed_c10d.PyCapsule._broadcast_coalesced. This function is either a Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python extension (perhaps created with pybind). If it is a Python builtin, please file an issue on GitHub so the PyTorch team can add support for it and see the next case for a workaround. If it is a third-party C/C++ Python extension, please either wrap it into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html for more details) or, if it is traceable, use torch.compiler.allow_in_graph.
  torch._dynamo.utils.warn_once(msg)
# (...)
# No such warnings unless both --compiled and --mask flags are set:
torchrun mha_mask_ddp.py --multiprocessing-distributed --compiled
torchrun mha_mask_ddp.py --multiprocessing-distributed --mask
torchrun mha_mask_ddp.py --multiprocessing-distributed
```
The warning is very specific and goes away if I call MHA without `attn_mask` or avoid compiled model for the eval-like step with the `no_grad()` context. I don't know how serious this is since it's just a warning, but it seems unintended regardless.

### Versions

The output below shows PyTorch version: 2.3.1 and torchvision==0.18.1, but I am actually running
```
$ pip freeze
(...)
torch==2.5.1
torchaudio==2.5.1
torchvision==0.20.1
(...)
```
I don't know why `python3 -mpip list --format=freeze` finds the old packages.

```
Collecting environment information...
PyTorch version: 2.3.1
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 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.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.8.0-47-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 550.90.12
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
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               124
On-line CPU(s) list:                  0-123
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7542 32-Core Processor
CPU family:                           23
Model:                                49
Thread(s) per core:                   1
Core(s) per socket:                   1
Socket(s):                            124
Stepping:                             0
BogoMIPS:                             5800.00
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr wbnoinvd arat npt nrip_save umip rdpid arch_capabilities
Virtualization:                       AMD-V
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            7.8 MiB (124 instances)
L1i cache:                            7.8 MiB (124 instances)
L2 cache:                             62 MiB (124 instances)
L3 cache:                             1.9 GiB (124 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-31,64-95
NUMA node1 CPU(s):                    32-63,96-123
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:               Mitigation; untrained return thunk; SMT disabled
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; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] flake8==4.0.1
[pip3] numpy==1.21.5
[pip3] optree==0.12.1
[pip3] torch==2.3.1
[pip3] torchvision==0.18.1
[pip3] triton==2.3.1
[conda] Could not collect
```

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @drisspg