diff --git a/docs/reference/index.md b/docs/reference/index.md index 30541b579..e03021c1e 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -68,6 +68,7 @@ stages based on the target device and precision. | `dtype` | `str \| null` | Data type (e.g., `float32`, `int64`). | | `shape` | `list[int \| str] \| null` | Tensor shape (e.g., `[1, 3, 224, 224]`). String entries declare symbolic dynamic axes and use size `1` for dummy inputs. | | `value_range` | `[float, float] \| null` | Min/max for dummy tensor generation. | +| `dummy_value_runs` | `list[[int, int \| float]] \| null` | Run-length encoded semantic dummy values as `[count, value]` pairs. Counts must exactly fill the concrete tensor shape. | --- @@ -137,6 +138,7 @@ Set to `null` (default) to skip evaluation. | `model_id` | `str \| null` | `null` | HuggingFace model ID for config resolution. | | `model_path` | `str \| dict[str, str] \| null` | `null` | Path to .onnx file, or a `{role: path}` dict for composite models. | | `task` | `str \| null` | `null` | Task type. | +| `prompt` | `str \| null` | `null` | Optional text prompt passed with each image for prompt-conditioned image-to-text models. | | `device` | `str` | `"auto"` | Inference device. | | `precision` | `str` | `"auto"` | Precision (`fp32`, `fp16`, `w8a16`, etc.). | | `ep` | `str \| null` | `null` | EP override. | diff --git a/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp16_config_decoder.json b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp16_config_decoder.json new file mode 100644 index 000000000..79abb5301 --- /dev/null +++ b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp16_config_decoder.json @@ -0,0 +1,299 @@ +{ + "precision": "fp16", + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "decoder_input_ids", + "dtype": "int32", + "shape": [ + 1, + 1 + ], + "value_range": [ + 0, + 51289 + ] + }, + { + "name": "encoder_hidden_states", + "dtype": "float32", + "shape": [ + 1, + 585, + 768 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "decoder_attention_mask", + "dtype": "int64", + "shape": [ + 1, + 1024 + ] + }, + { + "name": "cache_position", + "dtype": "int64", + "shape": [ + 1 + ] + }, + { + "name": "past_0_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_0_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_1_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_1_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_2_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_2_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_3_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_3_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_4_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_4_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_5_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_5_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + }, + { + "name": "present_0_key" + }, + { + "name": "present_0_value" + }, + { + "name": "present_1_key" + }, + { + "name": "present_1_value" + }, + { + "name": "present_2_key" + }, + { + "name": "present_2_value" + }, + { + "name": "present_3_key" + }, + { + "name": "present_3_value" + }, + { + "name": "present_4_key" + }, + { + "name": "present_4_value" + }, + { + "name": "present_5_key" + }, + { + "name": "present_5_value" + } + ] + }, + "optim": { + "gelu_fusion": true, + "layer_norm_fusion": true, + "matmul_add_fusion": true + }, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "task": "text2text-generation", + "model_id": "microsoft/Florence-2-base", + "model_type": "florence2", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "text2text-generation", + "model_class": "Florence2DecoderWrapper", + "model_type": "florence2" + } +} \ No newline at end of file diff --git a/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp16_config_encoder.json b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp16_config_encoder.json new file mode 100644 index 000000000..4bc65ce1c --- /dev/null +++ b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp16_config_encoder.json @@ -0,0 +1,109 @@ +{ + "precision": "fp32", + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int64", + "shape": [ + 1, + 585 + ], + "dummy_value_runs": [ + [ + 577, + 50265 + ], + [ + 1, + 0 + ], + [ + 1, + 2264 + ], + [ + 1, + 473 + ], + [ + 1, + 5 + ], + [ + 1, + 2274 + ], + [ + 1, + 6190 + ], + [ + 1, + 116 + ], + [ + 1, + 2 + ] + ] + }, + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 768, + 768 + ], + "dummy_value_runs": [ + [ + 1769472, + 0.0 + ] + ] + }, + { + "name": "attention_mask", + "dtype": "int64", + "shape": [ + 1, + 585 + ], + "dummy_value_runs": [ + [ + 585, + 1 + ] + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ] + }, + "optim": { + "gelu_fusion": true, + "layer_norm_fusion": true, + "matmul_add_fusion": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "image-feature-extraction", + "model_class": "Florence2EncoderWrapper", + "model_type": "florence2" + } +} \ No newline at end of file diff --git a/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp32_config_decoder.json b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp32_config_decoder.json new file mode 100644 index 000000000..0b5270a9e --- /dev/null +++ b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp32_config_decoder.json @@ -0,0 +1,277 @@ +{ + "precision": "fp32", + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "decoder_input_ids", + "dtype": "int32", + "shape": [ + 1, + 1 + ], + "value_range": [ + 0, + 51289 + ] + }, + { + "name": "encoder_hidden_states", + "dtype": "float32", + "shape": [ + 1, + 585, + 768 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "decoder_attention_mask", + "dtype": "int64", + "shape": [ + 1, + 1024 + ] + }, + { + "name": "cache_position", + "dtype": "int64", + "shape": [ + 1 + ] + }, + { + "name": "past_0_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_0_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_1_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_1_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_2_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_2_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_3_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_3_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_4_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_4_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_5_key", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "past_5_value", + "dtype": "float32", + "shape": [ + 1, + 12, + 1024, + 64 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + }, + { + "name": "present_0_key" + }, + { + "name": "present_0_value" + }, + { + "name": "present_1_key" + }, + { + "name": "present_1_value" + }, + { + "name": "present_2_key" + }, + { + "name": "present_2_value" + }, + { + "name": "present_3_key" + }, + { + "name": "present_3_value" + }, + { + "name": "present_4_key" + }, + { + "name": "present_4_value" + }, + { + "name": "present_5_key" + }, + { + "name": "present_5_value" + } + ] + }, + "optim": { + "gelu_fusion": true, + "layer_norm_fusion": true, + "matmul_add_fusion": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "text2text-generation", + "model_class": "Florence2DecoderWrapper", + "model_type": "florence2" + } +} \ No newline at end of file diff --git a/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp32_config_encoder.json b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp32_config_encoder.json new file mode 100644 index 000000000..4bc65ce1c --- /dev/null +++ b/examples/recipes/microsoft_Florence-2-base/cpu/cpu/image-to-text_fp32_config_encoder.json @@ -0,0 +1,109 @@ +{ + "precision": "fp32", + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int64", + "shape": [ + 1, + 585 + ], + "dummy_value_runs": [ + [ + 577, + 50265 + ], + [ + 1, + 0 + ], + [ + 1, + 2264 + ], + [ + 1, + 473 + ], + [ + 1, + 5 + ], + [ + 1, + 2274 + ], + [ + 1, + 6190 + ], + [ + 1, + 116 + ], + [ + 1, + 2 + ] + ] + }, + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 768, + 768 + ], + "dummy_value_runs": [ + [ + 1769472, + 0.0 + ] + ] + }, + { + "name": "attention_mask", + "dtype": "int64", + "shape": [ + 1, + 585 + ], + "dummy_value_runs": [ + [ + 585, + 1 + ] + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ] + }, + "optim": { + "gelu_fusion": true, + "layer_norm_fusion": true, + "matmul_add_fusion": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "image-feature-extraction", + "model_class": "Florence2EncoderWrapper", + "model_type": "florence2" + } +} \ No newline at end of file diff --git a/scripts/e2e_eval/testsets/models_all.json b/scripts/e2e_eval/testsets/models_all.json index df50dbf7b..50eef8718 100644 --- a/scripts/e2e_eval/testsets/models_all.json +++ b/scripts/e2e_eval/testsets/models_all.json @@ -3710,6 +3710,17 @@ "optimum_supported": true, "order": 8 }, + { + "hf_id": "microsoft/Florence-2-base", + "task": "image-to-text", + "model_type": "florence2", + "group": "microsoft", + "priority": "P1", + "downloads": 2651904, + "last_update_time": "2025-08-04T17:36:57+00:00", + "optimum_supported": false, + "order": 1 + }, { "hf_id": "microsoft/Phi-4-multimodal-instruct", "task": "visual-question-answering", diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index 733063492..dd5a6ddf8 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -958,6 +958,8 @@ def build( runtime_ep_value = cast("EPNameOrAlias", resolved_target.ep) logger.info("Auto-resolved device=%s, EP=%s", runtime_device, runtime_ep_value) + autogenerated_components: dict[str, str] | None = None + # Load or auto-generate config if config_file is not None: config_or_configs = _load_config( @@ -1021,8 +1023,28 @@ def build( ep=runtime_ep_value, ) else: + from ..loader.resolution import resolve_composite_components + + try: + autogenerated_components = resolve_composite_components( + model, + task=None, + trust_remote_code=trust_remote_code, + ) + except click.ClickException: + raise + except Exception as e: + raise click.ClickException( + f"Composite model detection failed unexpectedly: {e}" + ) from e + outer_task = ( + next(iter(autogenerated_components.values())) + if autogenerated_components + else None + ) config_or_configs = generate_build_config( model, + task=outer_task, trust_remote_code=trust_remote_code, device=runtime_device, precision=precision, @@ -1082,7 +1104,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: cfg.quant.rtn_accuracy_level = resolved_quant.rtn_accuracy_level # Store the original precision string for stage display if precision: - cfg.precision = precision.lower() # type: ignore[attr-defined] + cfg.precision = precision.lower() if cfg.compile is not None and cfg.compile.ep_config is not None: provider = cfg.compile.ep_config.provider patched = WinMLCompileConfig.for_provider(provider, device=runtime_device) @@ -1275,8 +1297,8 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: # Detect composite pipeline (registry-driven, same pattern as # export command). A composite fans out into one build per # sub-component; a plain model builds to the single output dir. - components = None - if model and not model_is_onnx: + components = autogenerated_components + if model and not model_is_onnx and config_file is not None: try: from ..loader.resolution import resolve_composite_components @@ -1336,6 +1358,16 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: ) completed: list[str] = [] + from ..loader.resolution import resolve_composite_precision_overrides + + composite_model_type = getattr(preloaded_hf_config, "model_type", None) + if composite_model_type is None and config.loader is not None: + composite_model_type = config.loader.model_type + precision_overrides = ( + resolve_composite_precision_overrides(composite_model_type, components) + if composite_model_type + else {} + ) try: for name, component_task in components.items(): console.print( @@ -1344,12 +1376,13 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: from ..config import generate_build_config as gen_cfg + component_precision = precision_overrides.get(name, precision) component_config = gen_cfg( model, task=component_task, trust_remote_code=trust_remote_code, device=runtime_device, - precision=precision, + precision=component_precision, ep=runtime_ep_value, export_policy_target=(request_device, request_ep_value), shape_config=shape_overrides, @@ -1361,7 +1394,9 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: # across sub-builds, and preserve the component-specific # quant metadata (task, model_id, model_type) that # generate_build_config populated for the sub-model. - if config.quant is None: + if name in precision_overrides: + pass + elif config.quant is None: component_config.quant = None else: # Overlay the outer quant settings, but always keep diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index 60317e09f..f0acd3cdc 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -653,6 +653,14 @@ def _generate_pipeline_configs( ) -> None: """Generate and save one config file per pipeline sub-component.""" from ..config import generate_hf_build_config + from ..loader.resolution import resolve_composite_precision_overrides + + precision_overrides = resolve_composite_precision_overrides( + model_type, + components, + hf_model=hf_model, + trust_remote_code=trust_remote_code, + ) for component_name, component_task in components.items(): console.print( @@ -660,6 +668,7 @@ def _generate_pipeline_configs( f"(task={component_task})...[/dim]" ) + component_precision = precision_overrides.get(component_name, precision) cfg = generate_hf_build_config( model_id=hf_model, task=component_task, @@ -669,7 +678,7 @@ def _generate_pipeline_configs( shape_config=shape_config, library_name=library_name, device=device, - precision=precision, + precision=component_precision, trust_remote_code=trust_remote_code, ep=ep, policy_overrides_config=policy_overrides_config, diff --git a/src/winml/modelkit/config/build.py b/src/winml/modelkit/config/build.py index e663e367f..05b4e4f20 100644 --- a/src/winml/modelkit/config/build.py +++ b/src/winml/modelkit/config/build.py @@ -141,6 +141,7 @@ class WinMLBuildConfig: compile: WinMLCompileConfig | None = field(default_factory=WinMLCompileConfig) eval: WinMLEvaluationConfig | None = None auto: bool = True + precision: str = "auto" # Skip ORT optimization. Pre-quantized inputs also clear ``quant``. skip_optimize: bool = False @@ -176,6 +177,7 @@ def from_dict(cls, config_dict: dict) -> WinMLBuildConfig: ), eval=eval_cfg, auto=config_dict.get("auto", True), + precision=config_dict.get("precision", "auto"), skip_optimize=config_dict.get("skip_optimize", False), ) @@ -186,6 +188,8 @@ def to_dict(self) -> dict: result["auto"] = False if self.skip_optimize: result["skip_optimize"] = True + if self.precision != "auto": + result["precision"] = self.precision result.update( { "export": self.export.to_dict() if self.export is not None else None, @@ -436,7 +440,7 @@ def _apply_target_policy( config.quant = None # Store resolved precision for multi-pass expansion. - config.precision = policy.precision # type: ignore[attr-defined] + config.precision = policy.precision if policy.compile_provider is not None: config.compile = WinMLCompileConfig.for_provider( diff --git a/src/winml/modelkit/eval/config.py b/src/winml/modelkit/eval/config.py index 1717ccdb5..507607999 100644 --- a/src/winml/modelkit/eval/config.py +++ b/src/winml/modelkit/eval/config.py @@ -151,6 +151,7 @@ class WinMLEvaluationConfig: input_data: str | None = None reference_path: str | None = field(default=None, metadata={"cli_name": "reference"}) task: str | None = None + prompt: str | None = None device: str = "auto" precision: str = "auto" ep: EPNameOrAlias | None = None @@ -187,6 +188,8 @@ def to_dict(self) -> dict: result["reference_path"] = self.reference_path if self.task is not None: result["task"] = self.task + if self.prompt is not None: + result["prompt"] = self.prompt result["device"] = self.device if self.precision != "auto": result["precision"] = self.precision @@ -239,6 +242,7 @@ def from_dict(cls, data: dict) -> WinMLEvaluationConfig: input_data=data.get("input_data"), reference_path=data.get("reference_path"), task=data.get("task"), + prompt=data.get("prompt"), device=data.get("device", "auto"), precision=data.get("precision", "auto"), ep=data.get("ep"), diff --git a/src/winml/modelkit/eval/image_to_text_evaluator.py b/src/winml/modelkit/eval/image_to_text_evaluator.py index db63404ed..89f191364 100644 --- a/src/winml/modelkit/eval/image_to_text_evaluator.py +++ b/src/winml/modelkit/eval/image_to_text_evaluator.py @@ -50,6 +50,7 @@ def __init__( cm = config.dataset.columns_mapping self._image_col = cm.get("input_column", get_default("image-to-text", "input_column")) self._label_col = cm.get("label_column", get_default("image-to-text", "label_column")) + self._prompt = config.prompt or "" super().__init__(config, model) def align_labels(self, dataset: Dataset, ds_config: DatasetConfig) -> Dataset: @@ -73,7 +74,7 @@ def compute(self) -> dict[str, Any]: continue try: - out = self.pipe(image, text="") + out = self.pipe(image, text=self._prompt) except Exception as e: logger.warning("Pipeline call failed (skipping): %s", e) skipped += 1 diff --git a/src/winml/modelkit/export/config.py b/src/winml/modelkit/export/config.py index 7649760ae..193e0f931 100644 --- a/src/winml/modelkit/export/config.py +++ b/src/winml/modelkit/export/config.py @@ -493,6 +493,7 @@ def _resolve_export_config_from_specs( ) value_ranges = io_specs.get("value_ranges", {}) + dummy_value_runs = io_specs.get("dummy_value_runs", {}) input_tensors = [ InputTensorSpec( @@ -500,6 +501,7 @@ def _resolve_export_config_from_specs( shape=shape, dtype=dtype, value_range=value_ranges.get(name), + dummy_value_runs=dummy_value_runs.get(name), ) for name, shape, dtype in zip(input_names, input_shapes, input_dtypes, strict=False) ] diff --git a/src/winml/modelkit/export/io.py b/src/winml/modelkit/export/io.py index 3bedca226..95976dbc8 100644 --- a/src/winml/modelkit/export/io.py +++ b/src/winml/modelkit/export/io.py @@ -41,6 +41,7 @@ ) from ..loader import to_optimum_task +from ..onnx import InputTensorSpec from .value_range import intercept_value_ranges @@ -486,13 +487,20 @@ def resolve_io_specs( input_shapes = [tuple(t.shape) for t in dummy_inputs.values()] input_dtypes = [str(t.dtype).replace("torch.", "") for t in dummy_inputs.values()] + dummy_value_runs = {} + if getattr(onnx_config, "PRESERVE_DUMMY_VALUE_RUNS", False): + dummy_value_runs = { + name: value_runs + for name, tensor in dummy_inputs.items() + if (value_runs := InputTensorSpec.compact_dummy_value_runs(tensor)) is not None + } # Build value_range dict: {name: (min, max)} from intercepted data value_range_tuples = { name: (info["min"], info["max"]) for name, info in value_ranges.items() } - return { + specs = { "inputs": onnx_config.inputs, "outputs": onnx_config.outputs, "input_names": list(onnx_config.inputs.keys()), @@ -502,3 +510,6 @@ def resolve_io_specs( "input_dtypes": input_dtypes, "value_ranges": value_range_tuples, } + if dummy_value_runs: + specs["dummy_value_runs"] = dummy_value_runs + return specs diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index e14de7af7..6a383c3be 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -547,11 +547,15 @@ def create_pipeline( if compatibility_factory is not None: pipe = compatibility_factory(model, model_id) else: + component_kwargs = _pipeline_component_kwargs(hf_task, model_id) + processor_loader = getattr(type(model), "load_pipeline_processor", None) + if processor_loader is not None and model_id is not None: + component_kwargs["processor"] = processor_loader(model, model_id) kwargs: dict[str, Any] = { # "device" is for HF pipeline tensor placement, not ORT EP. # WinMLSession handles device delegation internally. "device": "cpu", - **_pipeline_component_kwargs(hf_task, model_id), + **component_kwargs, } # transformers.pipeline has 60+ Literal overloads — runtime task strings can't diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index a1a372353..1de77fd1a 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -312,6 +312,39 @@ def resolve_composite(model_type: str, task: str) -> CompositeComponents | None: return dict(cls._SUB_MODEL_CONFIG) if cls is not None else None +def resolve_composite_precision_overrides( + model_type: str | None, + components: CompositeComponents, + *, + hf_model: str | None = None, + trust_remote_code: bool = False, +) -> dict[str, str]: + """Return component precision overrides declared by the matching composite.""" + if model_type is None and hf_model is not None: + from transformers import AutoConfig + + from ._autoconfig import load_hf_config + + model_type = load_hf_config( + AutoConfig, + hf_model, + trust_remote_code=trust_remote_code, + ).model_type + if model_type is None: + return {} + matching = { + tuple(sorted(cls._SUB_MODEL_PRECISION_OVERRIDES.items())) + for (registered_type, _), cls in _composite_registry().items() + if registered_type == model_type and components == cls._SUB_MODEL_CONFIG + } + if len(matching) > 1: + raise ValueError( + f"Composite registrations for {model_type!r} declare inconsistent " + "component precision overrides." + ) + return dict(next(iter(matching), ())) + + def resolve_composite_components( hf_model: str | None, *, diff --git a/src/winml/modelkit/loader/task.py b/src/winml/modelkit/loader/task.py index 0ae0ca6be..bf23d433d 100644 --- a/src/winml/modelkit/loader/task.py +++ b/src/winml/modelkit/loader/task.py @@ -128,6 +128,7 @@ # the live registry by `tests/unit/loader/test_composite_tasks.py`. COMPOSITE_TASKS: frozenset[str] = frozenset( { + "image-text-to-text", "image-to-text", "summarization", "table-question-answering", diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index b11e7b416..d12c5ff6d 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -48,6 +48,10 @@ from .depth_anything import DepthAnythingIOConfig as _DepthAnythingIOConfig # triggers registration from .depth_pro import DepthProIOConfig as _DepthProIOConfig # triggers registration from .detr import DETR_CONFIG +from .florence2 import FLORENCE2_CONFIG +from .florence2 import MODEL_CLASS_MAPPING as _FLORENCE2_CLASS_MAPPING +from .florence2 import Florence2DecoderIOConfig as _Florence2DecoderIOConfig +from .florence2 import Florence2EncoderIOConfig as _Florence2EncoderIOConfig from .layoutlm import LayoutLMQAIOConfig as _LayoutLMQAIOConfig # triggers registration from .layoutlmv3 import LAYOUTLMV3_CONFIG from .layoutlmv3 import LayoutLMv3IOConfig as _LayoutLMv3IOConfig # triggers registration @@ -123,6 +127,7 @@ _BART_CLASS_MAPPING, _BLIP_CLASS_MAPPING, _CLIP_CLASS_MAPPING, + _FLORENCE2_CLASS_MAPPING, _MARIAN_CLASS_MAPPING, _MU2_CLASS_MAPPING, _QWEN_CLASS_MAPPING, @@ -151,6 +156,7 @@ "clip-text-model": CLIP_CONFIG, "clip-vision-model": CLIP_CONFIG, "detr": DETR_CONFIG, + "florence2": FLORENCE2_CONFIG, "layoutlmv3": LAYOUTLMV3_CONFIG, "marian": MARIAN_CONFIG, "roberta": ROBERTA_FAMILY_CONFIG, diff --git a/src/winml/modelkit/models/hf/florence2.py b/src/winml/modelkit/models/hf/florence2.py new file mode 100644 index 000000000..f4098272c --- /dev/null +++ b/src/winml/modelkit/models/hf/florence2.py @@ -0,0 +1,441 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Native Florence-2 split image-to-text export.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar, cast + +import torch +import torch.nn as nn +from optimum.exporters.onnx import OnnxConfig +from optimum.utils import NormalizedConfig +from optimum.utils.input_generators import DummyInputGenerator +from transformers import Florence2Config, Florence2ForConditionalGeneration, Florence2Processor +from transformers.cache_utils import DynamicCache, EncoderDecoderCache +from transformers.conversion_mapping import register_checkpoint_conversion_mapping +from transformers.core_model_loading import Transpose, WeightConverter, WeightRenaming + +from ...config import WinMLBuildConfig +from ...export import register_onnx_overwrite +from ...optim import WinMLOptimizationConfig +from ..winml.composite_model import register_composite_model +from ..winml.encoder_decoder import EncoderDecoderInputGenerator, WinMLEncoderDecoderModel +from ..winml.kv_cache import PastKeyValueInputGenerator, WinMLStaticCache +from .decoder_wrapper import WinMLDecoderWrapper, WinMLStaticCacheDecoderIOConfig + + +if TYPE_CHECKING: + from transformers import GenerationConfig, PretrainedConfig + + +FLORENCE2_CONFIG = WinMLBuildConfig( + optim=WinMLOptimizationConfig( + gelu_fusion=True, + layer_norm_fusion=True, + matmul_add_fusion=True, + ), +) + +_FLORENCE2_IMAGE_TOKEN = "" # noqa: S105 - model vocabulary token +_FLORENCE2_LEGACY_IMAGE_TOKEN_ID = 50265 +_FLORENCE2_CAPTION_TOKEN_IDS = (0, 2264, 473, 5, 2274, 6190, 116, 2) +_FLORENCE2_ENCODER_SEQUENCE_LENGTH = 585 +_PROCESSOR_HUB_KWARGS = { + "cache_dir", + "force_download", + "local_files_only", + "proxies", + "revision", + "subfolder", + "token", +} + + +def _legacy_florence2_weight_conversions() -> list[Any]: + return [ + WeightRenaming(r"\.convs\.(\d+)\.proj\.", r".convs.\1.conv."), + WeightRenaming(r"\.conv1\.fn\.dw\.", ".conv1."), + WeightRenaming(r"\.conv2\.fn\.dw\.", ".conv2."), + WeightRenaming(r"\.ffn\.fn\.net\.", ".ffn."), + WeightRenaming(r"\.window_attn\.fn\.", ".window_attn."), + WeightRenaming(r"\.window_attn\.norm\.", ".norm1."), + WeightRenaming(r"\.channel_attn\.fn\.", ".channel_attn."), + WeightRenaming(r"\.channel_attn\.norm\.", ".norm1."), + WeightRenaming(r"\.ffn\.norm\.", ".norm2."), + WeightRenaming(r"^vision_tower\.", "model.vision_tower."), + WeightRenaming(r"^language_model\.model\.", "model.language_model."), + WeightRenaming( + r"^image_pos_embed\.", + "model.multi_modal_projector.image_position_embed.", + ), + WeightRenaming( + r"^image_proj_norm\.", "model.multi_modal_projector.image_proj_norm." + ), + WeightRenaming( + r"^visual_temporal_embed\.", + "model.multi_modal_projector.visual_temporal_embed.", + ), + WeightConverter( + source_patterns="^image_projection$", + target_patterns="model.multi_modal_projector.image_projection.weight", + operations=[Transpose()], + ), + ] + + +class _WinMLFlorence2ForConditionalGeneration(Florence2ForConditionalGeneration): + _keys_to_ignore_on_load_unexpected = [ # noqa: RUF012 - Transformers class contract + r"language_model\.final_logits_bias" + ] + + +register_checkpoint_conversion_mapping( + _WinMLFlorence2ForConditionalGeneration.__name__, + _legacy_florence2_weight_conversions(), +) + + +def _load_florence2_processor(model_name_or_path: str, **kwargs: Any) -> Florence2Processor: + processor_kwargs = {key: value for key, value in kwargs.items() if key in _PROCESSOR_HUB_KWARGS} + return Florence2Processor.from_pretrained( + model_name_or_path, + extra_special_tokens={"image_token": _FLORENCE2_IMAGE_TOKEN}, + **processor_kwargs, + ) + + +def _load_florence2_model( + model_name_or_path: str, **kwargs: Any +) -> Florence2ForConditionalGeneration: + model_kwargs = dict(kwargs) + model_kwargs.setdefault("dtype", torch.float32) + config = model_kwargs.pop("config", None) + if config is None: + config_kwargs = { + key: value for key, value in model_kwargs.items() if key in _PROCESSOR_HUB_KWARGS + } + config = Florence2Config.from_pretrained(model_name_or_path, **config_kwargs) + if config.image_token_id >= config.text_config.vocab_size: + config.image_token_id = _load_florence2_processor( + model_name_or_path, **model_kwargs + ).image_token_id + model_kwargs.pop("output_loading_info", None) + model, loading_info = cast( + "tuple[Florence2ForConditionalGeneration, dict[str, Any]]", + cast("Any", _WinMLFlorence2ForConditionalGeneration.from_pretrained)( + model_name_or_path, + config=config, + output_loading_info=True, + **model_kwargs, + ), + ) + unresolved = { + key: loading_info[key] + for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") + if loading_info.get(key) + } + if unresolved: + raise RuntimeError( + f"Florence-2 checkpoint conversion left unresolved weights: {unresolved}" + ) + return model + + +class _Florence2EncoderNormalizedConfig(NormalizedConfig): # type: ignore[misc] + def __init__(self, config: Any, **kwargs: Any) -> None: + super().__init__(config, **kwargs) + self.num_channels = config.vision_config.in_channels + self.image_size = 768 + self.image_token_id = ( + config.image_token_id + if config.image_token_id < config.text_config.vocab_size + else _FLORENCE2_LEGACY_IMAGE_TOKEN_ID + ) + self.bos_token_id = config.text_config.bos_token_id + self.eos_token_id = config.text_config.eos_token_id + + +class _Florence2EncoderInputGenerator(DummyInputGenerator): # type: ignore[misc] + SUPPORTED_INPUT_NAMES = ("input_ids", "pixel_values", "attention_mask") + + def __init__(self, task: str, normalized_config: Any, **kwargs: Any) -> None: + del task + self.batch_size = kwargs.get("batch_size", 1) + self.image_size = normalized_config.image_size + self.num_channels = normalized_config.num_channels + self.image_token_id = normalized_config.image_token_id + self.bos_token_id = normalized_config.bos_token_id + self.eos_token_id = normalized_config.eos_token_id + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ) -> torch.Tensor: + del framework, int_dtype, float_dtype + sequence_length = _FLORENCE2_ENCODER_SEQUENCE_LENGTH + if input_name == "input_ids": + image_tokens = torch.full( + (self.batch_size, 577), self.image_token_id, dtype=torch.long + ) + caption_tokens = torch.tensor( + _FLORENCE2_CAPTION_TOKEN_IDS, dtype=torch.long + ).repeat(self.batch_size, 1) + return torch.cat((image_tokens, caption_tokens), dim=1) + if input_name == "pixel_values": + return torch.zeros( + (self.batch_size, self.num_channels, self.image_size, self.image_size), + dtype=torch.float32, + ) + if input_name == "attention_mask": + return torch.ones((self.batch_size, sequence_length), dtype=torch.long) + raise ValueError(f"Unknown input: {input_name}") + + +class Florence2EncoderWrapper(nn.Module): + """Export the native Florence image/text encoder as one ONNX component.""" + + def __init__(self, model: Florence2ForConditionalGeneration) -> None: + super().__init__() + self.model = model + self.config = model.config + + @classmethod + def from_pretrained(cls, model_name_or_path: str, **kwargs: Any) -> Florence2EncoderWrapper: + """Load converted native weights and return an evaluation-mode wrapper.""" + model = _load_florence2_model(model_name_or_path, **kwargs) + wrapper = cls(model) + wrapper.eval() + return wrapper + + def forward( + self, + input_ids: torch.Tensor, + pixel_values: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + """Merge image features into placeholder embeddings and run the encoder.""" + inputs_embeds = self.model.model.get_input_embeddings()(input_ids) + image_features = self.model.get_image_features(pixel_values).pooler_output + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + placeholder_mask = self.model.get_placeholder_mask( + cast("torch.LongTensor", input_ids), + inputs_embeds=inputs_embeds, + image_features=image_features, + ) + inputs_embeds = inputs_embeds.masked_scatter(placeholder_mask, image_features) + outputs = self.model.model.language_model.encoder( + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + return_dict=True, + ) + return cast("torch.Tensor", outputs.last_hidden_state) + + +@register_onnx_overwrite("florence2", "feature-extraction", library_name="transformers") +class Florence2EncoderIOConfig(OnnxConfig): # type: ignore[misc] + """Declare semantic Florence encoder inputs and hidden-state output.""" + + NORMALIZED_CONFIG_CLASS = _Florence2EncoderNormalizedConfig + DUMMY_INPUT_GENERATOR_CLASSES = (_Florence2EncoderInputGenerator,) + PRESERVE_DUMMY_VALUE_RUNS = True + + @property + def inputs(self) -> dict[str, dict[int, str]]: + """Return encoder input dynamic-axis metadata.""" + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + """Return encoder output dynamic-axis metadata.""" + return {"last_hidden_state": {0: "batch_size", 1: "sequence_length"}} + + +class _Florence2DecoderNormalizedConfig(NormalizedConfig): # type: ignore[misc] + def __init__(self, config: Any, **kwargs: Any) -> None: + super().__init__(config, **kwargs) + self._text_config = config.text_config + + @property + def hidden_size(self) -> int: + return cast("int", self._text_config.hidden_size) + + @property + def num_layers(self) -> int: + return cast("int", self._text_config.decoder_layers) + + @property + def num_attention_heads(self) -> int: + return cast("int", self._text_config.decoder_attention_heads) + + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_attention_heads + + @property + def max_cache_len(self) -> int: + return cast("int", self._text_config.max_position_embeddings) + + @property + def vocab_size(self) -> int: + return cast("int", self._text_config.vocab_size) + + +class _Florence2DecoderInputGenerator(EncoderDecoderInputGenerator): + def __init__(self, task: str, normalized_config: Any, **kwargs: Any) -> None: + kwargs.setdefault("sequence_length", _FLORENCE2_ENCODER_SEQUENCE_LENGTH) + super().__init__(task, normalized_config, **kwargs) + + +@register_onnx_overwrite("florence2", "text2text-generation", library_name="transformers") +class Florence2DecoderIOConfig(WinMLStaticCacheDecoderIOConfig): + """Declare the six-layer Florence decoder static-cache contract.""" + + NORMALIZED_CONFIG_CLASS = _Florence2DecoderNormalizedConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + _Florence2DecoderInputGenerator, + PastKeyValueInputGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + """Return decoder and past-cache input metadata.""" + result: dict[str, dict[int, str]] = { + "decoder_input_ids": {0: "batch_size"}, + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + "decoder_attention_mask": {0: "batch_size"}, + "cache_position": {}, + } + for index in range(self._normalized_config.num_layers): + result[f"past_{index}_key"] = {0: "batch_size"} + result[f"past_{index}_value"] = {0: "batch_size"} + return result + + @property + def outputs(self) -> dict[str, dict[int, str]]: + """Return logits and present-cache output metadata.""" + result: dict[str, dict[int, str]] = {"logits": {0: "batch_size"}} + for index in range(self._normalized_config.num_layers): + result[f"present_{index}_key"] = {0: "batch_size"} + result[f"present_{index}_value"] = {0: "batch_size"} + return result + + +class Florence2DecoderWrapper(WinMLDecoderWrapper): + """Export the native Florence language decoder with a static KV cache.""" + + _HF_MODEL_CLS = Florence2ForConditionalGeneration + _IO_CONFIG_CLS = Florence2DecoderIOConfig + + @classmethod + def from_pretrained(cls, model_name_or_path: str, **kwargs: Any) -> Florence2DecoderWrapper: + """Load converted native weights and initialize shared decoder state.""" + model = _load_florence2_model(model_name_or_path, **kwargs) + wrapper = cls() + wrapper.model = model + wrapper.config = model.config + wrapper.onnx_config = cls._IO_CONFIG_CLS(model.config, task=cls._TASK) + wrapper.num_layers = wrapper.onnx_config._normalized_config.num_layers + wrapper.eval() + return wrapper + + def _invoke_hf(self, cache: Any, inputs: dict[str, torch.Tensor]) -> torch.Tensor: + model = cast("Any", self.model) + encoder_hidden_states = inputs["encoder_hidden_states"] + encoder_attention_mask = torch.ones( + encoder_hidden_states.shape[:2], dtype=torch.long, device=encoder_hidden_states.device + ) + outputs = model.model.language_model.decoder( + input_ids=inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"], + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=EncoderDecoderCache(cache, DynamicCache()), + use_cache=True, + cache_position=inputs["cache_position"], + return_dict=True, + ) + return cast("torch.Tensor", model.lm_head(outputs.last_hidden_state)) + + +@register_composite_model("florence2", "image-text-to-text") +@register_composite_model("florence2", "image-to-text") +class WinMLFlorence2ImageToText(WinMLEncoderDecoderModel): + """Run split Florence encoder and decoder components for image-to-text.""" + + main_input_name = "input_ids" + _SUB_MODEL_CONFIG: ClassVar[dict[str, str]] = { + "encoder": "image-feature-extraction", + "decoder": "text2text-generation", + } + _SUB_MODEL_PRECISION_OVERRIDES: ClassVar[dict[str, str]] = {"encoder": "fp32"} + + def __init__( + self, + sub_models: dict[str, Any], + config: PretrainedConfig, + device: str = "cpu", + ) -> None: + super().__init__(sub_models, config, device) + self.config.is_encoder_decoder = True + self.config.num_hidden_layers = config.text_config.decoder_layers + + @classmethod + def get_cache_class(cls) -> type: + """Use a fixed-capacity cache matching the exported decoder graph.""" + return WinMLStaticCache + + def load_pipeline_processor(self, model_id: str) -> Florence2Processor: + """Load the native processor with the Florence image token registered.""" + return _load_florence2_processor(model_id) + + @property + def generation_config(self) -> GenerationConfig: + """Return deterministic greedy generation defaults within cache capacity.""" + if not hasattr(self, "_generation_config"): + from transformers import GenerationConfig + + text_config = self.config.text_config + self._generation_config = GenerationConfig( + decoder_start_token_id=text_config.decoder_start_token_id, + bos_token_id=text_config.bos_token_id, + eos_token_id=text_config.eos_token_id, + pad_token_id=text_config.pad_token_id, + max_new_tokens=self._max_dec - 1, + num_beams=1, + do_sample=False, + ) + return self._generation_config + + @generation_config.setter + def generation_config(self, value: Any) -> None: + self._generation_config = value + + +MODEL_CLASS_MAPPING: dict[tuple[str, str | None], type] = { + ("florence2", None): WinMLFlorence2ImageToText, + ("florence2", "image-to-text"): WinMLFlorence2ImageToText, + ("florence2", "image-text-to-text"): WinMLFlorence2ImageToText, + ("florence2", "feature-extraction"): Florence2EncoderWrapper, + ("florence2", "text2text-generation"): Florence2DecoderWrapper, +} + + +__all__ = [ + "FLORENCE2_CONFIG", + "MODEL_CLASS_MAPPING", + "Florence2DecoderIOConfig", + "Florence2DecoderWrapper", + "Florence2EncoderIOConfig", + "Florence2EncoderWrapper", + "WinMLFlorence2ImageToText", +] diff --git a/src/winml/modelkit/models/winml/composite_model.py b/src/winml/modelkit/models/winml/composite_model.py index 1455ba737..ddb2aded8 100644 --- a/src/winml/modelkit/models/winml/composite_model.py +++ b/src/winml/modelkit/models/winml/composite_model.py @@ -107,6 +107,7 @@ class WinMLCompositeModel(PreTrainedModel): """ _SUB_MODEL_CONFIG: ClassVar[dict[str, str]] = {} + _SUB_MODEL_PRECISION_OVERRIDES: ClassVar[dict[str, str]] = {} def __init__( self, diff --git a/src/winml/modelkit/models/winml/kv_cache.py b/src/winml/modelkit/models/winml/kv_cache.py index 1bcfa1fe4..d13e8ac87 100644 --- a/src/winml/modelkit/models/winml/kv_cache.py +++ b/src/winml/modelkit/models/winml/kv_cache.py @@ -31,7 +31,7 @@ - ``build_decoder_mask(max_len)``: 2D attention mask for current step - ``get_query_cache_position(max_len)``: buffer indices of query tokens (used by HF's ``create_causal_mask`` and by T5's ``compute_bias``) -- ``update_all_layers(outputs)``: write present KV from ONNX output, advance step +- ``update_all_layers(outputs)``: write present KV and advance logical step - ``reset()``: zero out for new generation - ``create(config, kv_shape, dtype)``: factory from ONNX metadata @@ -114,6 +114,12 @@ def set_trace_position(self, position: torch.Tensor) -> None: """Provide the position tensor when a model omits cache update kwargs.""" self._trace_position = position + def get_seq_length(self, layer_idx: int = 0) -> int: + """Return the logical token count rather than fixed buffer occupancy.""" + if self._trace_position is not None: + return cast("int", self._trace_position.reshape(-1)[0]) + return self.step + def early_initialization( self, batch_size: int, @@ -187,23 +193,38 @@ def prepare_prefill_chunk( """ def update_all_layers(self, outputs: dict[str, Any]) -> None: - """Write present KV for all layers via ``update()`` and advance step. - - Step advances by N where N is the seq_len of the present KV tensors - (1 for gen, chunk_len for prefill). - """ + """Write present KV and advance by the logical decoder query length.""" import torch - n = 0 + logits = outputs.get("logits") + num_new_tokens = ( + int(logits.shape[1]) + if logits is not None + else int(outputs["present_0_key"].shape[2]) + ) for i in range(self.num_layers): k = outputs[f"present_{i}_key"] v = outputs[f"present_{i}_value"] k = k if isinstance(k, torch.Tensor) else torch.tensor(k) v = v if isinstance(v, torch.Tensor) else torch.tensor(v) - n = k.size(2) - ck = {"cache_position": torch.arange(self.step, self.step + n, dtype=torch.int64)} - self.update(k, v, i, cache_kwargs=ck) - self.step += n + layer = self._layer(i) + layer_keys = cast("torch.Tensor", layer.keys) + layer_values = cast("torch.Tensor", layer.values) + if k.shape == layer_keys.shape and v.shape == layer_values.shape: + layer_keys.copy_(k) + layer_values.copy_(v) + continue + if k.size(2) != num_new_tokens or v.size(2) != num_new_tokens: + raise ValueError( + "Present KV sequence length must match either the cache capacity " + "or the number of new decoder tokens." + ) + cache_position = self.get_query_cache_position( + layer_keys.size(2), + num_new_tokens, + ) + self.update(k, v, i, cache_kwargs={"cache_position": cache_position}) + self.step += num_new_tokens def reset(self) -> None: """Zero out all layers and reset step (start of new generation).""" diff --git a/src/winml/modelkit/onnx/io.py b/src/winml/modelkit/onnx/io.py index d2ad85878..4c95016fb 100644 --- a/src/winml/modelkit/onnx/io.py +++ b/src/winml/modelkit/onnx/io.py @@ -21,6 +21,7 @@ import json import logging +import math from dataclasses import dataclass from pathlib import Path from typing import Any @@ -31,6 +32,8 @@ logger = logging.getLogger(__name__) ShapeDim = int | str +DummyValue = int | float +DummyValueRun = tuple[int, DummyValue] # ============================================================================= @@ -69,6 +72,28 @@ class InputTensorSpec: dtype: str | None = None # "float32", "float16", "int64", "int32", etc. shape: tuple[ShapeDim, ...] | None = None value_range: tuple[float, float] | None = None # (min, max_exclusive) + dummy_value_runs: tuple[DummyValueRun, ...] | None = None + + def __post_init__(self) -> None: + """Normalize and validate serialized semantic dummy values.""" + if self.dummy_value_runs is None: + return + + normalized_runs: list[DummyValueRun] = [] + for run in self.dummy_value_runs: + if not isinstance(run, (tuple, list)) or len(run) != 2: + raise TypeError("dummy_value_runs entries must be (count, value) pairs") + count, value = run + if not isinstance(count, int) or isinstance(count, bool): + raise TypeError("dummy_value_runs counts must be integers") + if count <= 0: + raise ValueError("dummy_value_runs counts must be positive") + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise TypeError("dummy_value_runs values must be numeric") + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("dummy_value_runs values must be finite") + normalized_runs.append((count, value)) + self.dummy_value_runs = tuple(normalized_runs) def to_tensor(self) -> Any: """Generate a dummy tensor from this spec. @@ -102,6 +127,26 @@ def to_tensor(self) -> Any: concrete_shape = self.concrete_shape() + if self.dummy_value_runs is not None: + if sum(count for count, _ in self.dummy_value_runs) != math.prod(concrete_shape): + raise ValueError( + f"dummy_value_runs for '{self.name}' must fill the concrete tensor shape" + ) + if not torch_dtype.is_floating_point and any( + not isinstance(value, int) for _, value in self.dummy_value_runs + ): + raise TypeError( + f"dummy_value_runs for integer tensor '{self.name}' must contain integers" + ) + + tensor = torch.empty(concrete_shape, dtype=torch_dtype) + flat_tensor = tensor.reshape(-1) + start = 0 + for count, value in self.dummy_value_runs: + flat_tensor[start : start + count] = value + start += count + return tensor + if self.value_range is not None: lo, hi = self.value_range if torch_dtype.is_floating_point: @@ -146,6 +191,30 @@ def concrete_shape(self) -> tuple[int, ...]: ) return tuple(concrete) + @staticmethod + def compact_dummy_value_runs(tensor: Any) -> tuple[DummyValueRun, ...] | None: + """Run-length encode deterministic tensors without serializing random inputs.""" + import torch + + if not isinstance(tensor, torch.Tensor) or tensor.numel() < 2: + return None + if tensor.dtype == torch.bool or tensor.is_complex(): + return None + + flat_tensor = tensor.detach().cpu().reshape(-1) + values, counts = torch.unique_consecutive(flat_tensor, return_counts=True) + if values.numel() >= flat_tensor.numel(): + return None + + runs: list[DummyValueRun] = [] + for count, value in zip(counts.tolist(), values.tolist(), strict=True): + if not isinstance(value, (int, float)) or isinstance(value, bool): + return None + if isinstance(value, float) and not math.isfinite(value): + return None + runs.append((count, value)) + return tuple(runs) + def to_dict(self) -> dict[str, Any]: """Convert to dictionary, excluding None values.""" result: dict[str, Any] = {} @@ -157,6 +226,8 @@ def to_dict(self) -> dict[str, Any]: result["shape"] = self.shape if self.value_range is not None: result["value_range"] = list(self.value_range) + if self.dummy_value_runs is not None: + result["dummy_value_runs"] = [list(run) for run in self.dummy_value_runs] return result @classmethod @@ -172,6 +243,7 @@ def from_dict(cls, data: dict[str, Any]) -> InputTensorSpec: dtype=data.get("dtype"), shape=shape, value_range=value_range, + dummy_value_runs=data.get("dummy_value_runs"), ) diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 37f6bde79..eceeadeac 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -2394,6 +2394,68 @@ def _stamp(cfg: MagicMock, _path: Path) -> None: class TestBuildComposite: """Test build fans a composite model out with flat _ naming.""" + def test_composite_component_precision_override_is_preserved( + self, + runner: CliRunner, + tmp_path: Path, + ) -> None: + """A component override wins without changing sibling precision.""" + from winml.modelkit.commands.build import build + + components = { + "encoder": "image-feature-extraction", + "decoder": "text2text-generation", + } + generated_configs: list[MagicMock] = [] + + def fake_generate(*args, **kwargs): + cfg = MagicMock() + cfg.loader = MagicMock(task=kwargs.get("task"), model_type="florence2") + cfg.quant = None if kwargs["precision"] == "fp32" else MagicMock() + cfg.precision = kwargs["precision"] + cfg.compile = None + generated_configs.append(cfg) + return cfg + + with ( + patch( + "winml.modelkit.loader.resolution.resolve_composite_components", + return_value=components, + ), + patch( + "winml.modelkit.loader.resolution.resolve_composite_precision_overrides", + return_value={"encoder": "fp32"}, + ), + patch( + "winml.modelkit.config.generate_build_config", + side_effect=fake_generate, + ) as mock_generate, + patch("winml.modelkit.commands.build._run_single_build") as mock_build, + patch( + "winml.modelkit.commands.build._validate_loader_tasks_for_model", + return_value=None, + ), + ): + result = runner.invoke( + build, + [ + "-m", + "microsoft/Florence-2-base", + "-o", + str(tmp_path / "out"), + "--precision", + "fp16", + ], + obj={"debug": False}, + ) + + assert result.exit_code == 0, result.output + component_calls = mock_generate.call_args_list[1:] + assert [call.kwargs["precision"] for call in component_calls] == ["fp32", "fp16"] + built_configs = [call.kwargs["config"] for call in mock_build.call_args_list] + assert built_configs[0].quant is None + assert [cfg.precision for cfg in built_configs] == ["fp32", "fp16"] + def test_composite_builds_flat_with_cache_key( self, runner: CliRunner, @@ -2449,6 +2511,7 @@ def test_composite_builds_flat_with_cache_key( assert built_keys == set(components) # generate_build_config is called once for the outer auto-gen config, # then once per component. + assert mock_gen_cfg.call_args_list[0].kwargs["task"] == next(iter(components.values())) component_tasks = { call.kwargs["task"] for call in mock_gen_cfg.call_args_list @@ -2506,10 +2569,9 @@ def test_composite_autogen_passes_task_none_to_resolver( ) assert result.exit_code == 0, result.output - # No -c: task must be None (so the seq2seq bridge is applied), but - # model_type is still forwarded. + # No -c: task must be None so the seq2seq bridge is applied before + # generating the outer policy config. assert mock_resolve.call_args.kwargs["task"] is None - assert mock_resolve.call_args.kwargs["model_type"] == "t5" def test_composite_config_file_forwards_explicit_task( self, diff --git a/tests/unit/eval/test_image_to_text_evaluator.py b/tests/unit/eval/test_image_to_text_evaluator.py index a1f0b5d03..2d6c843e3 100644 --- a/tests/unit/eval/test_image_to_text_evaluator.py +++ b/tests/unit/eval/test_image_to_text_evaluator.py @@ -13,10 +13,8 @@ from winml.modelkit.inference.pipeline import _HF_PIPELINE_TASK_MAP -def make_evaluator(columns_mapping=None): +def make_evaluator(columns_mapping=None, prompt=None): """Instantiate evaluator with mocked dataset + pipeline.""" - import transformers - from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig mapping = columns_mapping or {} @@ -37,20 +35,29 @@ def make_evaluator(columns_mapping=None): config = WinMLEvaluationConfig( model_id="microsoft/trocr-base-handwritten", task="image-to-text", + prompt=prompt, dataset=DatasetConfig(path="Teklia/IAM-line", columns_mapping=mapping), ) - # Resolve the lazy Transformers export before patching it. - assert hasattr(transformers, "pipeline") with ( patch("datasets.load_dataset", return_value=mock_ds), - patch("transformers.pipelines.pipeline", return_value=mock_pipe), - patch.object(transformers, "pipeline", return_value=mock_pipe), + patch( + "winml.modelkit.eval.base_evaluator.WinMLEvaluator.prepare_pipeline", + return_value=mock_pipe, + ), ): return WinMLImageToTextEvaluator(config, model) class TestInit: + def test_config_serializes_prompt(self): + from winml.modelkit.eval import WinMLEvaluationConfig + + config = WinMLEvaluationConfig(task="image-to-text", prompt="") + + assert config.to_dict()["prompt"] == "" + assert WinMLEvaluationConfig.from_dict(config.to_dict()).prompt == "" + def test_uses_transformers_pipeline_task_name(self): assert _HF_PIPELINE_TASK_MAP["image-to-text"] == "image-text-to-text" @@ -115,6 +122,15 @@ def test_perfect_predictions(self): assert result["n_samples"] == 2 assert "cider" in result + def test_passes_configured_prompt_as_native_pipeline_text(self): + ev = make_evaluator(prompt="") + ev.data = [{"image": "img1", "text": "caption"}] + ev.pipe = MagicMock(return_value=[{"generated_text": "caption"}]) + + ev.compute() + + ev.pipe.assert_called_once_with("img1", text="") + def test_dict_output_shape(self): """Pipeline may also return a single dict (not a list).""" ev = make_evaluator() diff --git a/tests/unit/export/test_config_validation.py b/tests/unit/export/test_config_validation.py index d1f1ed621..3014ac095 100644 --- a/tests/unit/export/test_config_validation.py +++ b/tests/unit/export/test_config_validation.py @@ -13,6 +13,7 @@ import logging import pytest +import torch from winml.modelkit.export import ( InputTensorSpec, @@ -271,6 +272,29 @@ def test_name_only_roundtrip(self): assert restored.dtype is None assert restored.shape is None + def test_semantic_dummy_values_roundtrip(self): + original = InputTensorSpec( + name="input_ids", + dtype="int64", + shape=(1, 5), + dummy_value_runs=((3, 42), (1, 0), (1, 2)), + ) + restored = InputTensorSpec.from_dict(original.to_dict()) + + assert restored.dummy_value_runs == original.dummy_value_runs + assert torch.equal(restored.to_tensor(), torch.tensor([[42, 42, 42, 0, 2]])) + + def test_semantic_dummy_values_must_fill_shape(self): + spec = InputTensorSpec( + name="input_ids", + dtype="int64", + shape=(1, 4), + dummy_value_runs=((3, 42),), + ) + + with pytest.raises(ValueError, match="must fill the concrete tensor shape"): + spec.to_tensor() + # ============================================================================= # 9. InputTensorSpec.from_dict converts list shape to tuple diff --git a/tests/unit/export/test_florence2_onnx_config.py b/tests/unit/export/test_florence2_onnx_config.py new file mode 100644 index 000000000..4dc656b71 --- /dev/null +++ b/tests/unit/export/test_florence2_onnx_config.py @@ -0,0 +1,254 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for native Florence-2 split image-to-text export.""" + +from unittest.mock import patch + +import torch +from optimum.exporters.tasks import TasksManager +from transformers import Florence2Config + +from winml.modelkit.export import generate_dummy_inputs, resolve_io_specs +from winml.modelkit.export.config import _resolve_export_config_from_specs +from winml.modelkit.loader.resolution import ( + resolve_composite, + resolve_composite_components, + resolve_composite_precision_overrides, + resolve_task, +) +from winml.modelkit.models import HF_MODEL_CLASS_MAPPING +from winml.modelkit.models.hf.florence2 import ( + Florence2DecoderWrapper, + WinMLFlorence2ImageToText, + _legacy_florence2_weight_conversions, + _load_florence2_model, + _WinMLFlorence2ForConditionalGeneration, +) + + +def _convert_legacy_key(key: str) -> str: + for conversion in _legacy_florence2_weight_conversions(): + key, _ = conversion.rename_source_key(key) + return key + + +def test_florence2_legacy_checkpoint_key_conversion() -> None: + assert _convert_legacy_key("vision_tower.convs.0.proj.weight") == ( + "model.vision_tower.convs.0.conv.weight" + ) + assert _convert_legacy_key( + "vision_tower.blocks.0.0.spatial_block.window_attn.fn.qkv.weight" + ) == "model.vision_tower.blocks.0.0.spatial_block.window_attn.qkv.weight" + assert _convert_legacy_key( + "vision_tower.blocks.0.0.channel_block.ffn.norm.weight" + ) == "model.vision_tower.blocks.0.0.channel_block.norm2.weight" + assert _convert_legacy_key("language_model.model.encoder.layers.0.fc1.weight") == ( + "model.language_model.encoder.layers.0.fc1.weight" + ) + assert _convert_legacy_key("image_pos_embed.row_embeddings.weight") == ( + "model.multi_modal_projector.image_position_embed.row_embeddings.weight" + ) + assert _convert_legacy_key("image_projection") == ( + "model.multi_modal_projector.image_projection.weight" + ) + + +def test_florence2_legacy_checkpoint_loads_float32_source_model() -> None: + config = Florence2Config() + config.image_token_id = 50265 + config.text_config.vocab_size = 51289 + expected_model = object() + loading_info = { + "missing_keys": [], + "unexpected_keys": [], + "mismatched_keys": [], + "error_msgs": [], + } + + with patch.object( + _WinMLFlorence2ForConditionalGeneration, + "from_pretrained", + return_value=(expected_model, loading_info), + ) as from_pretrained: + model = _load_florence2_model("florence-checkpoint", config=config) + + assert model is expected_model + assert from_pretrained.call_args.kwargs["dtype"] is torch.float32 + + +def test_build_config_roundtrips_explicit_precision() -> None: + from winml.modelkit.config import WinMLBuildConfig + + config = WinMLBuildConfig(precision="fp32") + + assert config.to_dict()["precision"] == "fp32" + assert WinMLBuildConfig.from_dict(config.to_dict()).precision == "fp32" + + +def test_florence2_decoder_wrapper_initializes_shared_export_state() -> None: + config = Florence2Config() + loaded_model = torch.nn.Module() + loaded_model.config = config + + with patch( + "winml.modelkit.models.hf.florence2._load_florence2_model", + return_value=loaded_model, + ): + wrapper = Florence2DecoderWrapper.from_pretrained("florence-checkpoint") + + assert wrapper.model is loaded_model + assert wrapper.config is config + assert wrapper.num_layers == config.text_config.decoder_layers + assert wrapper.training is False + + +def test_florence2_family_registrations() -> None: + config = Florence2Config() + config._name_or_path = "microsoft/Florence-2-base" + + assert ("florence2", None) in HF_MODEL_CLASS_MAPPING + assert ("florence2", "image-text-to-text") in HF_MODEL_CLASS_MAPPING + assert ("florence2", "image-to-text") in HF_MODEL_CLASS_MAPPING + assert ("florence2", "feature-extraction") in HF_MODEL_CLASS_MAPPING + assert ("florence2", "text2text-generation") in HF_MODEL_CLASS_MAPPING + assert resolve_composite("florence2", "image-text-to-text") == { + "encoder": "image-feature-extraction", + "decoder": "text2text-generation", + } + assert resolve_composite_components(None, model_type="florence2") == { + "encoder": "image-feature-extraction", + "decoder": "text2text-generation", + } + assert resolve_composite_precision_overrides( + "florence2", + { + "encoder": "image-feature-extraction", + "decoder": "text2text-generation", + }, + ) == {"encoder": "fp32"} + assert resolve_task(config).task == "image-to-text" + + encoder = TasksManager.get_exporter_config_constructor( + exporter="onnx", + model_type="florence2", + task="feature-extraction", + library_name="transformers", + ) + decoder = TasksManager.get_exporter_config_constructor( + exporter="onnx", + model_type="florence2", + task="text2text-generation", + library_name="transformers", + ) + + assert encoder.func.__name__ == "Florence2EncoderIOConfig" + assert decoder.func.__name__ == "Florence2DecoderIOConfig" + + +def test_florence2_encoder_preserves_image_placeholder_tokens() -> None: + config = Florence2Config() + generated = generate_dummy_inputs("florence2", "feature-extraction", config) + specs = resolve_io_specs("florence2", "feature-extraction", config) + export_config = _resolve_export_config_from_specs("florence2", "feature-extraction", config) + + assert generated["input_ids"].shape == (1, 585) + assert torch.count_nonzero(generated["input_ids"] == 50265) == 577 + assert generated["input_ids"][0, 577:].tolist() == [0, 2264, 473, 5, 2274, 6190, 116, 2] + assert specs["dummy_value_runs"]["input_ids"][0] == (577, 50265) + assert torch.equal(export_config.generate_dummy_inputs()["input_ids"], generated["input_ids"]) + + +def test_florence2_decoder_declares_six_cache_layers() -> None: + config = Florence2Config( + text_config={ + "d_model": 768, + "decoder_attention_heads": 12, + "decoder_layers": 6, + "vocab_size": 51289, + } + ) + specs = resolve_io_specs("florence2", "text2text-generation", config) + generated = generate_dummy_inputs("florence2", "text2text-generation", config) + + assert specs["input_names"][:4] == [ + "decoder_input_ids", + "encoder_hidden_states", + "decoder_attention_mask", + "cache_position", + ] + assert [name for name in specs["input_names"] if name.endswith("_key")] == [ + f"past_{index}_key" for index in range(6) + ] + assert generated["encoder_hidden_states"].shape == (1, 585, 768) + assert specs["output_names"][0] == "logits" + + +def test_florence2_generation_keeps_encoder_prompt_out_of_decoder_history() -> None: + model = object.__new__(WinMLFlorence2ImageToText) + model.config = Florence2Config() + model.config.is_encoder_decoder = True + encoder_input_ids = torch.full((1, 585), 50265, dtype=torch.long) + pixel_values = torch.zeros(1, 3, 768, 768) + + inputs, input_name, model_kwargs = model._prepare_model_inputs( + None, + torch.tensor(model.config.text_config.bos_token_id), + { + "input_ids": encoder_input_ids, + "pixel_values": pixel_values, + }, + ) + decoder_input_ids, model_kwargs = model._prepare_decoder_input_ids_for_generation( + batch_size=1, + model_input_name=input_name, + model_kwargs=model_kwargs, + decoder_start_token_id=torch.tensor( + model.config.text_config.decoder_start_token_id + ), + device=encoder_input_ids.device, + ) + + assert input_name == "input_ids" + assert inputs is encoder_input_ids + assert model_kwargs["pixel_values"] is pixel_values + assert decoder_input_ids.tolist() == [[model.config.text_config.decoder_start_token_id]] + + +def test_full_static_cache_output_advances_by_logical_query_length() -> None: + from transformers import PretrainedConfig + + from winml.modelkit.models.winml.kv_cache import WinMLStaticCache + + config = PretrainedConfig() + config.num_hidden_layers = 1 + cache = WinMLStaticCache.create(config, [1, 12, 1024, 64], torch.float32) + present_key = torch.ones(1, 12, 1024, 64) + present_value = torch.full_like(present_key, 2) + + cache.update_all_layers( + { + "logits": torch.zeros(1, 1, 10), + "present_0_key": present_key, + "present_0_value": present_value, + } + ) + + assert cache.step == 1 + assert cache.get_seq_length() == 1 + assert torch.equal(cache._layer(0).keys, present_key) + assert torch.equal(cache._layer(0).values, present_value) + + +def test_static_cache_restores_exported_logical_position() -> None: + from transformers import PretrainedConfig + + from winml.modelkit.models.winml.kv_cache import WinMLStaticCache + + config = PretrainedConfig() + config.num_hidden_layers = 1 + cache = WinMLStaticCache.create(config, [1, 12, 1024, 64], torch.float32) + cache.set_trace_position(torch.tensor([7], dtype=torch.int64)) + + assert cache.get_seq_length() == 7 diff --git a/tests/unit/inference/test_pipeline.py b/tests/unit/inference/test_pipeline.py index 7755ecdb6..8b6357d43 100644 --- a/tests/unit/inference/test_pipeline.py +++ b/tests/unit/inference/test_pipeline.py @@ -16,7 +16,7 @@ import inspect import warnings from types import SimpleNamespace -from typing import Any +from typing import Any, ClassVar from unittest.mock import MagicMock, patch import pytest @@ -182,6 +182,30 @@ class ProcessorOnlyPipeline: class TestCreatePipeline: + def test_uses_model_supplied_pipeline_processor(self) -> None: + processor = MagicMock() + + class ModelWithProcessor: + io_config: ClassVar[dict[str, Any]] = {} + + def load_pipeline_processor(self, model_id: str) -> Any: + assert model_id == "test-model" + return processor + + pipe = MagicMock() + pipe.tokenizer = None + pipe.image_processor = None + with ( + patch( + "winml.modelkit.inference.pipeline._pipeline_component_kwargs", + return_value={"processor": "test-model"}, + ), + patch("transformers.pipeline", return_value=pipe) as pipeline, + ): + create_pipeline("image-to-text", ModelWithProcessor(), "test-model") + + assert pipeline.call_args.kwargs["processor"] is processor + def test_constructs_real_transformers_pipeline(self) -> None: from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor from transformers.pipelines import ImageClassificationPipeline