Repository: pytorch
Issue: [Dynamo] Eager fallback casued by graph breaks in module hooks
Labels: triaged,oncall: pt2,module: dynamo

### 🐛 Describe the bug

In deepspeed workloads, Zero parameter offload is implemented via module hooks. We find under torch.compile scenrio, if there any graph breaks happen in the pre/post hook of a module, the module will be fallbacked to the eager mode. Is it expected? It's currently hard to fix all the graph breaks in the hooks and is it possible to make the module stay in fx graph(only fallback the hooks)?

### Error logs

check the dynamo/inductor logs.

### Minified repro

```python
import os
os.environ["TORCH_COMPILE_DEBUG"] = "1"
os.environ["TORCHDYNAMO_VERBOSE"] = "1"
os.environ["TORCH_LOGS"] = "+dynamo,guards,graph_code"
import torch

class MyModule(torch.nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.fc0 = torch.nn.Linear(256, 256, bias=False)
        self.fc1 = torch.nn.Linear(256, 256, bias=False)
        self.dropout = torch.nn.Dropout(0.5)
    def forward(self, data, residual):
        output = residual + self.fc1(self.fc0(self.dropout(data))) * 0.5
        return output


my_mod = MyModule()

@torch._dynamo.disable()
def fn(t):
    return t.sin()

def collect(m, i, o):
  tmp = fn(o)
  return tmp.add(1)

my_mod.register_forward_hook(collect)
for child in my_mod.children():
    child.register_forward_hook(collect)
my_mod = torch.compile(my_mod, fullgraph=False)

atensor = torch.rand(256, 256, dtype=torch.float)
residual = torch.rand(256, 256, dtype=torch.float)
out = my_mod(atensor, residual)


```

### Versions

```
PyTorch version: 2.5.0.dev20240729
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.3.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.30.0
Libc version: N/A

Python version: 3.8.18 (default, Sep 11 2023, 08:17:16)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.3.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] onnx==1.16.1
[pip3] onnx2torch==1.5.14
[pip3] optree==0.11.0
[pip3] pytorch-lightning==1.7.7
[pip3] torch==2.5.0a0+git1f961ad
[pip3] torchaudio==2.4.0.dev20240729
[pip3] torchmetrics==0.10.0
[pip3] torchvision==0.12.0
[conda] numpy                     1.22.3                   pypi_0    pypi
[conda] numpy-base                1.24.3           py38h90707a3_0
[conda] onnx2torch                1.5.14                   pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] pytorch-lightning         1.7.7                    pypi_0    pypi
[conda] torch                     1.12.0.dev20220518          pypi_0    pypi
[conda] torchaudio                2.4.0.dev20240729          pypi_0    pypi
[conda] torchmetrics              0.10.0                   pypi_0    pypi
[conda] torchvision               0.12.0                   pypi_0    pypi
```

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec