[Vulkan] Cannot serialize a graph containing a non-finite float, so attention-mask models never lower

[msluszniak]

🐛 Describe the bug

The Vulkan backend cannot serialize a graph that contains a non-finite float, so any model whose graph carries one fails to lower. The -inf fill value of a transformer attention mask is the common source, which makes this hit ordinary HuggingFace encoders.

convert_to_flatbuffer dumps the graph with Python's json module and hands the result to flatc. The two disagree on how to spell the non-finite floats: Python emits Infinity / -Infinity / NaN, and flatc accepts none of them.

Minimal repro (no model download):

from executorch.backends.vulkan.serialization.vulkan_graph_schema import (
    Double, VkGraph, VkValue,
)
from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
    convert_to_flatbuffer,
)

convert_to_flatbuffer(
    VkGraph(version="1", chain=[], values=[VkValue(Double(float("-inf")))],
            input_ids=[], output_ids=[], constants=[], shaders=[])
)

error:
  /tmp/tmpcbwdsfuy/schema.json:1: 67: error: cannot parse value starting with: -

subprocess.CalledProcessError: Command '[... flatc --binary ...]' returned non-zero exit status 1

The offending value in the emitted JSON:

{"value": {"double_val": -Infinity}, "value_type": "Double"}

Real-model repro. sentence-transformers/all-MiniLM-L6-v2 partitions fine (14 delegate blobs) and then dies at serialization:

import torch, torch.nn.functional as F
from transformers import AutoModel
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower, EdgeCompileConfig

class Emb(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
    def forward(self, input_ids, attention_mask):
        tok = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]
        mask = attention_mask.unsqueeze(-1).expand(tok.size()).float()
        pooled = torch.sum(tok * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)
        return F.normalize(pooled, p=2, dim=1)

ep = torch.export.export(
    Emb().eval(),
    (torch.zeros(1, 128, dtype=torch.int64), torch.ones(1, 128, dtype=torch.int64)),
)
to_edge_transform_and_lower(
    {"forward": ep},
    partitioner=[VulkanPartitioner()],
    compile_config=EdgeCompileConfig(
        _core_aten_ops_exception_list=[torch.ops.aten.linear.default],
    ),
)

openai/clip-vit-base-patch32 fails the same way (12 non-finite values in its graph).

Why it is hard to diagnose

The error names a byte offset inside a deleted temporary file. Nothing in it points back at the model, the op, or even the concept of a non-finite value, so it reads as a schema or toolchain problem rather than a graph one.

Related: which partition the constant lands in varies with PYTHONHASHSEED, so at the default randomized seed the same script sometimes lowers and sometimes fails. That made this look flaky until the seed was pinned.

Fix

flatc does accept inf and -inf, verified against flatc 24.3.25:

token	flatc
Infinity / -Infinity	reject
NaN / nan / NAN	reject
inf / -inf	accept
infinity / -infinity	reject

So the infinities can be round-tripped exactly. The decompile direction needs the inverse rewrite, since flatc writes bare inf tokens that json.load will not parse. FlatBuffers JSON has no spelling for NaN at all, so that case is better reported directly than emitted as JSON flatc cannot read.

PR: #22307

Versions

ExecuTorch main @ c27baa8 (also reproduces on the v1.4.1 branch). flatc 24.3.25, Python 3.10, torch 2.13.0, macOS 15.5 / arm64.