Repository: pytorch
Issue: torch._inductor.triton_heuristics.cached_autotune is not thread safe / suited for multiprocessing via DDP/FSDP
Labels: oncall: distributed,triaged,module: multithreading,oncall: pt2,module: inductor,pt2d-triage-nov2024

### 🐛 Describe the bug

In multiprocessing mode (i.e. FSDP/DDP), there occur JSONDecodeErrors within torch._inductor.triton_heuristics.cached_autotune, if the filesystem does not lock the file itself.


### Error logs

Excerpts from the stack trace:

```
...lib/python3.11/site-packages/torch/_inductor/triton_heuristics.py", line 740, in cached_autotune
    best_config = load_cached_autotuning(cache_filename, configs_hash, configs)
```


```
.../lib/python3.11/json/decoder.py", line 355, in raw_decode     
    raise JSONDecodeError("Expecting value", s, err.value) from None                                        
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:                                         
JSONDecodeError: Expecting value: line 1 column 1 (char 0)      
```

### Minified repro

```
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.optim import DistributedOptimizer
from torch.optim import Adam

# Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)

def example(rank, world_size):
    # Initialize the distributed environment.
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    
    # Create the model and move it to GPU with the corresponding rank
    model = SimpleModel().to(rank)
    
    # Compile the model with torch.compile for optimization
    model = torch.compile(model, example_inputs=torch.randn(1, 10).to(rank))
    
    # Wrap the model with DDP
    ddp_model = DDP(model, device_ids=[rank])
    
    # Define an optimizer. Here we use Adam, but it can be any optimizer.
    optimizer = Adam(ddp_model.parameters(), lr=0.001)
    
    # Example training loop
    for _ in range(100):  # Assuming 100 iterations
        optimizer.zero_grad()
        outputs = ddp_model(torch.randn(64, 10).to(rank))  # Assuming a batch size of 64
        loss = outputs.mean()
        loss.backward()
        optimizer.step()

def main():
    world_size = torch.cuda.device_count()
    mp.spawn(example,
             args=(world_size,),
             nprocs=world_size,
             join=True)

if __name__ == "__main__":
    main()

```

Code generated by ChatGPT, did not reproduce the error with this as the error will only occur occasionally, e.g. in shared filesystems, where latency is high.

### Versions

Versions of relevant libraries:
[pip3] lion-pytorch==0.1.2
[pip3] numpy==1.26.4
[pip3] pytorch-lightning==2.2.0
[pip3] torch==2.2.0
[pip3] torchdata==0.7.1
[pip3] torchmetrics==1.3.0.post0
[pip3] torchtext==0.17.0
[pip3] torchvision==0.17.0
[pip3] triton==2.2.0
[conda] blas                      1.0                         mkl    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] numpy                     1.26.4          py311h64a7726_0    conda-forge
[conda] pytorch                   2.2.0           py3.11_cuda12.1_cudnn8.9.2_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-lightning         2.2.0                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchdata                 0.7.1                    pypi_0    pypi
[conda] torchmetrics              1.3.0.post0              pypi_0    pypi
[conda] torchtext                 0.17.0                   pypi_0    pypi
[conda] torchtriton               2.2.0                     py311    pytorch
[conda] torchvision               0.17.0              py311_cu121    pytorch

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @albanD @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov @XilunWu @ezyang @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @tianyu-l @msaroufim @bdhirsh @anijain2305 @zou3519 @peterbell10