Fix CANINE fp16 dtype mismatch in attention mask causing ONNX export failure

[lcheng321]

Fixes #47050

Problem

_create_3d_attention_mask_from_input_mask() always builds the mask in float32 through hardcoded .float() and dtype=torch.float32, ignoring the model dtype. In CanineSelfAttention.forward() the ndim==3 branch adds this float32 mask to attention_scores. When the model runs in float16 the addition promotes attention_scores to float32. After softmax attention_probs stays float32 while value_layer is still float16, so the final matmul crashes. Only the shallow character encoders hit this. The deep encoder uses create_bidirectional_mask() which is already dtype aware.

Fix

One line in CanineSelfAttention.forward():

- attention_scores = attention_scores + attention_mask
+ attention_scores = attention_scores + attention_mask.to(attention_scores.dtype)

Reproduction

import torch
from transformers import CanineModel, CanineTokenizer

tok = CanineTokenizer.from_pretrained("google/canine-s")
inputs = tok("hello world", return_tensors="pt")

model = CanineModel.from_pretrained("google/canine-s").half().eval()

torch.onnx.export(
    model,
    (inputs["input_ids"], inputs["attention_mask"]),
    "/tmp/canine_fp16.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["last_hidden_state"],
    opset_version=17,
)

Before

Runtime dtype trace at the crash point:

attention_scores before mask add : torch.float16
attention_mask (3D path)         : torch.float32
attention_scores after mask add  : torch.float32   <- promoted
attention_probs after softmax    : torch.float32
value_layer                      : torch.float16   <- mismatch

Result:

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

ONNX export never completes.

After

Same trace with the fix applied:

attention_scores before mask add : torch.float16
attention_mask after .to()       : torch.float16
attention_scores after mask add  : torch.float16
attention_probs after softmax    : torch.float16
value_layer                      : torch.float16   <- ok

Result:

ONNX export: succeeded

No crash. The dtype stays float16 through the whole attention path.

Numerical check

fp16 output compared against the fp32 baseline on cpu, same input:

max abs diff  : 0.007903
mean abs diff : 0.001273
any NaN       : False

Difference is within normal fp16 tolerance. The output is not corrupted.

Tests

Added test_model_fp16_with_attention_mask to CanineModelTest. It builds the model from config only, no pretrained weights, calls .half(), forces partial masking to hit the 3D mask path, and asserts the output dtype is float16 with no NaN.

Full file result:

66 passed, 137 skipped, 3 xfailed, 0 failed

Existing test_onnx_export_* and test_torch_export_* still pass, so the change introduces no regression in the export paths.

Out of scope

A second independent issue exists, a device mismatch in _repeat_molecules() that becomes visible only after this dtype fix. It is noted in #47050 and left for a follow up PR.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: canine

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 28732666907:2
Result: success | Jobs: 2 | Tests: 8 | Failures: 0 | Duration: 36s

[Rocketknight1]

It's a code agent PR but I'll allow it because I think the bug is real here.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.