Fix OnnxDiscrepancyCheck crash on BF16 models — format_data cannot convert bfloat16 to NumPy

[tadani3]

Describe your changes

Summary

Fix OnnxDiscrepancyCheck (and any other consumer of format_data) crashing on
BF16 models with TypeError: Got unsupported ScalarType BFloat16.

Background / previous logic

olive/common/utils.py::format_data builds the ONNX input feed by calling
torch.Tensor.numpy() directly on each input tensor and then value-casting to the
target ONNX dtype:

return {
    k: np.ascontiguousarray(
        data[k].cpu().numpy() if isinstance(data[k], torch.Tensor) else data[k],
        dtype=name_to_type[k],
    )
    for k in data
    if k in input_names
}

NumPy has no native bfloat16 dtype, so torch.Tensor.numpy() raises for a
torch.bfloat16 tensor before the trailing dtype=name_to_type[k] cast can run.
For a BF16 model, OnnxDiscrepancyCheck._prepare_dataloader generates dummy inputs at
the ONNX graph's input dtypes, so float inputs (e.g. past_key_values.*.key/value)
arrive as torch.bfloat16 and trip this path. The identical flow with FP16 succeeds,
which is the key signal that this is an input-feed-only gap.

BF16 is already handled correctly everywhere else in the pass:

• discrepancy_check.py::_onnx_output_to_torch reinterprets ORT's uint16 output back
  to torch.bfloat16.
• discrepancy_check.py::_onnx_dtype_to_torch maps onnx.TensorProto.BFLOAT16 → torch.bfloat16.
• common/utils.py::tensor_data_to_dtype already includes torch.bfloat16.

So format_data was the sole remaining gap.

Change

format_data now upcasts torch.bfloat16 tensors to float32 before .numpy(), then
lets the existing np.ascontiguousarray(..., dtype=name_to_type[k]) value-cast emit the
correct target dtype (bf16 via ml_dtypes, or fp16/fp32):

def _to_numpy(value):
    if isinstance(value, torch.Tensor):
        if value.dtype == torch.bfloat16:
            value = value.to(torch.float32)
        return value.cpu().numpy()
    return value

return {
    k: np.ascontiguousarray(_to_numpy(data[k]), dtype=name_to_type[k])
    for k in data
    if k in input_names
}

Why fp32 (not a uint16 bit-reinterpret like the output side): the trailing dtype=
is a value-preserving cast, so it needs an array whose values are correct. bf16 → fp32
is lossless (bf16 shares fp32's 8 exponent bits and has only 7 mantissa bits), so casting
back to bf16 returns the identical value. A uint16 .view() reinterpret would feed
integer bit-patterns into the value-cast and corrupt the data.

FP16/FP32/int and numpy-array inputs are untouched (no upcast branch taken), so their
feeds stay byte-identical.

Files changed

• olive/common/utils.py — format_data: upcast bf16 → fp32 before .numpy().
• test/common/test_format_data.py — new unit tests.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

Release note

Fixed a crash (TypeError: Got unsupported ScalarType BFloat16) when running the ONNX
discrepancy check / evaluation input feed on BF16 models.

(Optional) Issue link

[Copilot]

Pull request overview

This PR fixes a crash in Olive’s ONNX input-feed formatting path when inputs are torch.bfloat16, by upcasting BF16 tensors to FP32 before converting to NumPy, then letting the existing NumPy dtype cast produce the desired ONNX input dtype.

Changes:

• Updated olive.common.utils.format_data to upcast torch.bfloat16 tensors to torch.float32 prior to calling .cpu().numpy().
• Added unit tests to cover BF16 regression behavior and ensure non-BF16 inputs remain byte-identical.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
olive/common/utils.py	Adds a BF16-safe tensor→NumPy conversion path to prevent .numpy() from raising on BF16 tensors.
test/common/test_format_data.py	Adds coverage for BF16 inputs and regression tests for byte-identical behavior on other dtypes.

[xiaoyu-work]

can you rename this file to test_utils.py?