add fast test to check with bfloat16

olive/common/hf/utils.py

@@ -189,6 +189,17 @@
     ) and attn_implementation is not None:
         from_config_kwargs["attn_implementation"] = attn_implementation
     model = from_config(model_config, **from_config_kwargs)
+    # Re-initialise all floating-point parameters with N(0, 0.02) which is close to
+    # typical LLM weight distributions.  The default HuggingFace init (kaiming_uniform
+    # or xavier_uniform) produces weights with a much wider spread, leading to
+    # unrealistically large discrepancy-check errors after quantization.
+    if hasattr(model, "parameters"):
+        import torch
+
+        with torch.no_grad():
+            for param in model.parameters():
+                if param.is_floating_point():
+                    param.normal_(mean=0.0, std=0.02)
     logger.info("Generating test model class %s", type(model))
     return model
 

olive/common/utils.py

@@ -294,14 +294,29 @@ def format_data(data, io_config):
         data = dict(zip(input_names, [data]))
     elif not isinstance(data, dict):
         raise ValueError(f"Invalid input data format: {data}")
-    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
-    }
+
+    formatted = {}
+    for k in data:
+        if k not in input_names:
+            continue
+        v = data[k]
+        if isinstance(v, torch.Tensor):
+            v = v.cpu()
+            if v.dtype == torch.bfloat16:
+                import ml_dtypes
+
+                v = v.view(torch.uint16).numpy().view(ml_dtypes.bfloat16)
+            else:
+                v = v.numpy()
+
+        target_dtype = name_to_type[k]
+        # ONNX BFLOAT16 is commonly surfaced as "uint16" in io_config; preserve ml_dtypes.bfloat16
+        # so callers can detect bf16 and route through the IOBinding path.
+        if str(getattr(v, "dtype", "")) == "bfloat16" and str(target_dtype) == "uint16":
+            formatted[k] = np.ascontiguousarray(v)
+        else:
+            formatted[k] = np.ascontiguousarray(v, dtype=target_dtype)
+    return formatted
 
 
 def resolve_torch_dtype(dtype):

olive/passes/onnx/discrepancy_check.py

@@ -128,6 +128,58 @@ def _onnx_output_to_torch(onnx_output, reference_dtype):
     return onnx_tensor
 
 
