|
| 1 | +# ------------------------------------------------------------------------- |
| 2 | +# Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | +# Licensed under the MIT License. |
| 4 | +# -------------------------------------------------------------------------- |
| 5 | +import logging |
| 6 | +from typing import Optional |
| 7 | + |
| 8 | +from olive.data.config import DataConfig |
| 9 | +from olive.hardware import AcceleratorSpec |
| 10 | +from olive.model import ONNXModelHandler |
| 11 | +from olive.passes import Pass |
| 12 | +from olive.passes.pass_config import BasePassConfig, PassConfigParam |
| 13 | + |
| 14 | +logger = logging.getLogger(__name__) |
| 15 | + |
| 16 | + |
| 17 | +def _infer_shape(dynamic_shape, known_values=None): |
| 18 | + default_values = { |
| 19 | + "batch_size": 1, |
| 20 | + "past_sequence_length": 2, |
| 21 | + "total_sequence_length": 3, |
| 22 | + "sequence_length": 1, |
| 23 | + } |
| 24 | + if known_values: |
| 25 | + default_values.update(known_values) |
| 26 | + inferred_shape = [] |
| 27 | + for dim in dynamic_shape: |
| 28 | + if isinstance(dim, int): |
| 29 | + inferred_shape.append(dim) |
| 30 | + continue |
| 31 | + if dim not in default_values: |
| 32 | + raise KeyError( |
| 33 | + f"Unsupported symbolic dimension '{dim}' in shape {dynamic_shape}. " |
| 34 | + f"Known symbols are: {sorted(default_values)}. " |
| 35 | + "Update OnnxDiscrepancyCheck to handle this new case." |
| 36 | + ) |
| 37 | + inferred_shape.append(default_values[dim]) |
| 38 | + return tuple(inferred_shape) |
| 39 | + |
| 40 | + |
| 41 | +class OnnxDiscrepancyCheck(Pass): |
| 42 | + """Validates ONNX model outputs against a reference PyTorch model. |
| 43 | +
|
| 44 | + This pass does not transform the model. It runs inference on both the |
| 45 | + ONNX model and a reference PyTorch/HuggingFace model with the same inputs, |
| 46 | + then compares outputs element-wise. It reports: |
| 47 | + - Maximum absolute error (MaxAE) |
| 48 | + - Number of elements where the absolute difference exceeds 0.1 |
| 49 | + - Number of elements where the absolute difference exceeds 0.01 |
| 50 | +
|
| 51 | + The pass fails if any configured threshold is exceeded. |
| 52 | + """ |
| 53 | + |
| 54 | + @classmethod |
| 55 | + def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: |
| 56 | + return { |
| 57 | + "reference_model_path": PassConfigParam( |
| 58 | + type_=str, |
| 59 | + required=True, |
| 60 | + description="Path to the reference PyTorch/HuggingFace model to compare against.", |
| 61 | + ), |
| 62 | + "max_mae": PassConfigParam( |
| 63 | + type_=Optional[float], |
| 64 | + default_value=None, |
| 65 | + description=( |
| 66 | + "Maximum acceptable absolute error. " |
| 67 | + "If the max absolute difference exceeds this value, the pass fails." |
| 68 | + ), |
| 69 | + ), |
| 70 | + "max_elements_above_0_1": PassConfigParam( |
| 71 | + type_=Optional[int], |
| 72 | + default_value=None, |
| 73 | + description=( |
| 74 | + "Maximum acceptable number of elements with absolute difference > 0.1. If exceeded, the pass fails." |
| 75 | + ), |
| 76 | + ), |
| 77 | + "max_elements_above_0_01": PassConfigParam( |
| 78 | + type_=Optional[int], |
| 79 | + default_value=None, |
| 80 | + description=( |
| 81 | + "Maximum acceptable number of elements with absolute difference > 0.01. " |
| 82 | + "If exceeded, the pass fails." |
| 83 | + ), |
| 84 | + ), |
| 85 | + } |
| 86 | + |
| 87 | + def _run_for_config( |
| 88 | + self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str |
| 89 | + ) -> ONNXModelHandler: |
| 90 | + import numpy as np |
| 91 | + import torch |
| 92 | + |
| 93 | + from olive.common.config_utils import validate_config |
| 94 | + from olive.common.utils import format_data |
| 95 | + from olive.data.template import dummy_data_config_template |
| 96 | + from olive.model.config.io_config import is_io_config_static |
| 97 | + |
| 98 | + io_config = model.io_config |
| 99 | + if io_config: |
| 100 | + if is_io_config_static(io_config): |
| 101 | + input_shapes = io_config.get("input_shapes") |
| 102 | + else: |
| 103 | + input_shapes = [] |
| 104 | + known = {} |
| 105 | + for shape in io_config.get("input_shapes"): |
| 106 | + new_shape = _infer_shape(shape, known) |
| 107 | + input_shapes.append(new_shape) |
| 108 | + known.update(dict(zip(shape, new_shape))) |
| 109 | + data_config = dummy_data_config_template( |
| 110 | + input_shapes, io_config.get("input_names"), io_config.get("input_types") |
| 111 | + ) |
| 112 | + data_config = validate_config(data_config, DataConfig) |
| 113 | + data_config.load_dataset_config.params["max_samples"] = 1 |
| 114 | + else: |
| 115 | + raise RuntimeError( |
| 116 | + f"Model IO config is missing for {model.model_path}; cannot generate dummy inputs for discrepancy check." |
| 117 | + ) |
| 118 | + # Create dataloader |
| 119 | + dc = data_config.to_data_container() |
| 120 | + dataloader = dc.create_dataloader() |
| 121 | + |
| 122 | + # Load reference PyTorch model |
| 123 | + from transformers import AutoConfig, AutoModelForCausalLM |
| 124 | + |
| 125 | + ref_cfg = AutoConfig.from_pretrained(config.reference_model_path) |
| 126 | + architectures = getattr(ref_cfg, "architectures", None) or [] |
| 127 | + if not any("ForCausalLM" in arch for arch in architectures): |
| 128 | + raise ValueError( |
| 129 | + "OnnxDiscrepancyCheck currently supports only HuggingFace causal language models (ForCausalLM). " |
| 130 | + f"Got architectures={architectures}" |
| 131 | + ) |
| 132 | + |
| 133 | + ref_model = AutoModelForCausalLM.from_pretrained(config.reference_model_path) |
| 134 | + ref_model.eval() |
| 135 | + |
| 136 | + # Prepare ONNX session |
| 137 | + session = model.prepare_session() |
| 138 | + io_config = model.io_config |
| 139 | + |
| 140 | + # Run inference on both and compare |
| 141 | + all_max_abs_diff = [] |
| 142 | + all_count_above_0_1 = [] |
| 143 | + all_count_above_0_01 = [] |
| 144 | + total_elements = 0 |
| 145 | + |
| 146 | + with torch.no_grad(): |
| 147 | + for batch in dataloader: |
| 148 | + # Extract input data (batch may be (data, label) or just data) |
| 149 | + input_data = batch[0] if isinstance(batch, (tuple, list)) else batch |
| 150 | + |
| 151 | + # Run PyTorch inference |
| 152 | + if isinstance(input_data, dict): |
| 153 | + torch_inputs = {k: v.clone() for k, v in input_data.items()} |
| 154 | + else: |
| 155 | + torch_inputs = input_data |
| 156 | + |
| 157 | + torch_output = ref_model(**torch_inputs) |
| 158 | + torch_logits = torch_output.logits.detach().cpu().numpy() |
| 159 | + # Run ONNX inference |
| 160 | + onnx_input_feed = format_data(input_data, io_config) |
| 161 | + onnx_outputs = session.run(None, onnx_input_feed) |
| 162 | + onnx_logits = onnx_outputs[0] |
| 163 | + |
| 164 | + # Compute element-wise differences |
| 165 | + abs_diff = np.abs(torch_logits.astype(np.float64) - onnx_logits.astype(np.float64)) |
| 166 | + all_max_abs_diff.append(float(np.max(abs_diff))) |
| 167 | + all_count_above_0_1.append(int(np.sum(abs_diff > 0.1))) |
| 168 | + all_count_above_0_01.append(int(np.sum(abs_diff > 0.01))) |
| 169 | + total_elements += abs_diff.size |
| 170 | + |
| 171 | + max_abs_error = max(all_max_abs_diff) |
| 172 | + count_above_0_1 = sum(all_count_above_0_1) |
| 173 | + count_above_0_01 = sum(all_count_above_0_01) |
| 174 | + |
| 175 | + logger.info( |
| 176 | + "OnnxDiscrepancyCheck: max_abs_error=%.6f, elements_above_0.1=%d/%d, elements_above_0.01=%d/%d", |
| 177 | + max_abs_error, |
| 178 | + count_above_0_1, |
| 179 | + total_elements, |
| 180 | + count_above_0_01, |
| 181 | + total_elements, |
| 182 | + ) |
| 183 | + |
| 184 | + # Check thresholds |
| 185 | + failures = [] |
| 186 | + if config.max_mae is not None and max_abs_error > config.max_mae: |
| 187 | + failures.append(f"Max absolute error {max_abs_error:.6f} exceeds threshold {config.max_mae:.6f}") |
| 188 | + if config.max_elements_above_0_1 is not None and count_above_0_1 > config.max_elements_above_0_1: |
| 189 | + failures.append( |
| 190 | + f"Elements with diff > 0.1: {count_above_0_1} exceeds threshold {config.max_elements_above_0_1}" |
| 191 | + ) |
| 192 | + if config.max_elements_above_0_01 is not None and count_above_0_01 > config.max_elements_above_0_01: |
| 193 | + failures.append( |
| 194 | + f"Elements with diff > 0.01: {count_above_0_01} exceeds threshold {config.max_elements_above_0_01}" |
| 195 | + ) |
| 196 | + |
| 197 | + if failures: |
| 198 | + raise RuntimeError("ONNX model discrepancy check failed:\n" + "\n".join(f" - {f}" for f in failures)) |
| 199 | + |
| 200 | + # Return the model unchanged |
| 201 | + return model |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?v=4&size=40){kind=avatar}
0 commit comments