torch.compile for training

**Is your feature request related to a problem? Please describe.**

We should leverage torch.compile for training and fprop.

**Describe the solution you'd like**

Ideally we should be able to torch compile all calls to hf model (`train`, `get_logprobs`). The compile cost should be paid once at the init.
FYI there is a known issue with torch.compile + FSDP + BF16 and torch.compile was failing (with torch 2.5.1).

**Describe alternatives you've considered**

thunder compiler: https://github.com/Lightning-AI/lightning-thunder/ (we should investigate perf benefits of thunder over native compile)