CanineModel fp16 ONNX export fails with dtype mismatch in attention (expected scalar type Float but found Half)
Author: sgpublic
Pull request: no
Labels: name: bug

### 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
- [x] My own modified scripts

### Tasks

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

### Reproduction

versions:

+ onnx: 1.22.0
+ onnxscript: 0.5.7

use following script for reproduction.

<details>
<summary>script code</summary>

```python
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()
```
</details>

**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. 

Comments:
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:
  ```python
  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
  ```python
  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:
```python
attention_mask = attention_mask.to(attention_scores.dtype)
```
immediately before:
```python
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.

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.

Fixed in https://github.com/huggingface/transformers/pull/47083.