Repository navigation
Failure tu convert MultiHeadAttention #2311
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:
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:
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
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<?, ? ops/s]ERROR:coremltools:
ERROR - converting '_native_multi_head_attention' op (located at: 'text/transformer/0/attn'):
Converting PyTorch Frontend ==> 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 <module>
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:
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)
Reacted by Daniel Bourke
Reacted by Jian Liao
Activity
sinnuswong commentedon Sep 24, 2025
do you have any progress? I also got this error.
laclouis5 commentedon Sep 26, 2025
No but the workaround works fine.