+def _has_bfloat16(input_feed: dict) -> bool:
+    """Return True if any value in the input feed uses bfloat16 (ml_dtypes)."""
+    try:
+        import ml_dtypes
+
+        return any(getattr(v, "dtype", None) == ml_dtypes.bfloat16 for v in input_feed.values())
+    except ImportError:
+        return False
+
+
+def _run_onnx_session(session, input_feed: dict) -> list:
+    """Run ONNX inference, using IOBinding when bfloat16 inputs are present.
+
+    ``session.run()`` does not support bfloat16 numpy arrays because numpy has no native bf16
+    dtype.  When bfloat16 inputs are detected we fall back to IOBinding with
+    ``OrtValue.ortvalue_from_numpy_with_onnx_type`` which reinterprets a uint16 view as
+    ONNX BFLOAT16.  Outputs are extracted from the raw ``OrtValue`` buffer because neither
+    ``copy_outputs_to_cpu`` nor ``OrtValue.numpy`` support bfloat16.
+    """
+    if not _has_bfloat16(input_feed):
+        return session.run(None, input_feed)
+
+    import ctypes
+
+    import ml_dtypes
+    from onnxruntime import OrtValue
+
+    io_binding = session.io_binding()
+    for name, arr in input_feed.items():
+        if arr.dtype == ml_dtypes.bfloat16:
+            # ONNX TensorProto.BFLOAT16 == 16
+            ort_value = OrtValue.ortvalue_from_numpy_with_onnx_type(arr.view(np.uint16), 16)
+        else:
+            ort_value = OrtValue.ortvalue_from_numpy(arr)
+        io_binding.bind_ortvalue_input(name, ort_value)
+    for output in session.get_outputs():
+        # Ensure outputs are placed in host memory since we read them via data_ptr().
+        io_binding.bind_output(output.name, "cpu", 0)
+    io_binding.synchronize_inputs()
+    session.run_with_iobinding(io_binding)
+    io_binding.synchronize_outputs()
+
+    results = []
+    for ort_value in io_binding.get_outputs():
+        if ort_value.data_type() == "tensor(bfloat16)":
+            buf = (ctypes.c_uint8 * ort_value.tensor_size_in_bytes()).from_address(ort_value.data_ptr())
+            results.append(np.frombuffer(buf, dtype=np.uint16).view(ml_dtypes.bfloat16).reshape(ort_value.shape()))
+        else:
+            results.append(ort_value.numpy())
+    return results
+
+
 def _longest_common_token_sequence(seq_a: list[int], seq_b: list[int]) -> int:
     """Compute the length of the longest common token sequence starting from the beginning.
 
@@ -698,7 +750,7 @@ def _compute_logits_discrepancy(self, ref_model, session, dataloader, io_config,
                 torch_logits = torch_output.logits.detach()
                 # Run ONNX inference
                 onnx_input_feed = format_data(input_data, io_config)
-                onnx_outputs = session.run(None, onnx_input_feed)
+                onnx_outputs = _run_onnx_session(session, onnx_input_feed)
                 onnx_logits = _onnx_output_to_torch(onnx_outputs[0], torch_logits.dtype)
 
                 # Compute element-wise differences using torch in double precision
@@ -1139,12 +1191,12 @@ def _measure_speedup(
 
         # Warmup ONNX
         for _ in range(warmup_iterations):
-            session.run(None, onnx_input_feed)
+            _run_onnx_session(session, onnx_input_feed)
 
         # Time ONNX
         start = time.perf_counter()
         for _ in range(timing_iterations):
-            session.run(None, onnx_input_feed)
+            _run_onnx_session(session, onnx_input_feed)
         onnx_time = (time.perf_counter() - start) / timing_iterations
 
         speedup = pytorch_time / onnx_time if onnx_time > 0 else float("inf")

test/cli/test_cli_test_model_smoke.py

@@ -373,6 +373,82 @@ def test_save_local_tiny_qwen2_5_vl_supports_image_processor(self):
             assert "pixel_values" in inputs
             assert "image_grid_thw" in inputs
 
+    @staticmethod
+    def _bf16_cuda_supported():
+        import onnxruntime as ort
+        import torch
+
+        return torch.cuda.is_available() and "CUDAExecutionProvider" in ort.get_available_providers()
+
+    def test_bf16_precision(self):
+        """Verify that the optimize/run flow works when targeting bf16 precision.
+
+        Failures should be investigated as bf16 regressions. The test is skipped automatically
+        when the current environment does not provide CUDAExecutionProvider-backed bf16 support.
+        """
+        if not self._bf16_cuda_supported():
+            self.skipTest("bf16 smoke test requires CUDAExecutionProvider and torch.cuda support.")
+
+        if self.workdir is None:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                self._assert_bf16_precision(Path(temp_dir))
+        else:
+            workdir = Path(self.workdir)
+            workdir.mkdir(parents=True, exist_ok=True)
+            self._assert_bf16_precision(workdir)
+
+    def _assert_bf16_precision(self, tmp_path: Path):
+        model_id = self.model_ids[0]
+        model_name = model_id.replace("/", "--")
+        model_path = tmp_path / "models" / f"{model_name}-bf16"
+        config_output_dir = tmp_path / f"{model_name}-bf16-cfg"
+        run_output_dir = tmp_path / f"{model_name}-bf16-run"
+
+        _save_local_tiny_model(model_id, model_path)
+        _run_cli_main(
+            [
+                "optimize",
+                "-m",
+                str(model_path),
+                "--device",
+                "gpu",
+                "--provider",
+                "CUDAExecutionProvider",
+                "--precision",
+                "bf16",
+                "--output_path",
+                str(config_output_dir),
+                "--dry_run",
+            ]
+        )
+
+        config_path = config_output_dir / "config.json"
+        assert config_path.exists()
+        # run --config dump/config.json --test --test_metrics mae,speedup --output_path dump/run
+        _run_cli_main(
+            [
+                "run",
+                "--config",
+                str(config_path),
+                "--test",
+                "--test_metrics",
+                "mae,speedup",
+                "--output_path",
+                str(run_output_dir),
+            ]
+        )
+
+        assert (run_output_dir / TEST_OUTPUT_MARKER_FILE).exists(), (
+            f"Run output marker not found in {run_output_dir}; the bf16 run may have failed."
+        )
+
+        results_path = run_output_dir / "discrepancy_check_results.json"
+        assert results_path.exists(), f"discrepancy_check_results.json not found in {run_output_dir}"
+        results = json.loads(results_path.read_text())
+        assert "speedup" in results, f"'speedup' key missing from discrepancy results: {results}"
+        assert isinstance(results["speedup"], (int, float)), f"'speedup' is not a number: {results['speedup']!r}"
+        assert results["speedup"] > 0, f"'speedup' must be positive, got {results['speedup']}"
+
     def test_model_discrepancy(self):
         """Verify that OnnxDiscrepancyCheck runs successfully with the configured exporter."""
         if self.workdir is None:

test/passes/onnx/test_discrepancy_check.py

@@ -11,8 +11,10 @@
 
 from olive.passes.onnx.discrepancy_check import (
     _expand_genai_output_names,
+    _has_bfloat16,
     _longest_common_token_sequence,
     _reconcile_genai_speech_output_names,
+    _run_onnx_session,
 )
 
 
@@ -1804,3 +1806,102 @@ def test_save_results_adds_export_info_for_composite_model(self, tmp_path):
             "decoder": {"producer_name": "decoder-exporter", "producer_version": "3.1"},
         }
         assert model.model_attributes["discrepancy_check_results"]["export_info"] == results["export_info"]
+
+
+class TestRunOnnxSessionBfloat16:
+    """Tests for _run_onnx_session and _has_bfloat16 with bfloat16 data on CUDA."""
+
+    @pytest.fixture
+    def bfloat16_identity_onnx(self, tmp_path):
+        """Create a minimal ONNX model with bfloat16 input/output (identity)."""
+        import onnx
+        from onnx import TensorProto, helper
+
+        x_info = helper.make_tensor_value_info("X", TensorProto.BFLOAT16, [None, 4])
+        y_info = helper.make_tensor_value_info("Y", TensorProto.BFLOAT16, [None, 4])
+        node = helper.make_node("Identity", inputs=["X"], outputs=["Y"])
+        graph = helper.make_graph([node], "bf16_identity", [x_info], [y_info])
+        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 19)])
+        model_path = str(tmp_path / "bf16_identity.onnx")
+        onnx.save(model, model_path)
+        return model_path
+
+    def test_has_bfloat16_detects_bf16(self):
+        import ml_dtypes
+        import numpy as np
+
+        feed = {"x": np.ones((2, 3), dtype=ml_dtypes.bfloat16)}
+        assert _has_bfloat16(feed) is True
+
+    def test_has_bfloat16_false_for_float32(self):
+        import numpy as np
+
+        feed = {"x": np.ones((2, 3), dtype=np.float32)}
+        assert _has_bfloat16(feed) is False
+
+    @pytest.mark.skipif(
+        not __import__("torch").cuda.is_available(),
+        reason="CUDA not available",
+    )
+    def test_run_onnx_session_bfloat16_cuda(self, bfloat16_identity_onnx):
+        """_run_onnx_session should handle bfloat16 I/O via IOBinding on CUDA."""
+        import ml_dtypes
+        import numpy as np
+        import onnxruntime as ort
+
+        session = ort.InferenceSession(
+            bfloat16_identity_onnx,
+            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+        )
+        input_data = np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], dtype=ml_dtypes.bfloat16)
+        input_feed = {"X": input_data}
+
+        outputs = _run_onnx_session(session, input_feed)
+
+        assert len(outputs) == 1
+        assert outputs[0].dtype == ml_dtypes.bfloat16
+        assert outputs[0].shape == (2, 4)
+        np.testing.assert_array_equal(outputs[0].view(np.uint16), input_data.view(np.uint16))
+
+    def test_run_onnx_session_bfloat16_cpu(self, bfloat16_identity_onnx):
+        """_run_onnx_session should handle bfloat16 I/O via IOBinding on CPU."""
+        import ml_dtypes
+        import numpy as np
+        import onnxruntime as ort
+
+        session = ort.InferenceSession(
+            bfloat16_identity_onnx,
+            providers=["CPUExecutionProvider"],
+        )
+        input_data = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=ml_dtypes.bfloat16)
+        input_feed = {"X": input_data}
+
+        outputs = _run_onnx_session(session, input_feed)
+
+        assert len(outputs) == 1
+        assert outputs[0].dtype == ml_dtypes.bfloat16
+        assert outputs[0].shape == (1, 4)
+        np.testing.assert_array_equal(outputs[0].view(np.uint16), input_data.view(np.uint16))
+
+    def test_run_onnx_session_float32_uses_standard_run(self, tmp_path):
+        """_run_onnx_session should fall back to session.run() for float32 inputs."""
+        import numpy as np
+        import onnx
+        import onnxruntime as ort
+        from onnx import TensorProto, helper
+
+        x_info = helper.make_tensor_value_info("X", TensorProto.FLOAT, [None, 2])
+        y_info = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None, 2])
+        node = helper.make_node("Identity", inputs=["X"], outputs=["Y"])
+        graph = helper.make_graph([node], "fp32_identity", [x_info], [y_info])
+        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 19)])
+        model_path = str(tmp_path / "fp32_identity.onnx")
+        onnx.save(model, model_path)
+
+        session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+        input_data = np.array([[1.0, 2.0]], dtype=np.float32)
+
+        outputs = _run_onnx_session(session, {"X": input_data})
+
+        assert len(outputs) == 1
+        np.testing.assert_array_almost_equal(outputs[0], input_data)