OOM for ThunderFX and Thunder with DDP for Mistral-7B-v0.1

🐛 Bug

When running Mistral-7B-v0.1 we get OOM error. The same configuration passes for torch.compile.

To Reproduce

Steps to reproduce the behavior:

Please use:
1 node(s), each with 8 GPUs.
Image "INTERNAL_IMAGE:pjnl_20241112"
Training script:

python /opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py \
    --model_name Mistral-7B-v0.1 \
    --distributed_mode ddp \
    --shard_mode None \
    --compile dynamo_thunder \
    --checkpoint_activations False \
    --low_precision_mode none  \
    --micro_batch_size 1

Expected behavior

If we can run training with torch.compile we should be able to run it with Thunder as well.

Environment

system.device_product_name DGXH100
system.gpu_driver_version 535.129.03
libraries.cuda 12.6.98.001
libraries.pip.lightning 2.4.0.dev20240728
libraries.pip.lightning-thunder 0.2.0.dev0
libraries.pip.lightning-utilities 0.11.8
libraries.pip.litgpt 0.4.11
libraries.pip.nvfuser 0.2.22+gitba4f7d4
libraries.pip.pytorch-lightning 2.4.0
libraries.pip.torch 2.6.0a0+gita9b4989
libraries.pip.torchao 0.6.1
libraries.pip.torchmetrics 1.5.1
libraries.pip.torchvision 0.19.0a0+d23a6e1

Start by running the provided benchmark_litgpt.py command in the listed DGXH100 environment and compare the dynamo_thunder run with torch.compile. Investigate where the Thunder configuration exceeds memory during Mistral-7B-v0.1 DDP training; done means the Thunder run completes without OOM under the stated settings.