apple / apple/coremltools

Failure tu convert MultiHeadAttention

Open
#2,311 2 comments 2 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug

I'm trying to convert a PyTorch model to CoreML which contains multi head attention block but the conversion fails. The specific model is SigLIP from Hugging Face (`timm/ViT-SO400M-14-SigLIP`) loaded in PyTorch through OpenCLIP (I'm exporting the image encoder but the issue also occurs with the text encoder).

The issue seems to be in the conversion of the PyTorch module `MultiHeadAttention`. Depending on some conditions, the forward pass either calls `F.multi_head_attention_forward` or `torch._native_multi_head_attention`. The former seems to be convertible to CoreML while the later fails to be converted:

```text
RuntimeError: PyTorch convert function for op '_native_multi_head_attention' not implemented.
```

For now I found a simple workaround which is to disable the attention fast path in PyTorch:

```python
torch.backends.mha.set_fastpath_enabled(False)
```

The conversion of the model succeeds with this setting and the accuracy equivalent to the PyTorch model.

## Stack Trace

```text
Torch version 2.4.0 has not been tested with coremltools. You may run into unexpected errors. Torch 2.3.0 is the most recent version that has been tested.
/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/open_clip/factory.py:129: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
checkpoint = torch.load(checkpoint_path, map_location=map_location)
/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/torch/__init__.py:1777: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
assert condition, message
Converting PyTorch Frontend ==> MIL Ops: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 789/790 [00:00<00:00, 6323.29 ops/s]
Running MIL frontend_pytorch pipeline: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 5/5 [00:00<00:00, 83.09 passes/s]
Running MIL default pipeline: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 86/86 [00:08<00:00, 10.15 passes/s]
Running MIL backend_mlprogram pipeline: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 12/12 [00:00<00:00, 124.15 passes/s]
Converting PyTorch Frontend ==> MIL Ops: 0%| | 0/348 [00:00 MIL Ops: 6%|β–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 21/348 [00:00<00:00, 7452.44 ops/s]
Traceback (most recent call last):
File "/Users/louislac/Downloads/siglip-coreml/.venv/bin/convert", line 8, in
sys.exit(main())
^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
return self.main(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/click/core.py", line 1078, in main
rv = self.invoke(ctx)
^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
return ctx.invoke(self.callback, **ctx.params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/click/core.py", line 783, in invoke
return __callback(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/src/siglip_coreml/cli/convert.py", line 8, in main
convert()
File "/Users/louislac/Downloads/siglip-coreml/src/siglip_coreml/convert.py", line 133, in convert
ct_text_encoder = ct.convert(
^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/_converters_entry.py", line 635, in convert
mlmodel = mil_convert(
^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 188, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
proto, mil_program = mil_convert_to_proto(
^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 288, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/converter.py", line 108, in __call__
return load(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 87, in load
return _perform_torch_convert(converter, debug)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 131, in _perform_torch_convert
raise e
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 123, in _perform_torch_convert
prog = converter.convert()
^^^^^^^^^^^^^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 1293, in convert
convert_nodes(self.context, self.graph, early_exit=not has_states)
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 92, in convert_nodes
raise e # re-raise exception
^^^^^^^
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 87, in convert_nodes
convert_single_node(context, node)
File "/Users/louislac/Downloads/siglip-coreml/.venv/lib/python3.11/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 117, in convert_single_node
raise RuntimeError(
RuntimeError: PyTorch convert function for op '_native_multi_head_attention' not implemented.
```

## To Reproduce

The export of the SigLIP image encoder roughly ressemble this:

```python
import torch
import coremltools as ct
import open_clip

# NOTE: Module to isolate the image encoder.
class SigLIPImageEncoder(torch.nn.Module):
def __init__(self, model: torch.nn.Module) -> None:
super().__init__()

self.model = model

def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.model.encode_image(x)

model, _ = open_clip.create_model_from_pretrained("hf-hub:timm/ViT-SO400M-14-SigLIP")
image_encoder = SigLIPImageEncoder(model)
sample_image = torch.randn(1, 3, 224, 224)

# NOTE: Not using inference mode causes a weird crash.
with torch.inference_mode():
tr_image_encoder = torch.jit.trace(
func=image_encoder.eval(),
example_inputs=[sample_image],
)

ct_image_encoder = ct.convert(
model=tr_image_encoder,
inputs=[ct.TensorType(name="image", shape=sample_image.shape, dtype=np.float32)],
outputs=[ct.TensorType(name="embedding", dtype=np.float32)],
minimum_deployment_target=ct.target.iOS18,
)

ct_image_encoder.save("image-encoder.mlpackage")
```

## System environment (please complete the following information):
- coremltools version: 8.0b2
- OS (e.g. MacOS version or Linux type): macOS Sequoia 15.0
- Any other relevant version information (e.g. PyTorch or TensorFlow version): PyTorch 2.4 (also tested with other versions)

Contributor guide

Open the contributing guide

Research direction

Start in coremltools/converters/mil/frontend/torch/ops.py, where the conversion of _native_multi_head_attention raises the reported error, and reproduce it with the SigLIP export example and the listed PyTorch/coremltools versions. Done means the traced SigLIP encoder converts successfully without disabling PyTorch's MHA fast path; src/siglip_coreml/convert.py shows the calling conversion entry point.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, tooling
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.