Add guidance on DDP bucket_cap_mb for larger models & torch.compile

[chelsea0x3b]

The default is 25mb, which for even llama 8b is too small. Might be best to suggest use of max param size.

Unclear how this impacts things - also guidance on when to use torch.compile

[daire-byrne]

I also noticed that torch.compile was missing from the FSDP train_llm.py (but is used in all the others), so I compared with/without for FSDP.

On my test server with 4 x L40S (2 on each socket):

# no torch compile
torchrun --nnodes 1 --nproc-per-node 4 train_llm.py -b 6 -d tatsu-lab/alpaca -m meta-llama/Llama-3.2-3B
| 214/214 [06:12<00:00,  1.74s/it]

# with torch.compile (model + model.loss_function)
torchrun --nnodes 1 --nproc-per-node 4 train_llm.py -b 10 -d tatsu-lab/alpaca -m meta-llama/Llama-3.2-3B
| 128/128  [05:02<00:00,  2.36s/it]

Not only was it faster, but I was surprised that it allowed for s lightly large batch size (without OOM) too.

BTW it looks like bucket_cap_mb=500 in the latest version of the DDP script.

[daire-byrne]

Conversely, the TP scripts (06 & 07) both error out at the end when torch.compile is enabled (for meta-llama/Llama-3.2-3B).