To enable chunked loss, I refactor the `train_spec.loss_fn` to
`train_spec.build_loss_fn` as suggested in #996

 Also:
- Unified chunked loss implementation with the `cross_entropy_loss`
name, since default is a special case with chunk=1.
Identical loss and memory:

<img width="1257" alt="Screenshot 2025-03-30 at 13 03 44"
src="https://github.com/user-attachments/assets/11c07849-2340-4b42-ad5c-d96657ee53b9"
/>

- Enable compiling the loss function as it's working fine for me, and
result in reduction of memory and slightly better througput.

<img width="1275" alt="Screenshot 2025-03-30 at 16 47 00"
src="https://github.com/user-attachments/assets/09edaf16-1770-48da-87b2-e4de29e7519c"
/>

While we only have one type of loss, I create a field for `loss.name`
with default to 'cross_entropy' to be fork-friendly anyway.

Co-authored-by: tianyu-l <[email protected]>