Repository: pytorch
Issue: Torch compile OpenMP error with simple nn.Module
Labels: triaged,module: macos,module: openmp,has workaround,oncall: pt2

### 🐛 Describe the bug

I'm running the following program on python3.11.
It fails with both torch2.2 and torch2.3, but with different errors.

```python
import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.w1 = nn.Parameter(torch.empty(dim, dim))

    def forward(self, x: torch.Tensor):
        x = x[:, None, :] * self.w1
        x = torch.relu(x).sum(-1)
        return x

def test(batch_size, dim):
    print(batch_size, dim)
    x = torch.randn(batch_size, dim)
    func = torch.compile(Net(dim))
    result = func(x)
    result.sum().backward()

for dim in [2**i for i in range(10)]:
    test(1, dim)
```

With torch2.2.2 this fails with the error
```
[2024-06-04 15:45:00,474] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (8)
[2024-06-04 15:45:00,474] torch._dynamo.convert_frame: [WARNING]    function: 'forward' (/Users/ahle/kode/cuda_test/repro.py:9)
[2024-06-04 15:45:00,474] torch._dynamo.convert_frame: [WARNING]    last reason: ___guarded_code.valid
[2024-06-04 15:45:00,474] torch._dynamo.convert_frame: [WARNING] To log all recompilation reasons, use TORCH_LOGS="recompiles".
[2024-06-04 15:45:00,474] torch._dynamo.convert_frame: [WARNING] To diagnose recompilation issues, see https://pytorch.org/docs/master/compile/troubleshooting.html.
1 512
```

I upgraded to torch2.3.0, but now it fails with 
```
OMP: Error #15: Initializing libomp.dylib, but found libomp.dylib already initialized.
OMP: Hint This means that multiple copies of the OpenMP runtime have been linked into the program. That is dangerous, since it can degrade performance or cause incorrect results. The best thing to do is to ensure that only a single OpenMP runtime is linked into the process, e.g. by avoiding static linking of the OpenMP runtime in any library. As an unsafe, unsupported, undocumented workaround you can set the environment variable KMP_DUPLICATE_LIB_OK=TRUE to allow the program to continue to execute, but that may cause crashes or silently produce incorrect results. For more information, please see http://openmp.llvm.org/
Abort trap: 6
/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
```

I tried to disable to duplicate library check with `KMP_DUPLICATE_LIB_OK=True`, but even then it still fails with
```
OMP: Error #179: Function pthread_mutex_init failed:
OMP: System error #22: Invalid argument
Segmentation fault: 11
/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
```

### Versions

Collecting environment information...
PyTorch version: N/A
Is debug build: N/A
CUDA used to build PyTorch: N/A
ROCM used to build PyTorch: N/A

OS: macOS 13.4.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.28.1
Libc version: N/A

Python version: 3.11.5 (main, Sep 11 2023, 08:31:25) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.4.1-arm64-arm-64bit
Is CUDA available: N/A
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: N/A

CPU:
Apple M2 Pro

Versions of relevant libraries:
[pip3] No relevant packages
[conda] No relevant packages


cc @malfet @albanD @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang