Running Llama 2 on Apple Silicon GPUs - missing MPS types and operators

[Samyak2]

🚀 The feature, motivation and pitch

I have attempted to run Llama 2 on M-series (M1/M2) Mac GPUs here: https://github.com/Samyak2/llama-mps

Current status

The models loads correctly but inference fails because:

• The ComplexFloat dtype is not supported in MPS yet (Closest existing issue I found: FFT operators are not supported on MPS device #78044)
• The aten::view_as_complex operator is not supported in MPS yet (General MPS op coverage tracking issue #77764)
• The aten::polar.out operator is not supported in MPS yet. This can be worked around by setting PYTORCH_ENABLE_MPS_FALLBACK=1 which runs the operator on CPU instead. For full performance, this operator would need to be supported too.

There may be more operators and types that may need to be supported. I have not dug further on this since it crashes due to ComplexFloat not being supported.

Alternatives

There have been forks of Llama to make it work on CPU instead. Examples: https://github.com/b0kch01/llama-cpu
These will leave a lot of performance on the table though.

Additional context

Failure logs for context (from https://github.com/Samyak2/llama-mps):

<redacted>/llama/llama/model.py:55: UserWarning: The operator 'aten::polar.out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
  freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
Loaded in 11.68 seconds
<redacted>/llama/llama/model.py:72: UserWarning: 0The operator aten::view_as_complex appears to be a view operator, but it has no implementation for the backend "mps:0". View operators don't support falling back to run on the CPU, since the tensor's storage cannot be shared across devices. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/CPUFallback.cpp:181.)
  xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
<redacted>/llama/llama/model.py:73: UserWarning: 0The operator aten::view_as_complex appears to be a view operator, but it has no implementation for the backend "mps:0". View operators don't support falling back to run on the CPU, since the tensor's storage cannot be shared across devices. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/CPUFallback.cpp:181.)
  xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
libc++abi: terminating due to uncaught exception of type c10::TypeError: Trying to convert ComplexFloat to the MPS backend but it does not have support for that dtype.
Exception raised from getMPSScalarType at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/OperationUtils.mm:91 (most recent call first):
frame #0: at::native::mps::getMPSScalarType(c10::ScalarType) + 180 (0x116dc5954 in libtorch_cpu.dylib)
frame #1: invocation function for block in at::native::mps::binaryOpTensor(at::Tensor const&, at::Tensor const&, c10::Scalar const&, at::Tensor const&, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>>, MPSGraphTensor* (at::native::mps::BinaryOpCachedGraph*, MPSGraphTensor*, MPSGraphTensor*) block_pointer) + 108 (0x116de0814 in libtorch_cpu.dylib)
frame #2: invocation function for block in at::native::mps::MPSGraphCache::CreateCachedGraph(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&, at::native::mps::MPSCachedGraph* () block_pointer) + 216 (0x116ddb8d4 in libtorch_cpu.dylib)
frame #3: _dispatch_client_callout + 20 (0x185114400 in libdispatch.dylib)
frame #4: _dispatch_lane_barrier_sync_invoke_and_complete + 56 (0x18512397c in libdispatch.dylib)
frame #5: at::native::mps::MPSGraphCache::CreateCachedGraph(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&, at::native::mps::MPSCachedGraph* () block_pointer) + 160 (0x116dc99e0 in libtorch_cpu.dylib)
frame #6: at::native::mps::binaryOpTensor(at::Tensor const&, at::Tensor const&, c10::Scalar const&, at::Tensor const&, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>>, MPSGraphTensor* (at::native::mps::BinaryOpCachedGraph*, MPSGraphTensor*, MPSGraphTensor*) block_pointer) + 2352 (0x116ddf898 in libtorch_cpu.dylib)
frame #7: at::native::structured_mul_out_mps::impl(at::Tensor const&, at::Tensor const&, at::Tensor const&) + 128 (0x116de33f0 in libtorch_cpu.dylib)
frame #8: at::(anonymous namespace)::wrapper_MPS_mul_Tensor(at::Tensor const&, at::Tensor const&) + 140 (0x11457fea8 in libtorch_cpu.dylib)
frame #9: at::_ops::mul_Tensor::call(at::Tensor const&, at::Tensor const&) + 284 (0x1133bd898 in libtorch_cpu.dylib)
frame #10: torch::autograd::THPVariable_mul(_object*, _object*, _object*) + 396 (0x10726c2dc in libtorch_python.dylib)
frame #11: _object* torch::autograd::TypeError_to_NotImplemented_<&torch::autograd::THPVariable_mul(_object*, _object*, _object*)>(_object*, _object*, _object*) + 12 (0x1071c8330 in libtorch_python.dylib)
<omitting python frames>

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev

[albanD]

Hi,

Sounds like the exact same thing as llama1: you can just update the apply_rotary_emb function not to use complex. You will just have to write the equivalent of the complex multiplication by hand.

[Samyak2]

Hi,

Sounds like the exact same thing as llama1: you can just update the apply_rotary_emb function not to use complex. You will just have to write the equivalent of the complex multiplication by hand.

Thanks for the tip @albanD! I tried implementing this but something is clearly wrong because all the outputs are ?? ?? ??.

Here's my implementation:

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim - 1
    x_shape = x.shape[:-1]
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x_shape[1], x_shape[-1])
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x_shape)]
    return freqs_cis.view(*shape)


def _emulate_complex_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """Emulates complex number multiplication a tensor with the last dim size being 2.

    If you already have a complex tensor, use `torch.view_as_real` to get a tensor
    compatible with this function.

    This is useful for backends that do not support complex numbers yet (MPS for example).

    NOTE: currently hardcoded for tensors with 4 (+1) dimensions!
    """
    acad = a[:, :, :, :, 0:1].repeat(1, 1, 1, 1, 2) * b
    bdbc = (
        a[:, :, :, :, 1:2].repeat(1, 1, 1, 1, 2)
        * b[:, :, :, :, [1, 0]]
        * torch.tensor([-1, 1]).to(device)
    )

    return acad + bdbc


def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)
    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)
    freqs_cis = torch.view_as_real(reshape_for_broadcast(freqs_cis, xq_)).to(device)
    xq_out = _emulate_complex_mul(xq_, freqs_cis).flatten(3)
    xk_out = _emulate_complex_mul(xk_, freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

Here's the original code for reference.

[Samyak2]

After some more investigation, my re-written implementation of apply_rotary_emb seems correct. The output tensors are exactly the same as the original version. I'm not sure why the outputs are the way they are. The branch with these changes is here if anyone wants to try it out: Samyak2/llama-mps#1

[msaroufim]

I got something to work here btw in case anyone is interested karpathy/llama2.c#14

EDIT: @albanD @kulinseth I did try doing this rewrite but then M1 was about 10x slower than just doing complex on CPU so might need to look into this a bit more - lemme know if you see any obvious issues

[albanD]

Do you have sample inputs for which you see this slowdown?
Doing manual complex is expected to be slower than builtin when available in hw :/

[msaroufim]

Not sure how helpful it'll be but here's the inputs for the real and complex versions https://gist.github.com/msaroufim/57199b3dcfda2fcd10f8d7baf8a739f2

And just some quick benchmark showed

• Real version on MPS: apply_rotary_emb took 0.0017440319061279297 seconds
• Complex version on CPU: apply_rotary_emb took 5.5789947509765625e-05 seconds

Is complex support not supported by Apple hardware or is it just that someone needs to write more code in PyTorch? If it's the latter, could I presumably do it as a side project (1-2 dedicated days)

EDIT: I saw you already outlined how people can do it here #77764 lemme poke around

[kulinseth]

Not sure how helpful it'll be but here's the inputs for the real and complex versions https://gist.github.com/msaroufim/57199b3dcfda2fcd10f8d7baf8a739f2

And just some quick benchmark showed

• Real version on MPS: apply_rotary_emb took 0.0017440319061279297 seconds
• Complex version on CPU: apply_rotary_emb took 5.5789947509765625e-05 seconds

Is complex support not supported by Apple hardware or is it just that someone needs to write more code in PyTorch? If it's the latter, could I presumably do it as a side project (1-2 dedicated days)

EDIT: I saw you already outlined how people can do it here #77764 lemme poke around

@msaroufim thanks for the updates . We do have complex support in MPSGraph MacOS 14 . We need to enable it in PyTorch stack . If you are willing, it will be great if you take up the effort of adding complex support . If you have questions we are happy to help you with it

[DarrenAbramson]

@kulinseth would it be possible to add some more detail/link to docs about the complex support in MPSGraph MacOS 14?

There is the example pr for adding ops that already exist in the MPS back end -- do you think following this is still helpful?