Weight tying + FSDP = out of bounds

@kshitij12345 is already working on this.

Since May 23, 2024.

## 🐛 Bug

### To Reproduce

Code:

```python
import os
import torch
import torch.distributed as tdist
import thunder
from thunder.tests.litgpt_model import GPT, Config

if __name__ == "__main__":
tdist.init_process_group(backend="nccl")
LOCAL_RANK = int(os.environ["LOCAL_RANK"])
device = torch.device("cuda", LOCAL_RANK)
torch.set_default_device(device)

config = Config(block_size=256, padded_vocab_size=32000, n_layer=1, n_head=3, head_size=24, n_embd=144, rotary_percentage=1.0, parallel_residual=False, bias=False, norm_class_name='RMSNorm', mlp_class_name='LLaMAMLP', intermediate_size=384)
with device:
model = GPT(config)

model.transformer.wte.weight = model.lm_head.weight

model = thunder.distributed.fsdp(model)
model = thunder.jit(model, executors=["torch"])

input_ids = torch.randint(1, 30010, (128, 256), dtype=torch.long, device=device)
logits = model(input_ids)
print(logits.shape)
```

Run with:

```shell
CUDA_LAUNCH_BLOCKING=1 torchrun --nproc-per-node 2 --local-ranks-filter 0 repro.py
```

Error:

```text
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [86,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [87,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [88,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [89,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [90,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [91,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [92,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [93,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [94,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1289: indexSelectLargeIndex: block: [313,0,0], thread: [95,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
[rank0]:[E506 07:38:19.598156204 ProcessGroupNCCL.cpp:1432] [PG 0 (default_pg) Rank 0] Process group watchdog thread terminated with exception: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Exception raised from c10_cuda_check_implementation at ../c10/cuda/CUDAException.cpp:43 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7f7d6f779017 in /home/carlos/nightly-env/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7f7d6f728cd3 in /home/carlos/nightly-env/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x118 (0x7f7d6fb791f8 in /home/carlos/nightly-env/lib/python3.10/site-packages/torch/lib/libc10_cuda.so)
frame #3: c10d::ProcessGroupNCCL::WorkNCCL::finishedGPUExecutionInternal() const + 0x56 (0x7f7d22126926 in /home/carlos/nightly-env/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #4: c10d::ProcessGroupNCCL::WorkNCCL::isCompleted() + 0x58 (0x7f7d2212b2a8 in /home/carlos/nightly-env/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #5: c10d::ProcessGroupNCCL::watchdogHandler() + 0x1de (0x7f7d221322de in /home/carlos/nightly-env/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #6: c10d::ProcessGroupNCCL::ncclCommWatchdog() + 0x10c (0x7f7d221341bc in /home/carlos/nightly-env/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so)
frame #7: + 0xdc253 (0x7f7d6eeb0253 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #8: + 0x94ac3 (0x7f7d70986ac3 in /lib/x86_64-linux-gnu/libc.so.6)
frame #9: + 0x126850 (0x7f7d70a18850 in /lib/x86_64-linux-gnu/libc.so.6)

terminate called after throwing an instance of 'c10::DistBackendError'
what(): [PG 0 (default_pg) Rank 0] Process group watchdog thread terminated with exception: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

Removing one of:
- FSDP
- A high enough input_ids value (30010 in the example)
- weight tying

makes the problem not appear

cc @carmocca @awaelchli @crcrpar