Timeout for Platypus-30B and Thunder compile

## 🐛 Bug

With newest version of Docker image (tested on 2024-04-28 ) training with thunder.jit on 8xA100 it's not possible to run Platypus-30B and vicuna-33b-v1.3 models. This is the error:

> Time to instantiate model: 0.05 seconds.
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
Time to instantiate model: 0.05 seconds.
Time to instantiate model: 0.05 seconds.
Time to instantiate model: 0.05 seconds.
Time to instantiate model: 0.05 seconds.
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
/usr/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
self.pid = os.fork()
[rank6]:[W429 12:49:14.458854201 ProcessGroupNCCL.cpp:1113] WARNING: process group has NOT been destroyed before it is being destructed. On normal program exit, the application should call destroy_process_group to ensure that any pending NCCL data transfers have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. This constraint has always been present, but this warning has only been added since PyTorch 2.4
[rank3]:[E429 12:59:12.174771056 ProcessGroupNCCL.cpp:568] [Rank 3] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600059 milliseconds before timing out.
[rank5]:[E429 12:59:12.230443027 ProcessGroupNCCL.cpp:568] [Rank 5] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600085 milliseconds before timing out.
[rank1]:[E429 12:59:12.305260726 ProcessGroupNCCL.cpp:568] [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600071 milliseconds before timing out.
[rank4]:[E429 12:59:12.307015699 ProcessGroupNCCL.cpp:568] [Rank 4] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600032 milliseconds before timing out.
[rank7]:[E429 12:59:12.327170052 ProcessGroupNCCL.cpp:568] [Rank 7] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600097 milliseconds before timing out.
[rank2]:[E429 12:59:12.356285062 ProcessGroupNCCL.cpp:568] [Rank 2] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600025 milliseconds before timing out.
[rank0]:[E429 12:59:12.369438320 ProcessGroupNCCL.cpp:568] [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600021 milliseconds before timing out.
[rank4]:[E429 12:59:13.082814283 ProcessGroupNCCL.cpp:1602] [PG 0 (default_pg) Rank 4] Timeout at NCCL work: 721, last enqueued NCCL work: 722, last completed NCCL work: 720.
[rank4]:[E429 12:59:13.082874197 ProcessGroupNCCL.cpp:582] [Rank 4] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank4]:[E429 12:59:13.082891961 ProcessGroupNCCL.cpp:588] [Rank 4] To avoid data inconsistency, we are taking the entire process down.
[rank4]:[E429 12:59:13.082929112 ProcessGroupNCCL.cpp:1432] [PG 0 (default_pg) Rank 4] Process group watchdog thread terminated with exception: [Rank 4] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=721, OpType=_ALLGATHER_BASE, NumelIn=14909440, NumelOut=119275520, Timeout(ms)=600000) ran for 600032 milliseconds before timing out.
Exception raised from checkTimeout at /opt/pytorch/pytorch/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:570 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string, std::allocator >) + 0xae (0x7f5c2b47d84e in /usr/local/lib/python3.10/dist-packages/torch/lib/libc10.so)

### To Reproduce

Before each testing each compilation method I restarted the container:
```
mkdir -p output
docker run --pull=always --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v $PWD/output:/output -it INTERNAL_IMAGE:pjnl-20240427
```
**Thunder**
* Timeout error
* Command:
```
torchrun --nproc-per-node=8 /opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py --model_name Platypus-30B --compile thunder --distributed_mode fsdp --shard_mode zero3

```

**Inductor**
* Works fine
* Command:
```
torchrun --nproc-per-node=8 /opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py --model_name Platypus-30B --compile inductor --distributed_mode fsdp --shard_mode zero3
```

### Expected behavior

If we can run model using Torch Inductor we should be able to use it Thunder as well.

### Environment

As in the Docker image. This results come from single A100.
nvidia-smi output:

cc @carmocca @awaelchli @crcrpar