Different shapes, values of model weights and losses between FSDP training in Eager mode and with Thunder

🐛 Bug

After training Llama-3-8b on 8 A100 for 10 iterations with eager mode I printed the model weights:

torch_dist.barrier()
weights_after_training = benchmark.model.lm_head.weight[:10].data.to(device="cpu", dtype=torch.float32).numpy()
if global_rank in [0, None]:
    print(f"WEIGHTS:\n{weights_after_training.shape}\n{weights_after_training}")

when not using Thunder I got:

WEIGHTS:
(10,)
[ 0.01855469 0.00598145 0.01312256 0.01300049 0.00262451 0.0055542
-0.01104736 0.00076294 0.01202393 -0.00909424]

when using Thunder I got:

WEIGHTS:
(10, 4096)
[[ 0.01281738 0.00582886 0.01342773 ... -0.01196289 -0.00369263
-0.01287842]
[-0.00331116 0.01647949 -0.01452637 ... -0.01696777 -0.00650024
-0.00145721]
[ 0.01867676 0.00334167 0.00133514 ... -0.00531006 -0.00744629
0.01147461]
...
[-0.01019287 -0.00939941 0.00204468 ... 0.01184082 0.00201416
-0.01104736]
[-0.00643921 0.00318909 0.01623535 ... -0.00148773 0.01153564
-0.01086426]
[-0.00921631 -0.01452637 0.01586914 ... -0.01330566 0.00445557
0.00692749]]

So the shape and values are different. I checked and executing the training script multiple times gives consistent results (so it's not a problem with randomness).

To Reproduce

1. Start container by running:

docker run --pull=always --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it INTERNAL_IMAGE:pjnl-20240724

2. Later create in the container training script from the file linked to this issue or add in line 584 of the script
   lightning-thunder/thunder/benchmarks/benchmark_litgpt.py
   4 lines of code from the bug description.

3. Assuming the newly created script is called benchmark_litgpt.py call:

• For Eager

torchrun --standalone --max-restarts=0 --nproc-per-node=8  benchmark_litgpt.py  --model_name Llama-3-8B --max_iters 10 --warmup_iters 2 --distributed_mode fsdp --shard_mode zero3 --bucketing_mode block &> file_eager_1.txt

• For Thunder

torchrun --standalone --max-restarts=0 --nproc-per-node=8  benchmark_litgpt.py  --model_name Llama-3-8B --max_iters 10 --warmup_iters 2 --distributed_mode fsdp --shard_mode zero3 --bucketing_mode block --compile thunder &> file_thunder_1.txt

The results will be visible in file_eager_1.txt and file_thunder_1.txt

Expected behavior

Shapes and values of the model weights should be the same

Environment

Output from nvidia-smi

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15              Driver Version: 550.54.15      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA A100-SXM4-80GB          On  |   00000000:07:00.0 Off |                    0 |
| N/A   30C    P0             60W /  400W |       3MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA A100-SXM4-80GB          On  |   00000000:0F:00.0 Off |                    0 |
| N/A   29C    P0             58W /  400W |       3MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   2  NVIDIA A100-SXM4-80GB          On  |   00000000:47:00.0 Off |                    0 |
| N/A   29C    P0             58W /  400W |       3MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   3  NVIDIA A100-SXM4-80GB          On  |   00000000:4E:00.0 Off |                    0 |
| N/A   31C    P0             62W /  400W |    2757MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   4  NVIDIA A100-SXM4-80GB          On  |   00000000:87:00.0 Off |                    0 |
| N/A   33C    P0             59W /  400W |       3MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   5  NVIDIA A100-SXM4-80GB          On  |   00000000:90:00.0 Off |                    0 |
| N/A   33C    P0             62W /  400W |       3MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   6  NVIDIA A100-SXM4-80GB          On  |   00000000:B7:00.0 Off |                    0 |
| N/A   33C    P0             62W /  400W |       3MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   7  NVIDIA A100-SXM4-80GB          On  |   00000000:BD:00.0 Off |                    0 |
| N/A   33C    P0             60W /  400W |       3MiB /  81920MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+

Python packages:

lightning 2.3.3
lightning-thunder 0.2.0.dev0
lightning-utilities 0.11.6
litgpt 0.4.5
nvfuser 0.2.8+gitfa2bedc
nvidia-cudnn-frontend 1.5.2
nvidia-pyprof 3.11.0
pytorch-lightning 2.3.3
torch 2.5.0a0+git16a2a1a
torchmetrics 1.4.0.post0
torchvision 0.19.0a0+d23a6e1

Additional context

(This is py file, but to attach it here I had to change the extension)
benchmark_litgpt.txt

cc @carmocca @crcrpar