diff --git a/olive/common/hf/utils.py b/olive/common/hf/utils.py index 1e8687c62..9e21e3c47 100644 --- a/olive/common/hf/utils.py +++ b/olive/common/hf/utils.py @@ -189,6 +189,17 @@ class such as ``WhisperForConditionalGeneration`` (private ``_from_config``); bo ) 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 diff --git a/olive/common/utils.py b/olive/common/utils.py index 5f39996f5..4259825dd 100644 --- a/olive/common/utils.py +++ b/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): diff --git a/olive/passes/onnx/discrepancy_check.py b/olive/passes/onnx/discrepancy_check.py index 91fa7e008..e9dc48651 100644 --- a/olive/passes/onnx/discrepancy_check.py +++ b/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") diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index 7b492474e..286f796b1 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/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: diff --git a/test/passes/onnx/test_discrepancy_check.py b/test/passes/onnx/test_discrepancy_check.py index 404bdb3b4..ce32cf788 100644 --- a/test/passes/onnx/test_discrepancy_check.py +++ b/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)