Failure tu convert MultiHeadAttention
- 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
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