Debug MPS

[Aedelon]

EDIT: THIS IS A BAD WORK AROUND. SEE SOLUTION AND DISCUSSION HERE

Debug MPS Compatibility in Two Locations

This PR includes two fixes to ensure compatibility with the MPS backend (Apple Silicon) in PyTorch:

_1. Fix in backbone/_torch.py (Lines 137–138):

if DEFAULT_DEVICE == torch.device("mps"):
    q, k, v = q.cpu(), k.cpu(), v.cpu()

The F.scaled_dot_product_attention function is still not supported on MPS. This patch forces tensors to CPU in that case, avoiding runtime errors._

Switching from MPS to CPU and back to MPS significantly reduces training speed.
This is just a temporary workaround to uncover other issues. Please do not consider this PR for merging.

2. Fix in model.py (Lines 238–243):

if DEFAULT_DEVICE == torch.device("mps"):
    decode_one_token = torch.compile(
        decode_one_token, dynamic=True, backend="aot_eager", disable=cg or disable_torch_compile
    )
else:
    decode_one_token = torch.compile(decode_one_token, dynamic=True, disable=cg or disable_torch_compile)

The default backend (inductor) used by torch.compile is not supported on MPS. This change conditionally switches the backend to aot_eager when the selected device is MPS.

These adjustments allow the model to run properly on Apple devices using the MPS backend.

[Aedelon]

==> #190