[ONNX][dynamo_export] ONNX::Celu Half unsupported but export passed w/ invalid model when opmath disabled

[BowenBao]

In #113780 I'm experimenting turning off opmath type promotion for export since it inserts upcasts that are hard to reason for backends.

This issue can be reproed by below with the above branch removing the respective xfail.

TORCH_LOGS="onnx_diagnostics" pytest test/onnx/test_fx_op_consistency.py -v -n 10 -x -k test_output_match_nn_functional_celu_cpu_float16

Without opmath, aten::celu fp16 will be exported as is to onnx::celu. This should be an early export error since onnx::celu does not support half, however the export succeeds with invalid onnx model. Loading it in ORT later emits error

onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : This is an invalid model. Type Error: Type 'tensor(float16)' of input parameter (l_args_0_) of operator (aten_celu) in node (aten_celu_0) is invalid.

Looking at diagnostics, it turns out that exporter detected the absence of half celu in torchlib, but decides to fallback with the float32 impl.

• Should we tighten the export op dispatching fallback rules regarding dtype?
• Should we add upcast to torchlib celu such that half is supported?

@justinchuby @titaiwangms

[titaiwangms]

Looks like fallback needs to change anyway, because it's fallback to itself. @justinchuby In microsoft/onnxscript#563 (comment), you specifically say that we should support under case, is this part of the scenarios?

[justinchuby]

Yes it falls under ONNX point 1.

I am now inclined to have the converter examine opschema and insert casts for this case, because torchlib doesn’t have information on the output tensor type, only that of the inputs. So even though it may be able to add casts to inputs we still need to cast the outputs somewhere. The converter is in the best position to do so because it knows what the output type should be.

[titaiwangms]

Yes it falls under ONNX point 1.

I am now inclined to have the converter examine opschema and insert casts for this case, because torchlib doesn’t have information on the output tensor type, only that of the inputs. So even though it may be able to add casts to inputs we still need to cast the outputs somewhere. The converter is in the best position to do so because it knows what the output type should be.

Based on #113780, it seems that adding upcast in converter side is not preferred, since the optimization can't tell it's from models or type promotion?

[justinchuby]

What I understand is we don’t prefer always adding upcasts, but we will still need to cast when we actually need them (for onnx support). The current promotion logic doesn’t account for the opschema of the onnx function.

[titaiwangms]

Oops, it looks like it's not about fallback, it's only nearest match. If we added opschema into type promotion, should we disable nearest match?

[justinchuby]

Happy to discuss further offline