CanineModel fp16 ONNX export fails with dtype mismatch in attention (expected scalar type Float but found Half)

[sgpublic]

System Info

• transformers version: 5.12.1
• Platform: Linux-6.18.18-trim-x86_64-with-glibc2.41
• Python version: 3.12.13
• Huggingface_hub version: 1.22.0
• Safetensors version: 0.8.0
• Accelerate version: 1.14.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.9.0+cu129 (CUDA)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?: no
• Using GPU in script?: no
• GPU type: NVIDIA GeForce RTX 5060

Who can help?

@NielsRogge @sgugger

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

versions:

• onnx: 1.22.0
• onnxscript: 0.5.7

use following script for reproduction.

script code

import tempfile
from pathlib import Path

import onnx
import onnxruntime as ort
import torch
from transformers import CanineModel
import transformers


class CanineWrapper(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = CanineModel.from_pretrained("google/canine-s")

    def forward(self, input_ids, attention_mask):
        return self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
        ).last_hidden_state


def main():
    print("torch:", torch.__version__)
    print("transformers:", transformers.__version__)
    print("onnxruntime:", ort.__version__)

    model = CanineWrapper().eval()
    model.half()

    input_ids = torch.randint(
        low=0,
        high=256,
        size=(1, 128),
        dtype=torch.long,
        device=device,
    )
    attention_mask = torch.ones(
        1,
        128,
        dtype=torch.long,
        device=device,
    )

    output_path = Path(tempfile.gettempdir()) / "canine_fp16_repro.onnx"
    print("exporting to:", output_path)

    torch.onnx.export(
        model,
        (input_ids, attention_mask),
        str(output_path),
        input_names=["input_ids", "attention_mask"],
        output_names=["last_hidden_state"],
        opset_version=18,
        do_constant_folding=True,
        dynamo=False,
    )

    print("onnx export succeeded")

    onnx_model = onnx.load(str(output_path))
    onnx.checker.check_model(onnx_model)
    print("onnx checker succeeded")

    session = ort.InferenceSession(
        str(output_path),
        providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
    )
    print("onnxruntime session created")

    session.run(
        ["last_hidden_state"],
        {
            "input_ids": input_ids.cpu().numpy(),
            "attention_mask": attention_mask.cpu().numpy(),
        },
    )
    print("onnxruntime inference succeeded")


if __name__ == "__main__":
    main()

Actual behavior:

The export fails during torch.onnx.export(...):

torch: 2.9.0+cu129
transformers: 5.12.1
onnxruntime: 1.23.2
exporting to: /tmp/canine_fp16_repro.onnx

RuntimeError: expected scalar type Float but found Half

Relevant traceback:

File ".../transformers/models/canine/modeling_canine.py", line 353, in forward
    context_layer = torch.matmul(attention_probs, value_layer)
RuntimeError: expected scalar type Float but found Half

Expected behavior

CanineModel should be exportable to ONNX in fp16 after calling model.half(), or at least the forward pass used by ONNX tracing should keep attention tensors in consistent dtypes.

Specifically, attention_probs and value_layer should have compatible dtypes before torch.matmul(...).

If fp16 export for CANINE is not supported, the failure should ideally be documented or surfaced with a clearer error message.

[123RohitVarshit]

I was able to reproduce this independently on a fresh checkout of v5.12.1 using PyTorch 2.11.0+cu128, ONNX Runtime 1.23.2, CUDA 12.8, on a Tesla T4 - a different PyTorch version and GPU than the original report (PyTorch 2.9, RTX 5060), so this doesn't appear to be hardware- or version-specific.

Root cause of the first failure

I instrumented both CanineSelfAttention.forward() and _create_3d_attention_mask_from_input_mask() with dtype prints. Here's what I observed:

• After model.half(), query_layer, key_layer, and value_layer are all float16.
• _create_3d_attention_mask_from_input_mask() always creates a float32 mask, regardless of model dtype:

  to_mask = torch.reshape(to_mask, (batch_size, 1, to_seq_length)).float()
  broadcast_ones = torch.ones(
      size=(batch_size, from_seq_length, 1),
      dtype=torch.float32,
      device=to_mask.device,
  )
  mask = broadcast_ones * to_mask

• In CanineSelfAttention.forward(), adding this float32 mask to attention_scores (float16) promotes attention_scores to float32.
• Consequently, attention_probs (after softmax) is also float32, while value_layer remains float16.
• The subsequent

  context_layer = torch.matmul(attention_probs, value_layer)

  fails with:

  RuntimeError: expected scalar type Float but found Half
  

  which matches the original report.

This path is specific to the shallow character encoders (initial_char_encoder / final_char_encoder). The deep encoder uses create_bidirectional_mask(), which is already dtype-aware and doesn't hit this issue.

Debugging experiment (not a proposed fix)

I inserted:

attention_mask = attention_mask.to(attention_scores.dtype)

immediately before:

attention_scores = attention_scores + attention_mask

This keeps the attention computation entirely in float16, and export proceeds past the original failure.

I'm unsure whether the correct fix is:

• making _create_3d_attention_mask_from_input_mask() construct the mask in the model's dtype directly, or
• having CanineSelfAttention defensively cast any incoming mask to attention_scores.dtype.

Happy to submit a PR for whichever approach is preferred.

Second issue

After the dtype mismatch is resolved, export progresses further but then fails inside _repeat_molecules() with:

RuntimeError: Expected all tensors to be on the same device, but got index is on cpu, different from other tensors on cuda:0

I haven't fully investigated this yet, so I don't want to claim a root cause prematurely - my current guess is that a shape-derived value gets traced as a CPU tensor during ONNX export and later interacts with a CUDA tensor in torch.repeat_interleave, but that's unconfirmed.

Summary

There appear to be at least two independent FP16 ONNX export issues in CanineModel:

1. a confirmed dtype promotion caused by the 3D attention mask in the shallow encoders, and
2. a separate, not-yet-diagnosed device-placement issue in _repeat_molecules() that becomes visible once the first is fixed.

Let me know if a PR addressing the mask issue is welcome on its own, or if you'd prefer both resolved together.

[sgpublic]

Hi @123RohitVarshit , thanks for trying to reproduce this.

I realized that the issue is not actually specific to CUDA. The underlying problem is the fp16 dtype mismatch during ONNX export, and it can also be reproduced on CPU.

I’ve updated the issue title and description to remove the unnecessary CUDA-specific wording and keep the report focused on the actual problem.

Sorry for the confusion.

[lcheng321]

Fixed in #47083.