diff --git a/.azure_pipelines/job_templates/olive-test-cpu-template.yaml b/.azure_pipelines/job_templates/olive-test-cpu-template.yaml index 589da1b364..bdde89da8d 100644 --- a/.azure_pipelines/job_templates/olive-test-cpu-template.yaml +++ b/.azure_pipelines/job_templates/olive-test-cpu-template.yaml @@ -42,7 +42,13 @@ jobs: - script: | python -m pip install pytest python -m pip install -r $(Build.SourcesDirectory)/test/$(requirements_file) + displayName: Install test dependencies + + - script: | python -m pip list + displayName: Show environment (pip list) + + - script: | coverage run --source=$(Build.SourcesDirectory)/olive -m pytest -v -s -p no:warnings --disable-warnings --log-cli-level=WARNING --junitxml=$(Build.SourcesDirectory)/logs/test-TestOlive.xml -m "$(pytest_marker)" $(Build.SourcesDirectory)/$(test_path) --basetemp $(PYTEST_BASETEMP) coverage xml displayName: Test Olive diff --git a/.azure_pipelines/olive-ci.yaml b/.azure_pipelines/olive-ci.yaml index cc132d3d2e..5c7580e571 100644 --- a/.azure_pipelines/olive-ci.yaml +++ b/.azure_pipelines/olive-ci.yaml @@ -12,6 +12,7 @@ trigger: - mcp/** - notebooks/** - scripts/** + - skills/** pr: branches: include: @@ -27,6 +28,7 @@ pr: - mcp/** - notebooks/** - scripts/** + - skills/** variables: runCodesignValidationInjection: false diff --git a/.azure_pipelines/scripts/run_test.sh b/.azure_pipelines/scripts/run_test.sh index caa5fcb1c4..45baaaba00 100644 --- a/.azure_pipelines/scripts/run_test.sh +++ b/.azure_pipelines/scripts/run_test.sh @@ -42,7 +42,9 @@ BUILD_CUDA_EXT=0 pip install --no-build-isolation "git+https://github.com/PanQiW pip install huggingface-hub hf auth login --token "$7" +echo "===== Environment (pip list) =====" pip list +echo "===================================" # Step 4: Run tests with or without coverage tracking XML_PATH="/logs/TestOlive.xml" diff --git a/.github/workflows/test-model-fast.yml b/.github/workflows/test-model-fast.yml index fb03db2e10..138da791ab 100644 --- a/.github/workflows/test-model-fast.yml +++ b/.github/workflows/test-model-fast.yml @@ -7,9 +7,15 @@ on: push: branches: - main + paths-ignore: + - "mcp/**" + - "skills/**" pull_request: branches: - main + paths-ignore: + - "mcp/**" + - "skills/**" jobs: ubuntu-test-model-fast: diff --git a/README.md b/README.md index 8473c08c5b..5b495c938e 100644 --- a/README.md +++ b/README.md @@ -115,7 +115,8 @@ The sample chat app to run is found as [model-chat.py](https://github.com/micros - [Recipes](https://github.com/microsoft/olive-recipes) ## Data/Telemetry -Distributions of this project may collect usage data and send it to Microsoft to help improve our products and services. See the [privacy statement](docs/Privacy.md) for more details. + +This project may collect usage data and send it to Microsoft to help improve our products and services. See the [privacy statement](docs/Privacy.md) for details. ## 🤝 Contributions and Feedback - We welcome contributions! Please read the [contribution guidelines](./CONTRIBUTING.md) for more details on how to contribute to the Olive project. @@ -133,4 +134,3 @@ Licensed under the [MIT](./LICENSE) License. [![Build Status](https://dev.azure.com/aiinfra/PublicPackages/_apis/build/status%2FOlive%20CI?label=Olive-CI)](https://dev.azure.com/aiinfra/PublicPackages/_build/latest?definitionId=1240) [![Build Status](https://dev.azure.com/aiinfra/PublicPackages/_apis/build/status%2FOlive-ORT-Nightly?label=Olive-ORT-Nightly)](https://dev.azure.com/aiinfra/PublicPackages/_build/latest?definitionId=1279) - diff --git a/docs/Privacy.md b/docs/Privacy.md index 95aee00b0b..cc0d1defbc 100644 --- a/docs/Privacy.md +++ b/docs/Privacy.md @@ -6,11 +6,28 @@ The software may collect information about you and your use of the software and *** ## Technical Details -Olive uses the [OpenTelemetry](https://opentelemetry.io/) API for its implementation. Telemetry is turned ON by default. Based on user consent, this data may be periodically sent to Microsoft servers following GDPR and privacy regulations for anonymity and data access controls. Application, device, and version information is collected automatically. +Telemetry is turned ON by default. Based on user consent, this data may be periodically sent to Microsoft servers following GDPR and privacy regulations for anonymity and data access controls. Application, device, and version information is collected automatically. In addition, Olive may collect additional telemetry data such as: - Invoked commands - Performance data - Exception information -Collection of this additional telemetry can be disabled by adding the `--disable_telemetry` flag to any Olive CLI command, or by setting the `OLIVE_DISABLE_TELEMETRY` environment variable to `1` before running. Telemetry is also automatically disabled when a CI/CD environment is detected (e.g., GitHub Actions, Azure Pipelines, Jenkins). If telemetry is enabled, but cannot be sent to Microsoft, it will be stored locally and sent when a connection is available. You can override the default cache location by setting the `OLIVE_TELEMETRY_CACHE_DIR` environment variable to a valid directory path. +You can fully disable telemetry by adding the `--disable_telemetry` flag to any Olive CLI command, setting `OLIVE_DISABLE_TELEMETRY=1` or `ORT_DISABLE_TELEMETRY=1` before running, or calling `olive.telemetry.disable_telemetry()`. Each option suppresses every subsequent Olive telemetry event for the remainder of the process, including Olive workflow containers started by that process. When the opt-out is active before first telemetry use, Olive does not construct the telemetry singleton or create the telemetry queue, uploader, or persistent device identifier. Disabling at runtime stops this process's uploader, retains already queued unsent rows unchanged for a later telemetry-enabled process, and does not enqueue another Heartbeat. The environment variables accept `1`, `true`, `yes`, `on`, or `y` after trimming and without regard to case. + +In CI/CD environments (e.g., GitHub Actions, Azure Pipelines, Jenkins), Olive suppresses the device-id heartbeat and the action/error events and only emits the `OliveRecipe` event. Any full opt-out takes precedence and sends nothing. The `OliveRecipe` event may include recipe metadata such as pass types, explicitly configured target settings, the host system type (including the default `LocalSystem` host) and any explicitly configured host accelerator settings, whether a custom package config was provided, a redacted snapshot of custom package-config overrides, and a redacted snapshot of explicitly supplied config overrides. + +Telemetry is implemented using only the Python standard library. In enabled local runs, one `OliveHeartbeat` per process and detailed events are written to a local per-user SQLite queue before a background uploader sends them to Microsoft over HTTPS. Olive first reserves a minimal Heartbeat durably, then adds available operating-system metadata before making it eligible for upload; if enrichment is interrupted, the minimal Heartbeat remains eligible for a later delivery attempt. CI recipe events use a separate recipe-only queue and receive a bounded shutdown delivery attempt. Events that cannot be sent remain in the applicable queue for a later run. The transmitted `deviceId` is the `c:`-prefixed SHA-256 hash of a shared persistent UUID; the raw UUID is not transmitted. + +### Event schemas + +All events include Olive-assigned `appName`, `LibraryVersion`, and `AppSessionGuid` values that event callers cannot override; `initTs` is included when supplied by the caller. Olive emits only the following event-specific fields: + +| Event | Fields | +| --- | --- | +| `OliveHeartbeat` | `deviceId`, `deviceIdStatus`; `os`, `osVersion`, `osRelease`, and `osArchitecture` when enrichment succeeds | +| `OliveAction` | `invokedFrom`, `actionName`, `durationMs`, `success` | +| `OliveError` | `exceptionType`, `exceptionMessage` | +| `OliveRecipe` | `recipeName`, `recipeHash`, `recipeSource`, `recipeFormat`, `recipeCommand`, `executionMode`, `workflowId`, `configOverrides`, `success`, `inputModelType`, `inputModelSource`, `modelTask`, `targetSystemType`, `targetDevice`, `targetExecutionProvider`, `targetExecutionProviders`, `hostSystemType`, `hostDevice`, `hostExecutionProvider`, `hostExecutionProviders`, `passTypes`, `passCount`, `dataConfigCount`, `searchEnabled`, `packageConfigProvided`, `packageConfigOverrides`, `isCI` | + +Free-text values, paths, URLs, query secrets, credential-bearing configuration keys, environment-variable values, and nested configuration metadata are recursively redacted at the serialization boundary and capped at 40,960 UTF-8 bytes. `recipeHash` is computed only after credential, environment-value, and path redaction. Error messages may contain sanitized exception and frame metadata but never source-code lines. diff --git a/docs/source/features/quantization.md b/docs/source/features/quantization.md index a5b5ec6fbe..50b6761693 100644 --- a/docs/source/features/quantization.md +++ b/docs/source/features/quantization.md @@ -71,6 +71,161 @@ This pass supports ONNX models and can quantize `MatMul` and `Gather` nodes to 4 } ``` +## PyTorch Native RTN + +The `Rtn` pass applies RTN weight quantization directly to a PyTorch (Hugging Face) model, before any ONNX +export. Unlike `OnnxBlockWiseRtnQuantization` (which operates on an already-exported ONNX graph), `Rtn` runs on +the `HfModelHandler` and replaces the weight *storage* of quantizable parameters in place with a quantized +tensor representation, while keeping the surrounding modules (`nn.Linear` / `nn.Embedding`, and MoE experts) +otherwise unchanged. This lets you compose it with other PyTorch quantization passes (see the `Gptq` example +below) and share settings via per-module `overrides`. + +By default `Rtn` quantizes `nn.Linear` weights and leaves the embeddings, the language-model head, and any +Mixture-of-Experts (MoE) experts at full precision. Three independent category flags opt those groups in: + +| Flag | Default | Effect when `true` | +| --- | --- | --- | +| `lm_head` | `false` | Also quantize the language-model head. | +| `embeds` | `false` | Also quantize the input embeddings. | +| `moe` | `false` | Also quantize MoE expert weights: classic per-expert `nn.ModuleList` experts (e.g. Mixtral, PhiMoE on older `transformers`) are always supported, and fused-expert modules are supported whenever the resolved experts class reports `is_transposed=False` (the K-last `(E, OUT, K)` layout used by most fused-experts architectures). Architectures that store a transposed `(E, K, OUT)` layout, such as gpt-oss, are rejected. Architectures whose experts class does not report `is_transposed` at all -- either because it predates the `transformers` fused-experts refactor, or because it has not adopted it (e.g. llama4, aria) -- are also rejected, since Olive cannot independently verify their layout. | + +The `moe` flag is **fail-closed**: when `moe` is `false`, every module under an experts subtree is skipped even +if it looks like a plain `nn.Linear`. If the model config indicates an MoE architecture but Olive cannot resolve +the experts subtree for that (unrecognized) architecture, the pass raises a clear error *before* modifying any +parameter, rather than silently quantizing the experts. + +Only weight parameters are quantized. Fused expert 2D bias parameters, when present, remain in full precision. + +### `moe` and ONNX export + +MoE quantization in Olive is **storage-only**: Olive does not export 3D fused-expert `QuantTensor`s to ONNX. +Attempting to `torch.onnx.export` a model with 3D-quantized experts raises a clear error directing you to +Mobius / ORT GenAI `ModelBuilder` for the experts. Non-MoE parts (attention projections, router/gate, +embeddings, lm_head) still export through the existing `MatMulNBits` / `GatherBlockQuantized` path. + +### `moe` and native PyTorch inference: force the `"eager"` experts implementation + +`transformers` lets a loaded MoE model pick its runtime forward strategy independently of the +checkpoint's on-disk layout, via `model.set_experts_implementation(...)` / `config._experts_implementation` +(`"eager"`, `"grouped_mm"`, `"batched_mm"`, ...). Some non-`"eager"` strategies (e.g. `"grouped_mm"`, which +`transformers` may auto-select even on CPU) call `weight.transpose(-2, -1)` on the fused-experts weight +before dispatching to their matmul kernel. Because Olive's 3D fused-expert `QuantTensor` is storage-only +(see above) and cannot represent a transpose without a lossy unpack/re-quantize round trip, this raises a +`RuntimeError` at inference time -- even for architectures whose checkpoint layout (`is_transposed=False`) +is fully supported for quantization. + +**Workaround**: after loading a `moe=True`-quantized checkpoint for native PyTorch inference (as opposed to +consuming it via Mobius / ORT GenAI `ModelBuilder`), force the eager path once, before running any forward +pass: + +```python +model.set_experts_implementation("eager") +``` + +This is tracked as a follow-up in [#2619](https://github.com/microsoft/Olive/issues/2619). + +### `modules_to_not_convert` and `overrides` + +`modules_to_not_convert` lists module-name patterns to exclude entirely, and `overrides` maps module-name +patterns to per-module `{"bits", "symmetric", "group_size"}` settings. Both accept two key styles: + +- **Plain strings** keep the existing Hugging Face semantics (substring match for `modules_to_not_convert`, + literal match for `overrides`). +- **`re:`-prefixed keys** are treated as regular expressions matched with `re.fullmatch` (e.g. + `"re:model\\.layers\\.\\d+\\.mlp\\..*"`). + +For safety, `re:` patterns are validated before use: overly long patterns and patterns with nested unbounded +quantifiers (catastrophic-backtracking / ReDoS shapes such as `(a+)+`) are rejected with a clear error. + +### Resolution order and override precedence + +When multiple rules could apply to the same target, the **first** rule that matches wins, in this order: + +1. `modules_to_not_convert` (hard exclude) +2. category flags (`lm_head` / `embeds` / `moe`) — hard excludes; `overrides` can never re-include what a + category flag skipped +3. `overrides` +4. pass-level defaults (`bits` / `group_size` / `sym`) + +When several `overrides` entries match the same target, precedence is **insertion order in the config, first +match wins** — not "longest / most specific pattern". Order your `overrides` from most specific to least +specific accordingly. + +### Example Configuration +```json +{ + "type": "Rtn", + "bits": 4, + "group_size": 128, + "sym": false, + "embeds": true, + "moe": true, + "overrides": { + "re:.*\\.experts\\..*": { "bits": 4, "group_size": 32 } + } +} +``` + +### Composing with `Gptq` + +`Rtn` can run on an already-quantized model, so you can quantize the transformer `nn.Linear` layers with a +calibration-based pass such as `Gptq` first, then cover the parts `Gptq` doesn't handle (embeddings, lm_head, +MoE experts) with `Rtn`: + +```json +[ + { "type": "Gptq" }, + { "type": "Rtn", "moe": true, "embeds": true } +] +``` + +The reverse order is not supported: calibration-based passes assume a clean full-precision starting point and +will reject an already-quantized model. + +### Migration note: removal of `QuantLinear` / `QuantEmbedding` + +The previous `nn.Module` wrappers `olive.common.quant.nn.QuantLinear` and `QuantEmbedding` have been **removed** +as a sanctioned breaking change. Quantized weights are now stored as a `QuantTensor` on the parameter itself +rather than by swapping the parent module. **Checkpoints produced by the old `QuantLinear` / `QuantEmbedding` +classes cannot be reloaded through Olive's own HF quantizer** — this includes both: + +- models persisted with `torch.save(model)` (pickling live `QuantLinear` / `QuantEmbedding` instances), and +- safetensors/state-dict checkpoints, since the buffer naming convention changed from bare `.qweight` / + `.scales` / `.qzeros` to `.weight_qweight` / `weight_scales` / `weight_qzeros`. + +There is no migration shim for either case (consistent with every prior packing-format change to this module). +Re-run the `Rtn` pass on the original full-precision model to regenerate a checkpoint in the current format. + +### 2-bit quantization is not exportable to ONNX + +`Rtn` supports `bits` in `{2, 4, 8}` for the PyTorch quantized-checkpoint path, but the ONNX export-compat path +(`QuantLinearNbit`) only supports 4-bit and 8-bit packing. Attempting to export a 2-bit `QuantTensor` to ONNX +raises a clear `ValueError` at export time rather than silently producing an incorrect graph; 2-bit quantization +remains usable for PyTorch-only workflows. + +## PyTorch Native KQuant + +The `KQuant` pass is a calibration-free weight quantizer that applies llama.cpp's +iterative weighted-least-squares k-quant search to PyTorch (Hugging Face) model +weights. It supports the same `lm_head`, `embeds`, and `moe` category flags as +`Rtn`; all default to `false`. + +With `moe=true`, classic per-expert `nn.ModuleList` layouts are supported, as are +fused expert weights whose experts module reports `is_transposed=false` (K-last +`(E, OUT, K)`). Transposed `(E, K, OUT)` layouts and fused implementations with a +missing or non-boolean `is_transposed` attribute are rejected rather than risking +quantization along the wrong dimension. Only direct 3D expert weight parameters +are quantized; expert biases remain in full precision. + +```json +{ + "type": "KQuant", + "bits": 4, + "group_size": 32, + "moe": true +} +``` + ## HQQ `HQQ (Half-Quadratic Quantization)` is a fast, calibration-free weight quantization method that enables low-bit quantization of large models without relying on gradient-based optimization. Unlike data-dependent approaches like GPTQ, [HQQ](https://dropbox.github.io/hqq_blog/) uses half-quadratic splitting to minimize weight quantization error efficiently. @@ -365,4 +520,3 @@ Configurations: ``` Please refer to [AimetQuantization](aimet_quantization) for more details about the pass and its config parameters. - diff --git a/mcp/uv.lock b/mcp/uv.lock index dd3d4d32c0..5b5bf07764 100644 --- a/mcp/uv.lock +++ b/mcp/uv.lock @@ -163,62 +163,59 @@ wheels = [ [[package]] name = "cryptography" -version = "48.0.1" +version = "50.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/45/870e7f4bef50e5f53b9f51d4428aee5290eedf58ba443f16b1ebb7ab8e66/cryptography-48.0.1.tar.gz", hash = "sha256:266f4ee051abb2f725b74ef8072b521ce1feacf685a3364fa6a6b45548db791a", size = 832989, upload-time = "2026-06-09T22:32:31.8Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1b/bc/ee4137cbbe105652c0ee4252792b78fc8e7afa4b8e61d9d5dc05a7f45731/cryptography-48.0.1-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:3e4a1a3232eef2e6c732827d5722db29a0cc8b27af2a4d865b094cf954be9ca1", size = 8008324, upload-time = "2026-06-09T22:31:00.702Z" }, - { url = "https://files.pythonhosted.org/packages/d5/85/6379d42181bfc713094f081360fc5784d6c816b599d45e7f082502d173ce/cryptography-48.0.1-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:32143b24adb918f078134e1e230f1eb8cc04886b92c28b5f0041aaf3e5699225", size = 4696243, upload-time = "2026-06-09T22:32:33.446Z" }, - { url = "https://files.pythonhosted.org/packages/9c/87/c85d147b53323c7eb4d850920c8901377323c2a0ff8d79c262d4fee89aa2/cryptography-48.0.1-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f0d27a5696721ef7a672b8c810f6aded391058e0b9486e63e6d93baf765da691", size = 4713235, upload-time = "2026-06-09T22:31:40.141Z" }, - { url = "https://files.pythonhosted.org/packages/79/58/67cbf8cf1ee7c54b439ca07bbecf8362c07afc11a3724fea70f745784add/cryptography-48.0.1-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:eb86ce1af36fe65041b6db9a8bb064ee621a7e5fded0f80d475ec243477cd242", size = 4702323, upload-time = "2026-06-09T22:31:42.191Z" }, - { url = "https://files.pythonhosted.org/packages/89/c6/24266ac10c47f6cd2a865f4446062b466da1d1f10b27189eac00e61bf0c9/cryptography-48.0.1-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:b024e784ad6c077ee0147b35ea9cbfc1e34e1fd4c1dcca214c2794d73a12df08", size = 5300085, upload-time = "2026-06-09T22:31:58.703Z" }, - { url = "https://files.pythonhosted.org/packages/d2/bb/cc4b78784f97efc8c5874c2a9743708d172be6663024b34a0467885ae0c8/cryptography-48.0.1-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3752f2dbc8f07a30aad2932c986cea495b03bb554887828225da104f732852b6", size = 4746137, upload-time = "2026-06-09T22:31:31.01Z" }, - { url = "https://files.pythonhosted.org/packages/1f/52/0c44de3f5267f8fbe8e835138017522a333436166e406f0db9b9e6e3033f/cryptography-48.0.1-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:bd81490cd5801d755cf97bb68ac191f14b708470b1c7cf4580f669b9c9264cd8", size = 4333867, upload-time = "2026-06-09T22:32:28.096Z" }, - { url = "https://files.pythonhosted.org/packages/9a/2e/772d7adbfa931537bc401640b7cac9976bff689bda187833e5d63b428e49/cryptography-48.0.1-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:66fd0771e7b9c6dcd44cf1120690d2338d16d72795cf40cae2786a39eba65429", size = 4701805, upload-time = "2026-06-09T22:31:38.284Z" }, - { url = "https://files.pythonhosted.org/packages/f8/a3/b06844f303873493c963caf581c04df31c7035e0c1b0f02c4814d319ec80/cryptography-48.0.1-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:3fd2ca57062b241c856670b073487d2e86c4637937ca5601e48f97bf8e11fc8f", size = 5258461, upload-time = "2026-06-09T22:31:04.187Z" }, - { url = "https://files.pythonhosted.org/packages/9f/13/8b765e2e12b07c74941caadb9d1c8fdc006c4dfbf2b8f2d610519758954d/cryptography-48.0.1-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:0ee6ea481db1ab889cba043ec1eda17bb9c1ea79db6722f779c3667f9f70322f", size = 4745488, upload-time = "2026-06-09T22:32:30.07Z" }, - { url = "https://files.pythonhosted.org/packages/2e/aa/48972bce55049b32a94f4907eda4d75fa385aad8a39506cc2fc72196ecf0/cryptography-48.0.1-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:f2ceef93cb096aa3c4cc4b5c94ca6131f9196d28c64d6111533402a9b2054d41", size = 4830256, upload-time = "2026-06-09T22:31:43.868Z" }, - { url = "https://files.pythonhosted.org/packages/47/a2/e5079a032fb85cf6005046ca92bbd78b0c82dad2b5751ab8c311659da06f/cryptography-48.0.1-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:9bd3f92d76217892b15df84ca256c2c113d386fdda7a7d8691aeeced976507c6", size = 4979117, upload-time = "2026-06-09T22:31:05.845Z" }, - { url = "https://files.pythonhosted.org/packages/b7/a0/8f50cae9c74e718ed769d63ed5c74bd0ea830c9550a74629cebd1b9c7bc7/cryptography-48.0.1-cp311-abi3-win32.whl", hash = "sha256:b9a32b876490d66c8bcc9963ef220199569748434ab01a9d6aaeabf88e7f5158", size = 3304154, upload-time = "2026-06-09T22:32:16.845Z" }, - { url = "https://files.pythonhosted.org/packages/c5/69/0572c77dbace6fef72f33755bd52ea399c71367250d366237f8691826b9e/cryptography-48.0.1-cp311-abi3-win_amd64.whl", hash = "sha256:39489bfca54c7a1f6b297efcd8bc608ab92d16c4ca631b0cad4da46724588b24", size = 3817138, upload-time = "2026-06-09T22:32:00.388Z" }, - { url = "https://files.pythonhosted.org/packages/42/06/3e768b4c3bc78201583fa35a0e18f640dd782ff41afba88f8545481a8874/cryptography-48.0.1-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:f817adc181390bd54f2f700107a7419040fb7c1bdf2fc26f36551a06a68c3345", size = 7989830, upload-time = "2026-06-09T22:31:07.8Z" }, - { url = "https://files.pythonhosted.org/packages/8a/13/6476736484b94041110c8340a3eb63962fea4975baea8cb4a512adb44d4d/cryptography-48.0.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d5d30989c6917b478b5817902e85fddaea2261efa8648383d965381ccb9e1ac4", size = 4689201, upload-time = "2026-06-09T22:31:09.745Z" }, - { url = "https://files.pythonhosted.org/packages/79/62/65a87f34d2a431546e2509b85d55e8c90df86d668f6731da64d538512ac2/cryptography-48.0.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:df637c05205ea7c1d7fbcbe54bbfea648a52951155f997af13d895d0ecc96991", size = 4702822, upload-time = "2026-06-09T22:32:24.409Z" }, - { url = "https://files.pythonhosted.org/packages/7f/59/810b5204b0a9b10f4b6bc06bd551a8b609803cd931806bc3b71884b225e5/cryptography-48.0.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:869c3b8a53bfe27147832df48b32adadf558249d50e76cb3769d40e986b13265", size = 4694875, upload-time = "2026-06-09T22:32:08.737Z" }, - { url = "https://files.pythonhosted.org/packages/24/dc/d8ca05ffea724eec6d232ea6f18e74c269eb6bdfdcc9bfba689790d1325f/cryptography-48.0.1-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:e361afba8918070d376df76f408a4f67fec0ee9cff81a99e48fe9a233ef59e17", size = 5290385, upload-time = "2026-06-09T22:31:15.212Z" }, - { url = "https://files.pythonhosted.org/packages/03/8c/3be6cb4da181f5bb6c19cf560c2359d60644a6b5fc5b57854e528f47b296/cryptography-48.0.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:d069066deead00ac7f090be101be875a06855908f7ec004c27b8fefb4acfb411", size = 4737082, upload-time = "2026-06-09T22:32:22.66Z" }, - { url = "https://files.pythonhosted.org/packages/aa/f6/d5f60a5a1434dbfd949e227fd0065d194c7e6b6ac526b17f5c06152b8231/cryptography-48.0.1-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:09f73a725d582cef64b91281a322cd798d14a33b2b6f2b7ad9531dc336d84c02", size = 4325328, upload-time = "2026-06-09T22:32:10.777Z" }, - { url = "https://files.pythonhosted.org/packages/17/b7/ba75dd947a14b6ad907b01ae8f6b5b348cdd1b48142f0063dee9e20c1d9d/cryptography-48.0.1-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:15254441469dd6bf027039453288e2072124f8b6603563f5d759e1c9b69273fa", size = 4694530, upload-time = "2026-06-09T22:31:53.105Z" }, - { url = "https://files.pythonhosted.org/packages/62/29/50d6b9e8aff12d8b67afaeb3569335e32dc83a5723e3bbded24fdac9f809/cryptography-48.0.1-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:8ace4507d1e6533c125f4fac754f8bb8b6a74c08e92179dabd7e16571a3efbf3", size = 5245046, upload-time = "2026-06-09T22:31:25.774Z" }, - { url = "https://files.pythonhosted.org/packages/9f/04/618f4115cfc0add0838c82507aa18a346089428da8653ad38b3ff36f5cb3/cryptography-48.0.1-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:b4e391975f038e66432328639620a4aff2d307513b004f1ca06d6225bced815c", size = 4736660, upload-time = "2026-06-09T22:32:12.676Z" }, - { url = "https://files.pythonhosted.org/packages/24/9c/06e062462a0de28a3b3911322eded4c16deb9f441b1b7575d3dc59488ab5/cryptography-48.0.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:42fcd8e26fe555d9b3577a135f5091fefa0aa4e99129c23fb56787a1bd4ada72", size = 4822229, upload-time = "2026-06-09T22:31:17.062Z" }, - { url = "https://files.pythonhosted.org/packages/f4/be/0561971eaaee4b8a0e7d5113c536921063ab91aaf23278ac374eaf881e11/cryptography-48.0.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:c1400da5e32a43253392277eac7490a60e497d810a63dd5608d71bbd7af507c9", size = 4966364, upload-time = "2026-06-09T22:31:32.842Z" }, - { url = "https://files.pythonhosted.org/packages/a4/27/728c77876f12b000820b69ae490f3c4083775e79e07827e9e60be07ad209/cryptography-48.0.1-cp314-cp314t-win32.whl", hash = "sha256:0df56b056bc17c1b7d6821dfa65216e62bd232d8ab05eb3db44e71d235651471", size = 3278498, upload-time = "2026-06-09T22:31:29.154Z" }, - { url = "https://files.pythonhosted.org/packages/06/e3/79a612c6d7b1e6ee0edd43633d53035bec2cfb78c82b76f7864f39e36f34/cryptography-48.0.1-cp314-cp314t-win_amd64.whl", hash = "sha256:9de21387aa95e2a895823d0745b430bed4f33503ba9ab5e0b5311f33e37d66d2", size = 3798790, upload-time = "2026-06-09T22:31:56.697Z" }, - { url = "https://files.pythonhosted.org/packages/ca/6c/00fa2a95997164c8b2072ce327c23d4ab20809ccc323ea5fab91e53a4bba/cryptography-48.0.1-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:4fdc69f8e4316bcf0c8c8ec1f26f285d12e8142d88d96c876a59a03be3f6ae67", size = 7987408, upload-time = "2026-06-09T22:32:20.777Z" }, - { url = "https://files.pythonhosted.org/packages/b0/d9/45f309a7e4e5f3f8f121d6d3be9e94024a7726ec598d6e08ae04edb2f04d/cryptography-48.0.1-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:48fe40804d4caa2288f24e70ca8c64c42dd826da0ad7e4f1b41b2128d679e6c8", size = 4690196, upload-time = "2026-06-09T22:31:54.74Z" }, - { url = "https://files.pythonhosted.org/packages/5f/9f/a1bc8bcc798811b8527eb374bbccf30a3f3e806829d967118222bf1125eb/cryptography-48.0.1-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:86be3b1b0b6bf09482fb50a979c508d2950ed95f5621ec77f4e385962006b83a", size = 4696782, upload-time = "2026-06-09T22:31:45.615Z" }, - { url = "https://files.pythonhosted.org/packages/66/c2/81a4fb4e4373c500bb526bc337ac5719dd31dd15b970b84a238168c6aa08/cryptography-48.0.1-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:4ab0a343c807bbcd90c971cd1ecf072937cd01847a9e002bef88fb47ac6be577", size = 4696618, upload-time = "2026-06-09T22:31:11.564Z" }, - { url = "https://files.pythonhosted.org/packages/e5/0b/aa68b221dde92d09cb29a024ede17550ee21e77a404e59fc093c82bb51e1/cryptography-48.0.1-cp39-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:9621de99d2da096006b629979efd8ae7eb2d8b822488d0c89ee4000c306c59b1", size = 5289970, upload-time = "2026-06-09T22:31:20.368Z" }, - { url = "https://files.pythonhosted.org/packages/78/13/fba657f958d2af66ea959a4ba01212632089249d34af1ae48054136344d7/cryptography-48.0.1-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:88c852a0ae366e262e5a1744b685e6a433dc8788dd2a277e418bf4904203609d", size = 4731873, upload-time = "2026-06-09T22:31:22.253Z" }, - { url = "https://files.pythonhosted.org/packages/4c/4c/9a964756d24a26b3e34dfcb16f961b89838786e6700b635b0d1e3adff4b6/cryptography-48.0.1-cp39-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:43c5835e2cb98c8733d86f57d6fc879b613f5c3478607281c3e36daffc6dd8a6", size = 4330804, upload-time = "2026-06-09T22:31:36.56Z" }, - { url = "https://files.pythonhosted.org/packages/4b/0f/a10f3a6eb12950a10e3a874070283aa2dd5875b2bfd15fad8a3e17b3f13e/cryptography-48.0.1-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:fe0180af5bf9236518a087e35bf2d9a347d5f5f51e63c579d683ddff424e3d46", size = 4696217, upload-time = "2026-06-09T22:31:13.351Z" }, - { url = "https://files.pythonhosted.org/packages/f3/6f/5cd12f951165ea73ef85266775d97e4c763b2474ccfd816dd69d3a18d6f8/cryptography-48.0.1-cp39-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:b7a2d1a937a738a881737cec135a38bb61470589b17515b9f73f571d0ae10401", size = 5245252, upload-time = "2026-06-09T22:32:02.193Z" }, - { url = "https://files.pythonhosted.org/packages/68/ab/8aaa12e4516ec4464033ab79b6f3b592bd5a92102467c4ace8a0d970203f/cryptography-48.0.1-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:b74ca3b8e5ecdd833bf6a002ca41b4793bb27fb8f1c06ffaf2643c9e9140e31b", size = 4731388, upload-time = "2026-06-09T22:32:04.019Z" }, - { url = "https://files.pythonhosted.org/packages/1b/24/50027ea4dca85ec1f40688f3c24fb32ccacd520583c9592c3cc95628e6fb/cryptography-48.0.1-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:2c37f2461406063b417837f5f3daab668652acd82423efcd7f0a9f04be972de1", size = 4824186, upload-time = "2026-06-09T22:32:18.707Z" }, - { url = "https://files.pythonhosted.org/packages/52/41/04cb5eb17085ade6f50cc611fb657df6a0f5885350de8764ece89c050197/cryptography-48.0.1-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:86fe77abb1bd87afb251d4d02ada7ecf53a32cee9b67d976abb2e45a13297475", size = 4964539, upload-time = "2026-06-09T22:31:18.793Z" }, - { url = "https://files.pythonhosted.org/packages/36/bf/ed70785c496e89d7e73b7cda2d21f2447fd6d4e821714b8d04ff217fed92/cryptography-48.0.1-cp39-abi3-win32.whl", hash = "sha256:6b2c0c3e6ccf3ade7750f836ef3ee36eea250cc467d45c256895573ac08cc6f1", size = 3282307, upload-time = "2026-06-09T22:30:53.162Z" }, - { url = "https://files.pythonhosted.org/packages/b3/ff/371ea7d252656ee1eb6d83eeeef3d1d0c6baf1d6497687d081ea03814670/cryptography-48.0.1-cp39-abi3-win_amd64.whl", hash = "sha256:9a49ca6c81417f6a5edb50375a60cccdd70fa0a91a5211829dbea74eba94d2ac", size = 3793408, upload-time = "2026-06-09T22:32:15.191Z" }, - { url = "https://files.pythonhosted.org/packages/a9/d3/eb4e394e587341fdad09a09101fa76478ead3a78b0ad63e55c22f0d75c02/cryptography-48.0.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:08a597acce1ff37f347400087776599e2348a3a8bc53b44120e463cd274efe4a", size = 3951747, upload-time = "2026-06-09T22:31:23.871Z" }, - { url = "https://files.pythonhosted.org/packages/e0/4a/3f43451b4f858bfceaaaffc649e6e787e8d4fb332a1d443af39ab02cc8f1/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:735824ec41b7f74a7c45fb1591349333e4c696cb6c044e5f46356e560143e4cd", size = 4641226, upload-time = "2026-06-09T22:31:02.532Z" }, - { url = "https://files.pythonhosted.org/packages/73/4e/855584c2c23b09e4ce2d3b9c30e983e679cd60b068c513c6bbdb91e11782/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:92a46e1d638daa264ba2971c0b0489c9409787943efae4d60ffda3d091ef832c", size = 4668958, upload-time = "2026-06-09T22:32:06.213Z" }, - { url = "https://files.pythonhosted.org/packages/42/3b/d35750e41d803d1e516fd6d6011f065424924da7af1748cef4cc9cb3ede1/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:7e234ac052af99f2700826a5c29ea99d9c1b1f80341cde62d11c8154dc8e0bd9", size = 4640793, upload-time = "2026-06-09T22:32:26.331Z" }, - { url = "https://files.pythonhosted.org/packages/ca/aa/cdb7181fe865285e87e96825aaab239400f1de0c3bfba9bd9769b79f1a92/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:33842cf0888951cef5bc7ac724ab844a42044c1727b967b7f8997289a0464f92", size = 4668505, upload-time = "2026-06-09T22:31:27.534Z" }, - { url = "https://files.pythonhosted.org/packages/5d/8c/ce3823c06c2804f194f9e64f0d67fa3f4094a39f2bb1a990cd03603af8fc/cryptography-48.0.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:6184ca7b174f28d7c703f1290d4b297217c45355f77a98f67e9b7f14549ac54a", size = 3742204, upload-time = "2026-06-09T22:31:34.773Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/de/41/6cbdcf9142d00fe82836fbb51e503e58088575cf7a0fe1dbff6695bf0840/cryptography-50.0.0.tar.gz", hash = "sha256:eeac2acb5a20ed25e0ad6d1df9891a520b78b404266b6d11778f25d5d691a6c9", size = 880201, upload-time = "2026-07-31T14:25:10.11Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/5c/59086b4aac5e879d38ddbcf74e4be7ade89cebc3eb199a55da998c3bb46a/cryptography-50.0.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:031e2d5dd4bb9caa3ca9c82e5a197fd8ae680232cee62603d1a813f3f07e3d03", size = 4001252, upload-time = "2026-07-31T14:23:33.331Z" }, + { url = "https://files.pythonhosted.org/packages/57/ef/8f2df13c7216bcad3e1c74e07f6e193d93e998e114f524a53877c9af27ad/cryptography-50.0.0-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fd9192b7b70c573d7f214eb1ae35e00d359f6f5e4b27c7e21e30de1fc6204645", size = 4719554, upload-time = "2026-07-31T14:23:35.611Z" }, + { url = "https://files.pythonhosted.org/packages/d9/41/029086c34d91052fc3b88bcc8056f709a7c915c7a23b235a54eb800b1c97/cryptography-50.0.0-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:06a32a980526a6ab9a4b9bf8f7385800791e2bb960903cb6b530e4817509a3b7", size = 4702130, upload-time = "2026-07-31T14:23:37.635Z" }, + { url = "https://files.pythonhosted.org/packages/7d/ff/b6ce0954962e7f7b969f850a883744197bb3910bdfd7b6da162eab7d9f68/cryptography-50.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a1b30560f2acc95aa8b2e06e716a13dbfc97314747b80d9707e307f77b40d6b3", size = 4725244, upload-time = "2026-07-31T14:23:39.471Z" }, + { url = "https://files.pythonhosted.org/packages/06/1e/63a1027cb7fec360a182208e1b7767d5aa1fe57be3d6aa856e69a321edc0/cryptography-50.0.0-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:8d89f3976b10b4ce31118de72329025f70d2c6ead14a8217c5514dd2c6d5a78f", size = 5342265, upload-time = "2026-07-31T14:23:41.286Z" }, + { url = "https://files.pythonhosted.org/packages/6b/72/a1116d683a6d7ece94590013882515de087edf9ef0e6292aae615a44df73/cryptography-50.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b42a28c1844fd9de8f3f7d540e36b66f3a9c83fceac7170ebc7a6a19edd9dcae", size = 4734609, upload-time = "2026-07-31T14:23:43.139Z" }, + { url = "https://files.pythonhosted.org/packages/15/37/36a9c479bbe49acea2636c7fd3360d20f7b7e079c300352011c44850b181/cryptography-50.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:900131fafd8aead39ac7dd3a7e833be754c17a95cfd91221636949fe4eb0aa8a", size = 4356517, upload-time = "2026-07-31T14:23:44.939Z" }, + { url = "https://files.pythonhosted.org/packages/32/98/8a151d64367204cbc63ec65d37502f1d9c53cf4bfc6ec3c532614dbec60d/cryptography-50.0.0-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:07949c449a1abcf60d1ee6e88956d89404c7df3c8258f46589e912988e551987", size = 4724529, upload-time = "2026-07-31T14:23:46.93Z" }, + { url = "https://files.pythonhosted.org/packages/22/f6/ec13b470172126464a86bf54d2294a46d29837fc51ba3e45d4047946fb5e/cryptography-50.0.0-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:f89831ef99dd7dd169ab06d63a831adb9e20a87aac6d380266bbda5823349169", size = 5299852, upload-time = "2026-07-31T14:23:48.851Z" }, + { url = "https://files.pythonhosted.org/packages/da/3a/f05e32c99d440c9bb891ea0e36c9091891e36be5a9a87ab2ee6ea20729f6/cryptography-50.0.0-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:82148ec5bddac30b51a5b3c1945075f896fa022cb93f8e4a01e9f6ee95292c5f", size = 4734462, upload-time = "2026-07-31T14:23:50.861Z" }, + { url = "https://files.pythonhosted.org/packages/ca/dc/bd72b26be8953f80625f63151efd38eee71c76ca6cf591c08ff34615a79e/cryptography-50.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1489e263a8048bb8b6a8bac662eb2d402ea5d2b7b4699b72f385f1e2772db105", size = 4852708, upload-time = "2026-07-31T14:23:52.715Z" }, + { url = "https://files.pythonhosted.org/packages/27/20/c930314a2ab476d15dec966ec87e2e9637bb02b06106b12c0396c57bb603/cryptography-50.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7cec5b856506da6defb290f30c9ee687d5f5e8cb0bd3f6459dde43b0b4fa40ef", size = 5004179, upload-time = "2026-07-31T14:23:54.887Z" }, + { url = "https://files.pythonhosted.org/packages/32/2e/c9db68a0c4bfa28e310707527c0ee3a2bd254104d2e02e68f368e197aa4c/cryptography-50.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:bd1c592e4d5974f0d08d4888e432157adba757c66da0246918e43677fafa2d30", size = 3840395, upload-time = "2026-07-31T14:23:56.677Z" }, + { url = "https://files.pythonhosted.org/packages/c3/fb/951032a3bf22a5697c83183fb6294a4843772947a70e616c57b3ff5f522e/cryptography-50.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:49e7d93abdbd2990caced757e5fade25302f719c3c8fb6e6fff2dde98999fc41", size = 3989258, upload-time = "2026-07-31T14:23:58.881Z" }, + { url = "https://files.pythonhosted.org/packages/d4/67/91eb047e69c5e845f2f14b8a2e4a1aab0f283cb885531e9e22c8adb176bc/cryptography-50.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:19736989797678c6af1e55cd49055cdbcb55d8f6b5583ac5335f933aba9101dc", size = 4700648, upload-time = "2026-07-31T14:24:00.702Z" }, + { url = "https://files.pythonhosted.org/packages/30/82/85f0f7425c856b9f96459411eb12e74ef72df9caf6f8f15bf23a33ff131f/cryptography-50.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:80b63928fa35083b33966ce1efb70e5b9607181e49dcd1c22c8c005e319f667f", size = 4682442, upload-time = "2026-07-31T14:24:02.538Z" }, + { url = "https://files.pythonhosted.org/packages/1a/28/b555a365adff1cca2fbe7b9e487d68a40de6bc67ff2cb587473eb43de0e7/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:d58c3db7cd6eed54e6c06744db55456b65ebd7492ddeae9c1e93cfca7aa857d3", size = 4707596, upload-time = "2026-07-31T14:24:04.394Z" }, + { url = "https://files.pythonhosted.org/packages/72/d8/f52538140cc719df62a01cf87d1c7142318d235817109d6f4054d7c352d6/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:df2a58a472f332225671c35b0a830208b86d004f82baa8530fa3782c85646533", size = 5314552, upload-time = "2026-07-31T14:24:06.31Z" }, + { url = "https://files.pythonhosted.org/packages/38/14/6120e5bd7c5aa022ad15424ba4d5c5269d0d9448ed4d55e492ea91e3c1c4/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:11b74db56cdbe3cdee6e3f6982ecb70334fa10dce99ed58bf7894aaaa3b2a037", size = 4717113, upload-time = "2026-07-31T14:24:08.349Z" }, + { url = "https://files.pythonhosted.org/packages/fa/71/190bf38c3ee2e0f8efc9860ae100c9df4169742eef274b91e7aa1cb133b9/cryptography-50.0.0-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:f59e38625469987d7ef6d495323c55e7db6c212eaf6112267e0d3b565a2e9c9f", size = 4338580, upload-time = "2026-07-31T14:24:10.227Z" }, + { url = "https://files.pythonhosted.org/packages/3a/63/504ccfbbe61fd8aa983f7f146399cdf034c72c2fc55f5b2dfdcdcdb20c99/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:ecfed7367f965a0328cfbdd70da860f15441f002f613185668c6e6ebf5a0ac11", size = 4707038, upload-time = "2026-07-31T14:24:12.169Z" }, + { url = "https://files.pythonhosted.org/packages/01/77/2cf79bbfc4d12ca106437a6e170d6aaa01a373e93093118aaaef0e801bd4/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:9aa87839c383bdbab6ef865787a1fb877af8dd03464c4400322726feaaadfc6d", size = 5273110, upload-time = "2026-07-31T14:24:14.38Z" }, + { url = "https://files.pythonhosted.org/packages/e5/45/8aae2972c520145377ea3559a605a899bebe227bf070b33cdb445929a9b9/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:6ba6a53445bd3cfa809ef3ef5f1589aa6ba08784a1d962bf47d0940e871dab1c", size = 4716439, upload-time = "2026-07-31T14:24:16.415Z" }, + { url = "https://files.pythonhosted.org/packages/7b/20/4fe50b619a48c2525cc46e2dbc1ac490708d704be5d467bdaac6dc955682/cryptography-50.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3f5735ffe4996d28b809371756219f5354864902a3b9e7c0b9ee87041209fc9c", size = 4837383, upload-time = "2026-07-31T14:24:18.553Z" }, + { url = "https://files.pythonhosted.org/packages/92/91/3a31366e183343d3703f8995c095f5734676bd6938118047e50fcf279eb4/cryptography-50.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:1b4a266766514614f8aa60416e71f2fc6e575d36e7bdc90f644fadb2f4b75b95", size = 4985772, upload-time = "2026-07-31T14:24:20.385Z" }, + { url = "https://files.pythonhosted.org/packages/74/9a/02ffe35b2853d121689871eb5dce862092562b3a1ed5cc98f1aaed441506/cryptography-50.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:12b9c6996425c76ea6c457ace4f3073e715b8c545add07cd1a8f3a4f90691269", size = 3816291, upload-time = "2026-07-31T14:24:22.125Z" }, + { url = "https://files.pythonhosted.org/packages/03/37/73d005be173aff344af30e9fd2a576575cb2391a7101d9cd3842e1fa8cce/cryptography-50.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:ccdc4a71a4dabae05de219404f9f4abc38e3b58422177ff93d0da05967dafa07", size = 4036009, upload-time = "2026-07-31T14:24:24.122Z" }, + { url = "https://files.pythonhosted.org/packages/ff/c6/7a6202a534e32103a285b7834a120869557fe198d51d7cfe59754c8bda9c/cryptography-50.0.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:910e1d2668e7de9648f2bcee30e180db2a6b15c30f887d7c4c93ddf96e3992e3", size = 4745252, upload-time = "2026-07-31T14:24:26.118Z" }, + { url = "https://files.pythonhosted.org/packages/85/4f/0fa8c2f4428198f15d9ff8d63400e27afbf94ce833f6108da1eb3753f945/cryptography-50.0.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a91296cb61e8df6f86d0c19cc4068228da256bf59bf86049fbd821084565327f", size = 4728939, upload-time = "2026-07-31T14:24:27.994Z" }, + { url = "https://files.pythonhosted.org/packages/d1/63/54dd723490ba2dc09b299682c10b38db38f159728bcaae8c591b8af2f22d/cryptography-50.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:e722f16708d854fe924790e051061f6704a472c3bac347b6fd88033ea8dd0dc5", size = 4748483, upload-time = "2026-07-31T14:24:30.254Z" }, + { url = "https://files.pythonhosted.org/packages/1d/dd/7c77d26285cc7f6991efce64a0f5b4f9383bfa5dd8c5033003eaf7db4cdb/cryptography-50.0.0-cp39-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:d764dcf130c428ef66786f866dd750f53182bc608813489915e9fc106bb0c82f", size = 5367599, upload-time = "2026-07-31T14:24:32.457Z" }, + { url = "https://files.pythonhosted.org/packages/46/c9/f60aed34c013f317f92817b6c171c2d22a78270fa41109bd4b08af26b194/cryptography-50.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:105110f43a471dbd0060b9c9516cb8a6a79233631a04cc2ba16f28323ac6e025", size = 4762647, upload-time = "2026-07-31T14:24:34.599Z" }, + { url = "https://files.pythonhosted.org/packages/be/f3/f9a0173b139372c3a48ed98154b45cc6b9de17c789d5ab552e621c293609/cryptography-50.0.0-cp39-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:828743d939e9629bc267b8e2d08d8bb67cd4319c771a33d4b18b22dd8fb7440a", size = 4385197, upload-time = "2026-07-31T14:24:36.647Z" }, + { url = "https://files.pythonhosted.org/packages/d8/36/83bb81f6e569bc38e1e4a7bc80f29b46bb9601920bc455fc8e888f5d5742/cryptography-50.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:2a8183b489dc1f7f80f135780fadc1108f14b31b8a40411c7a5b17425f65f28b", size = 4748095, upload-time = "2026-07-31T14:24:39.493Z" }, + { url = "https://files.pythonhosted.org/packages/6b/16/d3008eff98c764979865834c3d386d4fd041b5f52e7f34fc29ac1a5eb515/cryptography-50.0.0-cp39-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:6e7d61120573a7f2cd94cc095f9e81f6967c61ccdf194285aa143ecec8e0b708", size = 5325948, upload-time = "2026-07-31T14:24:41.556Z" }, + { url = "https://files.pythonhosted.org/packages/9c/f8/d97f9603efda3888187bfdb893f26c41be4735c10631d05d284ee6b047c4/cryptography-50.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:37fdb0d0111f1e2ff07139dfb79f1b49531f8e213c46f1163dd7642979b58c47", size = 4762400, upload-time = "2026-07-31T14:24:43.636Z" }, + { url = "https://files.pythonhosted.org/packages/64/a2/4615c8f7d81a00b1d6e6afe19f694e1543582349fb5f4076f6cb5dc36485/cryptography-50.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c87f62a3d3b9888ed0fdde100ec06aa61ca9cd44bad9057d1dff9a516b5f5bb9", size = 4878208, upload-time = "2026-07-31T14:24:45.522Z" }, + { url = "https://files.pythonhosted.org/packages/d2/1a/efcfb02f91407149a0dacffffab791f7e19bf6385f63b3666dc8b5e5c9c8/cryptography-50.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:65c2c3add92b45fd0709db8594536aea39c2a67af0e27ffcf049c498501140b7", size = 5037050, upload-time = "2026-07-31T14:24:47.697Z" }, + { url = "https://files.pythonhosted.org/packages/57/30/4a22984d4f1bdfb8c054f07a92bc176b97a3134cc1d6c4b3bffb1f3688b4/cryptography-50.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:d24fead1d4d076e1bfb006dcec392074a3cd8d7b4fc8a595aa64073b2b7a96ba", size = 3874135, upload-time = "2026-07-31T14:24:50.085Z" }, + { url = "https://files.pythonhosted.org/packages/9d/3e/e54cde8c01631a5a8226ccd617eab9e57fd5cfdad90f1a9e6bb570794631/cryptography-50.0.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:5e34edd123674534acd70147f0ca331eaa2c74e6325fb2028c886aa26ba0b68c", size = 3963170, upload-time = "2026-07-31T14:24:51.968Z" }, + { url = "https://files.pythonhosted.org/packages/01/b6/0b9e125e90f3d2dcf599a218a899cda7326a3158cfa258723f0b398b08f6/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:8eb5e1172eb569ea8a872796576e6a67c276351728b6455d5beb01242b027c6a", size = 4692441, upload-time = "2026-07-31T14:24:53.743Z" }, + { url = "https://files.pythonhosted.org/packages/53/c9/a5151588710785a96d7bc4de27d4cd62f263bbbcb203cfe29df537eb6505/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:910d11e1a385c654bf738bf3e6b8e6ed5de0f5610fcae2be9e5b398d8081d20e", size = 4699810, upload-time = "2026-07-31T14:24:55.746Z" }, + { url = "https://files.pythonhosted.org/packages/c7/1a/15b92b25eb6ce3089cd49377ae990a0f3ad485a510f968aed1f19dbdcdf2/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:62598a8a57f815db4c6259a4e97d857dab56697e7de8e8ab02352ab74da1995d", size = 4691924, upload-time = "2026-07-31T14:24:58.082Z" }, + { url = "https://files.pythonhosted.org/packages/62/15/219075012ab13e8905f3cd572204f4acb4b111df787104346b9bc0cea789/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:07479a1cb08219ab719147e742e76090c9c773321959bb94946fffdd397a6437", size = 4699593, upload-time = "2026-07-31T14:24:59.951Z" }, + { url = "https://files.pythonhosted.org/packages/8e/b5/c2c5fce26f0ee40d21bafe7f191d29a34b35a65ac4fe8a1191d1983612e9/cryptography-50.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c99c003e088647b8a5b7c145d6f78c335f6348332b62e142d411c4b63d1460b9", size = 3813796, upload-time = "2026-07-31T14:25:02.298Z" }, ] [[package]] @@ -226,7 +223,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ diff --git a/olive/assets/io_configs/defaults.yaml b/olive/assets/io_configs/defaults.yaml index a6c055a464..d92ece4c92 100644 --- a/olive/assets/io_configs/defaults.yaml +++ b/olive/assets/io_configs/defaults.yaml @@ -3,12 +3,34 @@ # Attribute aliases (same concept, different naming across models) aliases: - # Layer count - num_layers: [num_hidden_layers, n_layer, n_layers] - # Hidden dimensions - hidden_size: [dim, d_model, n_embd] - num_attention_heads: [num_heads, n_head, n_heads, encoder_attention_heads] - num_kv_heads: [num_key_value_heads] + # Layer count. ``text_config.num_hidden_layers`` covers composite VL configs (e.g. + # Qwen3.5/3.6-MoE-VL, GLM4V-MoE) whose top level only has ``vision_config``/``text_config`` + # and no flat decoder attributes of its own. + num_layers: [num_hidden_layers, n_layer, n_layers, text_config.num_hidden_layers] + # Hidden dimensions. See ``num_layers`` above for why the ``text_config.*`` fallbacks exist. + hidden_size: [dim, d_model, n_embd, text_config.hidden_size] + num_attention_heads: [num_heads, n_head, n_heads, encoder_attention_heads, text_config.num_attention_heads] + num_kv_heads: [num_key_value_heads, text_config.num_key_value_heads] + # Explicit head_dim, so composite VL configs whose head_dim doesn't equal + # hidden_size // num_attention_heads (e.g. Qwen3.5/3.6-MoE-VL) resolve correctly instead of + # falling back to that derived (and here incorrect) value. + head_dim: [text_config.head_dim] + # MoE expert count. Covers flat configs (Mixtral / gpt-oss / PhiMoE + # ``num_local_experts``, DeepSeek ``n_routed_experts``, Ernie4.5 / Aria + # ``moe_num_experts``) and nested sub-configs (DBRX ``ffn_config``, Llama4 / + # GLM4V-MoE / Qwen3-VL-MoE ``text_config``, Qwen3-Omni ``thinker_config.text_config``). + num_experts: + - num_local_experts + - n_routed_experts + - moe_num_experts + - ffn_config.moe_num_experts + - ffn_config.num_experts + - text_config.num_experts + - text_config.num_local_experts + - text_config.n_routed_experts + - text_config.moe_num_experts + - thinker_config.text_config.num_experts + - thinker_config.text_config.num_local_experts # Image dimensions height: [sample_size, image_size, vision_config.image_size] width: [sample_size, image_size, vision_config.image_size] diff --git a/olive/cache.py b/olive/cache.py index 7f27a50de1..654a44766d 100644 --- a/olive/cache.py +++ b/olive/cache.py @@ -18,7 +18,7 @@ from olive.common.config_utils import ConfigBase, convert_configs_to_dicts, validate_config from olive.common.constants import DEFAULT_CACHE_DIR, DEFAULT_WORKFLOW_ID from olive.common.container_client_factory import AzureContainerClientFactory -from olive.common.utils import hash_dict, hf_repo_exists, set_nested_dict_value +from olive.common.utils import hardlink_copy_file, hash_dict, hf_repo_exists, set_nested_dict_value from olive.model.config.model_config import ModelConfig from olive.resource_path import ResourcePath, create_resource_path, find_all_resources @@ -451,10 +451,68 @@ def _rebase_additional_files(config: dict, fallback_dir: Path): else: copied_components = [] saved_external_files = {} + + # replace resources with their local cache paths once per component so the same + # resolved model json/paths can be used both to detect shared external-data files + # (below) and to actually copy the components (further below) + processed_components = [] for component_name, component in zip( model_json_config["model_component_names"], model_json_config["model_components"] ): if component["type"].lower() != "onnxmodel": + processed_components.append((component_name, component, None, None)) + else: + component_model_json, component_local_resource_names = self._replace_with_local_resources( + component, only_cache_files=only_cache_files + ) + processed_components.append( + (component_name, component, component_model_json, component_local_resource_names) + ) + + # onnx components of a composite model can share a single external-data file (e.g. + # QNN GPU static LLM prefill/decode models reusing one model.onnx.data). Detect + # that sharing upfront and reserve+copy the original file name for it, so it keeps + # that name in the output folder instead of being renamed after whichever component + # happens to be processed first. Opt-in via model_attributes so this scan only runs + # for passes that are known to rely on it (currently QNN GPU StaticLLM). + if model_attributes.get("keep_shared_external_data_names"): + from olive.passes.onnx.common import get_external_data_file_names + + external_file_usage: dict[str, int] = {} + for _, _, component_model_json, _ in processed_components: + if component_model_json is None: + continue + component_model_path = ( + ModelConfig.model_validate(component_model_json).create_model().model_path + ) + if not component_model_path or not Path(component_model_path).is_file(): + continue + for external_name in get_external_data_file_names(component_model_path): + external_file_path = str(Path(component_model_path).parent / external_name) + external_file_usage[external_file_path] = external_file_usage.get(external_file_path, 0) + 1 + + reserved_external_names = set() + for external_file_path, usage_count in external_file_usage.items(): + if usage_count <= 1: + continue + original_name = Path(external_file_path).name + if original_name in reserved_external_names: + continue + # resave_model only copies an external-data file the first time it sees its + # source path; since we're deciding the target name ahead of that, copy it here + # so the file actually exists once resave_model reuses this reserved name. + actual_output_dir.mkdir(parents=True, exist_ok=True) + hardlink_copy_file(external_file_path, actual_output_dir / original_name, follow_symlinks=True) + saved_external_files[external_file_path] = original_name + reserved_external_names.add(original_name) + + for ( + component_name, + component, + component_model_json, + component_local_resource_names, + ) in processed_components: + if component_model_json is None: # save each component with a prefix # e.g. "component_1" -> "component_1_{resource_name}" copied_components.append( @@ -468,10 +526,6 @@ def _rebase_additional_files(config: dict, fallback_dir: Path): ) else: # save all onnx files into the same directory - component_model_json, component_local_resource_names = self._replace_with_local_resources( - component, only_cache_files=only_cache_files - ) - for resource_name in component_local_resource_names: if resource_name != "model_path": # this case does not exist in the current code @@ -541,9 +595,11 @@ def _save_model( output_file = output_dir actual_output_dir = output_dir.parent else: - # Otherwise, create model.onnx in the directory + # Otherwise, create model.onnx in the directory. + # Preserve the source onnx_file_name stem (e.g. model_ctx) so the output + # filename matches what genai_config.json references. actual_output_dir = output_dir - model_file_name = "model" + model_file_name = Path(onnx_file_name).stem if has_additional_files and onnx_file_name else "model" if path_prefix: model_file_name = f"{path_prefix}_{model_file_name}" output_file = output_dir / f"{model_file_name}.onnx" diff --git a/olive/cli/api.py b/olive/cli/api.py index cae2963264..86bbe36f5e 100644 --- a/olive/cli/api.py +++ b/olive/cli/api.py @@ -19,6 +19,7 @@ from olive.cli.run import WorkflowRunCommand from olive.cli.session_params_tuning import SessionParamsTuningCommand from olive.engine.output import WorkflowOutput +from olive.telemetry import disable_telemetry # pylint: disable=W0212 @@ -127,6 +128,8 @@ def _run_unified_command(command_class, **kwargs) -> Any: known_kwargs = {k: v for k, v in kwargs.items() if k in args_schema} args = _create_unified_args(command_class, known_kwargs) + if getattr(args, "disable_telemetry", False): + disable_telemetry() # Check if command handles unknown_args (like FineTuneCommand) constructor_params = inspect.signature(command_class.__init__).parameters diff --git a/olive/cli/base.py b/olive/cli/base.py index a64a9bf473..ef19fc8968 100644 --- a/olive/cli/base.py +++ b/olive/cli/base.py @@ -31,7 +31,7 @@ def _parse_test_metrics(value: str) -> list: Accepts values like ``'mae'``, ``'mae,speedup'``, or ``'mae speedup'`` and returns a flat list of validated metric names. Raises ``argparse.ArgumentTypeError`` - for any unrecognised name. + for any unrecognized name. """ import argparse @@ -338,7 +338,9 @@ def _run_workflow(self): mark_test_output_path(self.args.output_path) print("Dry run mode enabled. Configuration file is generated but no optimization is performed.") return None - workflow_output = olive_run(run_config) + workflow_output = olive_run( + run_config, recipe_telemetry_metadata=self._get_recipe_telemetry_metadata(), emit_error_telemetry=False + ) if is_test: mark_test_output_path(self.args.output_path) save_discrepancy_check_results(workflow_output, self.args.output_path) @@ -349,6 +351,17 @@ def _run_workflow(self): print(f"Model is saved at {self.args.output_path}") return workflow_output + def _get_recipe_telemetry_metadata(self) -> dict[str, str]: + recipe_name = self.__class__.__name__ + if recipe_name.endswith("Command"): + recipe_name = recipe_name[: -len("Command")] + return { + "recipe_name": recipe_name, + "recipe_command": recipe_name, + "recipe_source": "generated_cli", + "recipe_format": "generated", + } + @staticmethod def _parse_extra_options(kv_items): from onnxruntime_genai import __version__ as OrtGenaiVersion @@ -360,9 +373,17 @@ def _parse_extra_options(kv_items): "Please either upgrade to onnxruntime-genai version > 0.9.0 or use the model builder pass directly in the config file." ) - from onnxruntime_genai.models.builder import parse_extra_options - - return parse_extra_options(kv_items) # pylint: disable=no-value-for-parameter + # The model builder also validates the options and loads the Hugging Face config, so its + # `parse_extra_options` needs the model, precision and execution provider that are only + # known once the workflow runs. Just split the pairs here and let the ModelBuilder pass + # run the validation later. + extra_options = {} + for item in kv_items: + key, separator, value = item.partition("=") + if not separator: + raise ValueError(f"Invalid extra option '{item}'. Expected key=value.") + extra_options[key.strip()] = value.strip() + return extra_options @staticmethod def _save_config_file(config: dict, output_dir: Optional[str] = None, file_name: str = "config.json"): @@ -995,7 +1016,11 @@ def add_search_options(sub_parser: ArgumentParser): def add_telemetry_options(sub_parser: ArgumentParser): """Add telemetry options to the sub_parser.""" - sub_parser.add_argument("--disable_telemetry", action="store_true", help="Disable telemetry for this command.") + sub_parser.add_argument( + "--disable_telemetry", + action="store_true", + help="Disable all Olive telemetry for this process.", + ) return sub_parser diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 262a29ab08..3966b798c3 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -82,9 +82,9 @@ def register_subcommand(parser: ArgumentParser): "--use_mobius_builder", action="store_true", help=( - "Whether to use MobiusBuilder (mobius-ai) to capture ONNX model. " + "Whether to use MobiusBuilder (mobius-onnx) to capture ONNX model. " "Supports multi-component multimodal models (VLMs). " - "Requires 'pip install mobius-ai'." + "Requires 'pip install mobius-onnx'." ), ) diff --git a/olive/cli/init/__init__.py b/olive/cli/init/__init__.py index 26ae8d3562..213de7aa0b 100644 --- a/olive/cli/init/__init__.py +++ b/olive/cli/init/__init__.py @@ -4,7 +4,8 @@ # -------------------------------------------------------------------------- from argparse import ArgumentParser -from olive.cli.base import BaseOliveCLICommand +from olive.cli.base import BaseOliveCLICommand, add_telemetry_options +from olive.telemetry import action class InitCommand(BaseOliveCLICommand): @@ -21,8 +22,10 @@ def register_subcommand(parser: ArgumentParser): default="./olive-output", help="Default output directory for the generated command. Default is ./olive-output.", ) + add_telemetry_options(sub_parser) sub_parser.set_defaults(func=InitCommand) + @action def run(self): from olive.cli.init.wizard import InitWizard diff --git a/olive/cli/launcher.py b/olive/cli/launcher.py index 55e6ffdeb4..369c93092a 100644 --- a/olive/cli/launcher.py +++ b/olive/cli/launcher.py @@ -24,7 +24,7 @@ from olive.cli.run_pass import RunPassCommand from olive.cli.session_params_tuning import SessionParamsTuningCommand from olive.cli.shared_cache import SharedCacheCommand -from olive.telemetry import Telemetry +from olive.telemetry import Telemetry, disable_telemetry def get_cli_parser(called_as_console_script: bool = True) -> ArgumentParser: @@ -66,18 +66,20 @@ def main(raw_args=None, called_as_console_script: bool = True): args, unknown_args = parser.parse_known_args(raw_args) - telemetry = Telemetry() - if args.disable_telemetry: - telemetry.disable_telemetry() - if not hasattr(args, "func"): parser.print_help() sys.exit(1) - # Run the command - service = args.func(parser, args, unknown_args) - service.run() - telemetry.shutdown() + if getattr(args, "disable_telemetry", False): + disable_telemetry() + telemetry = None + try: + telemetry = Telemetry.get_or_create_if_enabled() + service = args.func(parser, args, unknown_args) + service.run() + finally: + if telemetry is not None: + telemetry.shutdown() def legacy_call(deprecated_module: str, command_name: str, *args): diff --git a/olive/cli/model_package.py b/olive/cli/model_package.py index 6cc5f90f22..79b0ca05fb 100644 --- a/olive/cli/model_package.py +++ b/olive/cli/model_package.py @@ -11,32 +11,48 @@ ``CompositeModel`` with ONNX components). Single-source packages are allowed: a single variant under one component is a normal, valid package. -Output layout (per the ORT model-package proposal):: +Output layout (per the ONNX Runtime model-package specification):: / ├── manifest.json - ├── configs/ - │ └── # tokenizer, genai_config, ... - └── models/ - └── / - ├── metadata.json - └── / - ├── genai_config_overlay.json # optional: per-variant runtime fields - ├── model.onnx - └── ... # external-data blobs (inline) + ├── models/ + │ └── model/ # the one component + │ ├── component.json + │ └── / + │ ├── genai_config.json # complete, self-contained + │ ├── text.onnx # every role's graph + │ ├── vision.onnx + │ └── ... # external-data blobs (inline) + └── shared_assets/ + └── sha256-/ # content-addressed tokenizer assets + ├── tokenizer.json + └── ... Notes: -- ``metadata.json`` is selection-only. Each variant declares a single - execution provider inline (``ep``) plus optional ``device`` and opaque - ``compatibility_string``. -- Each variant directory is self-contained: the ONNX file and any external-data - blobs it references are copied inline so stock ORT can load it directly. -- ``genai_config.json`` is canonicalized into ``/configs/``: variant- - specific runtime fields (``filename``, ``session_options``) are stripped from - the base and each role gets a ``component`` pointer so ORT GenAI can map - roles to ``models//`` at load time. The stripped fields are - re-injected per variant as a ``genai_config_overlay.json`` (an RFC 7386 JSON - Merge Patch applied on top of ``configs/genai_config.json``). +- The package declares exactly **one** component. ORT-GenAI selects a single + component and loads the complete ``genai_config.json`` from the variant + directory it picks, resolving every role's ``filename`` against that one + directory; it rejects a package declaring more than one component. So all + of a model's roles (``decoder``, ``embedding``, ``vision``, ...) live in + one variant directory even though each is a separate ORT session at + runtime. Variants are the per-hardware builds of that same model. +- ``manifest.json`` maps component name -> component directory. ORT reads + ``component.json`` from that directory; the file is selection-only. Each + variant declares a single execution provider inline (``ep``) plus optional + ``device`` and opaque ``compatibility_string``. +- Each variant directory is self-contained: the ONNX files, any external-data + blobs they reference, and a **complete** ``genai_config.json`` are placed + inline so stock ORT and ORT-GenAI can load the variant directly. ORT-GenAI + loads ``/genai_config.json`` and has no notion of a + package-level base config or a merge-patch overlay. +- Role graphs keep the relative path their source declared + (``vision_encoder/model.onnx``), so producers that name every role's graph + ``model.onnx`` don't collide inside the shared variant directory. +- Tokenizer assets are shared across variants through a content-addressed + ``shared_assets/sha256-/`` directory; each variant config points at it + with ``model.tokenizer_dir = "sha256:"``. Processor configs stay inside + the variant directory because ORT-GenAI resolves those relative to + ``genai_config.json`` rather than through the package resolver. """ @@ -60,23 +76,46 @@ logger = logging.getLogger(__name__) # Files inside an Olive output dir that always belong next to the ONNX model -# rather than under /configs/. +# rather than in a shared asset. _MODEL_SUFFIXES = {".onnx", ".bin", ".data", ".xml"} -# Schema versions emitted in the package JSON files. Keep in sync with the -# ORT model-package schema. -_MANIFEST_SCHEMA_VERSION = 1 -_METADATA_SCHEMA_VERSION = 1 +# Schema version emitted in manifest.json. The ORT model-package schema +# expects a "." string; the major gates compatibility. +_MANIFEST_SCHEMA_VERSION = "1.0" -# Directory under the package root that holds consumer-shared config assets -# (genai_config base, tokenizer, processor configs, chat templates). -_CONFIGS_DIR = "configs" +# Directory under the package root that holds content-addressed shared assets. +# ORT discovers ``/shared_assets/sha256-/`` at open time; variants +# reference them with the ``sha256:`` scheme. +_SHARED_ASSETS_DIR = "shared_assets" + +# Filename ORT reads when a manifest component entry points at a directory. +# The name is fixed by the ORT model-package specification. +_COMPONENT_FILENAME = "component.json" + +# Config files ORT-GenAI resolves relative to ``genai_config.json`` (via +# ``config_path / filename``) rather than through the package resolver. These +# must be copied into every variant directory instead of a shared asset. +_VARIANT_LOCAL_CONFIG_NAMES = frozenset({"processor_config.json", "audio_processor_config.json"}) + +# The one file onnxruntime-extensions requires to exist under +# ``model.tokenizer_dir``; its absence is a hard tokenizer-load failure. Used to +# decide whether the shared asset directory is worth pointing ``tokenizer_dir`` +# at in the first place. +_TOKENIZER_SENTINEL = "tokenizer_config.json" # Directory under the package root that holds per-component subdirectories. -# Required by the ORT model-package schema; ORT's model-package loader -# discovers components via ``/models//metadata.json``. +# The manifest maps each component name to ``models/``; ORT reads +# ``component.json`` from that directory. _MODELS_DIR = "models" +# Name of the single component every generated package declares. ORT-GenAI +# opens a package by selecting "the" component and requires the package to +# declare exactly one ("declares N components; onnxruntime-genai requires +# exactly one"), so every genai_config role must live in one component whose +# variant directory holds a complete genai_config.json. ``model`` is the name +# used by the ORT-GenAI model-package documentation. +_GENAI_COMPONENT_NAME = "model" + # Conventional directory suffix for an ORT model package. Not enforced by # ORT/ORT-GenAI loaders (they probe structure, not filenames), but matches # the canonical naming used in ORT's model-package documentation and the @@ -233,38 +272,50 @@ def run(self): def _build_variants(self, targets: list[tuple[str, Path, dict]]) -> list["VariantSpec"]: variants: list[VariantSpec] = [] for target_name, source_path, source_genai in targets: - # Each role under ``genai_config.model`` is an independent ORT - # inference session at runtime, so each role becomes its own - # package component. A text-only model has one role (``decoder``) - # → one component. A VLM has three roles (``vision``, - # ``embedding``, ``decoder``) → three components. A QNN - # pipeline-shaped role becomes ONE component whose variant - # directory holds every stage's ONNX flat. + # Every role under ``genai_config.model`` (``vision``, + # ``embedding``, ``decoder``, ...) is a separate ORT inference + # session at runtime, but they all belong to ONE package + # component: ORT-GenAI selects a single component and loads the + # complete ``genai_config.json`` from that component's variant + # directory, resolving every role's ``filename`` against it. + # Splitting roles across components would leave each component's + # config missing its siblings' graphs, and ORT-GenAI rejects a + # multi-component package outright. So one source = one variant + # holding all of its roles' ONNX files. artifacts_by_role = _collect_artifacts_per_role(source_path, source_genai) + + onnx_files: list[Path] = [] + onnx_rel_paths: list[str] = [] + onnx_rel_paths_by_role: dict[str, list[str]] = {} for role_name, role_artifacts in artifacts_by_role.items(): - onnx_files = [a.source_path for a in role_artifacts] - onnx_rel_paths = [a.package_rel_path for a in role_artifacts] - - ep = _resolve_ep_for_role(source_genai, role_name) - # ``ep_compatibility_info`` metadata is conventionally - # written on the role's first ONNX (the primary stage for - # a pipeline role); probe that file for the EP-scoped - # compatibility string. - raw_compat = _extract_ep_compatibility_from_onnx(onnx_files[0], ep) if onnx_files else None - compatibility_string = raw_compat.strip() if raw_compat and raw_compat.strip() else None - - variants.append( - VariantSpec( - component_name=role_name, - variant_name=target_name, - role_name=role_name, - onnx_files=onnx_files, - onnx_rel_paths=onnx_rel_paths, - ep=ep, - compatibility_string=compatibility_string, - source_genai=source_genai, - ) + onnx_rel_paths_by_role[role_name] = [a.package_rel_path for a in role_artifacts] + onnx_files.extend(a.source_path for a in role_artifacts) + onnx_rel_paths.extend(a.package_rel_path for a in role_artifacts) + + ep = _resolve_ep_for_variant(source_path, source_genai, artifacts_by_role) + # ``ep_compatibility_info`` metadata is conventionally written on + # the ONNX compiled for the EP. With every role in one variant the + # producer may tag only the role that actually targets the EP, so + # probe each graph and take the first declaration. + compatibility_string = None + for onnx_file in onnx_files: + raw_compat = _extract_ep_compatibility_from_onnx(onnx_file, ep) + if raw_compat and raw_compat.strip(): + compatibility_string = raw_compat.strip() + break + + variants.append( + VariantSpec( + component_name=_GENAI_COMPONENT_NAME, + variant_name=target_name, + onnx_files=onnx_files, + onnx_rel_paths=onnx_rel_paths, + onnx_rel_paths_by_role=onnx_rel_paths_by_role, + ep=ep, + compatibility_string=compatibility_string, + source_genai=source_genai, ) + ) return variants # ------------------------------------------------------------------ @@ -381,7 +432,7 @@ class OnnxArtifact: to disambiguate against. Direct callers that construct a VariantSpec by hand may supply a nested subpath when the variant truly needs a multi-file layout under one component. The rel path is the same - string the variant's ``genai_config_overlay.json`` emits under + string the variant's ``genai_config.json`` emits under ``model..filename``, so the on-disk layout and the loader's view stay aligned. """ @@ -411,16 +462,25 @@ class VariantSpec: # supplied, must be index-aligned with ``onnx_files`` and contain a # safe relative path for every ONNX. When empty, the writer falls back # to the legacy flat layout (each ONNX placed at - # ``/``) — kept for direct callers and tests that - # predate multi-component sources. + # ``/``) — kept for direct callers that supply + # only ``onnx_files``. onnx_rel_paths: list[str] = field(default_factory=list) + # Per-role view of ``onnx_rel_paths``, keyed by genai_config role + # (``decoder``, ``vision``, ...). The overlay writer needs to know which + # in-package paths belong to which role so each role's ``filename`` / + # pipeline stage filenames point at the files actually written for it. + # Empty for direct callers that supply only ``onnx_files``. + onnx_rel_paths_by_role: dict[str, list[str]] = field(default_factory=dict) # The genai_config role this variant represents (e.g. ``decoder``, # ``vision``, ``embedding``). When set, the overlay writer scopes its # lift to just this role rather than the whole ``model`` block, so a - # multi-role source (Mobius VLM) produces one VariantSpec per role - # under one component per role. When unset, the writer falls back to - # the legacy multi-role lift for direct callers / tests that predate - # per-role components. + # caller can split a multi-role model across one component per role. + # The ``generate-model-package`` CLI never sets this: ORT-GenAI opens + # exactly one component per package and resolves every role's + # ``filename`` against that component's selected variant directory, so + # a per-role split would be unloadable. It stays available for direct + # ``write_model_package`` callers targeting non-GenAI consumers, which + # the plain ORT model-package spec does allow. role_name: Optional[str] = None def __post_init__(self) -> None: @@ -469,13 +529,16 @@ def write_model_package( previous run. :param variants: Ordered list of variants. Component insertion order is the order each component first appears in this list. - :param config_files: Map from filename (basename) to source path; copied - into ``/configs/``. Same-named files contributed by - different sources should be byte-identical; the first wins on - conflict and a warning is logged. + :param config_files: Map from filename (basename) to source path. Tokenizer + assets land in a content-addressed ``shared_assets/sha256-/`` + directory; processor configs are copied into every variant directory; + ``genai_config.json`` becomes the base each variant's complete config + is derived from. Same-named files contributed by different sources + should be byte-identical; the first wins on conflict and a warning is + logged. :param producer_info: Olive-specific provenance recorded under - ``manifest.producer``. Schema-tolerated extra field; producers may add - namespaced extras. + ``manifest.additional_metadata.producer``. The ORT schema rejects + unknown top-level manifest keys, so provenance is namespaced there. :param package_name: Name recorded under ``manifest.package_name``. Defaults to the output directory name. :param package_version: Version recorded under ``manifest.package_version``. @@ -512,14 +575,10 @@ def write_model_package( ) seen.add(v.variant_name) - for comp_name, comp_variants in components.items(): - _write_component(output_dir, comp_name, comp_variants, component_to_role.get(comp_name, comp_name)) - - # Build the role -> component map needed by _copy_config_files so it can - # inject ``model..component`` markers into the base genai_config. ORT - # requires every role-block to declare which package component it loads; - # without those markers ORT-GenAI's variant auto-selection fails with - # "the genai config does not reference any package components". + # Role -> component mapping. Kept as a consistency check: a genai_config + # role must belong to exactly one package component, otherwise the + # per-variant configs derived below would disagree about which variant + # directory serves that role. role_to_component: dict[str, str] = {} def _assign_role(role: str, component: str) -> None: @@ -533,18 +592,18 @@ def _assign_role(role: str, component: str) -> None: "belong to exactly one package component." ) - # Preferred path: per-role variants (each VariantSpec carries the role - # it represents). One variant = one role = one component, so the - # mapping is direct and conflicts are easy to surface. + # Per-role variants (each VariantSpec carries the role it represents). + # One variant = one role = one component, so the mapping is direct and + # conflicts are easy to surface. Only direct ``write_model_package`` + # callers take this path; the CLI packs every role into one component. for v in variants: if v.role_name: _assign_role(v.role_name, v.component_name) - # Legacy / multi-role path: variants without ``role_name`` aggregate - # several roles under one component. Seed from each variant's source - # genai_config so every role that appears under ``model.`` - # (vision, embedding, decoder, ...) still points back at that - # variant's component_name. Preserves backward compatibility for - # direct ``write_model_package`` callers and tests. + # Multi-role path (the CLI default): variants without ``role_name`` + # aggregate several roles under one component. Seed from each variant's + # source genai_config so every role that appears under ``model.`` + # (vision, embedding, decoder, ...) points back at that variant's + # component_name. for v in variants: if v.role_name: continue @@ -562,19 +621,56 @@ def _assign_role(role: str, component: str) -> None: if explicit_role not in role_to_component: role_to_component[explicit_role] = comp_name - if config_files: - _copy_config_files(output_dir, config_files, role_to_component) + # Shared assets and the base genai_config must exist before the variants + # are written: each variant embeds a complete genai_config.json derived + # from the base, pointing at the tokenizer asset by URI. + shared = _write_shared_assets(output_dir, config_files or {}) + + for comp_name, comp_variants in components.items(): + _write_component( + output_dir, + comp_name, + comp_variants, + component_to_role.get(comp_name, comp_name), + shared, + ) _write_manifest( output_dir, list(components.keys()), producer_info, package_name or output_dir.name, package_version ) +@dataclass +class SharedConfigAssets: + """Package-level config material resolved before variants are written. + + :param base_genai: The package's base ``genai_config.json`` (parsed) with + variant-specific keys stripped. Each variant's complete config is this + merged with that variant's own fields. Empty when no source supplied a + ``genai_config.json``. + :param model_level_defaults: The model-level scalars in + ``_VARIANT_LEVEL_MODEL_KEYS`` as declared by the base source. Applied + under a variant's own values so a variant that does not declare them + still ends up with a complete config. Per-role keys are deliberately + excluded: a variant must never inherit another role's ``filename``. + :param tokenizer_asset_uri: ``sha256:`` URI of the shared tokenizer + asset directory, or ``None`` when no tokenizer files were contributed. + :param variant_local_files: Config files ORT-GenAI resolves relative to + ``genai_config.json``; copied into every variant directory. + """ + + base_genai: dict[str, Any] = field(default_factory=dict) + model_level_defaults: dict[str, Any] = field(default_factory=dict) + tokenizer_asset_uri: Optional[str] = None + variant_local_files: dict[str, Path] = field(default_factory=dict) + + def _write_component( output_dir: Path, component_name: str, comp_variants: list[VariantSpec], component_role: str, + shared: "SharedConfigAssets", ) -> None: component_dir = output_dir / _MODELS_DIR / component_name component_dir.mkdir(parents=True, exist_ok=True) @@ -677,10 +773,23 @@ def _write_component( dst = dst_dir / entry_name _copy_with_collision_check(entry, dst, skip_if_identical=True) - # Per-variant runtime fields flow through genai_config_overlay.json. - _write_genai_config_overlay(variant_dir, component_role, v) + # Config files ORT-GenAI resolves relative to genai_config.json must + # sit next to it inside the variant directory. + for name, src in shared.variant_local_files.items(): + src_path = Path(src) + if src_path.is_dir(): + dst = variant_dir / name + if not dst.exists(): + shutil.copytree(str(src_path), str(dst)) + elif src_path.is_file(): + _copy_with_collision_check(src_path, variant_dir / name, skip_if_identical=True) + + # Each variant carries a complete, self-contained genai_config.json: + # ORT-GenAI loads /genai_config.json directly and + # never merges a package-level base config. + _write_variant_genai_config(variant_dir, component_role, v, shared) - _write_metadata(component_dir, component_name, comp_variants) + _write_component_json(component_dir, component_name, comp_variants) def _copy_with_collision_check(src: Path, dst: Path, *, skip_if_identical: bool = False) -> None: @@ -709,7 +818,7 @@ def _copy_with_collision_check(src: Path, dst: Path, *, skip_if_identical: bool shutil.copy2(str(src), str(dst)) -def _write_metadata(component_dir: Path, component_name: str, comp_variants: list[VariantSpec]) -> None: +def _write_component_json(component_dir: Path, component_name: str, comp_variants: list[VariantSpec]) -> None: variants_payload: dict[str, Any] = {} for v in comp_variants: # EP fields are inline on the variant object; a variant targets a @@ -720,10 +829,11 @@ def _write_metadata(component_dir: Path, component_name: str, comp_variants: lis if v.compatibility_string: variant_obj["compatibility_string"] = v.compatibility_string variants_payload[v.variant_name] = variant_obj + # The ORT schema allows only component_name / variants / additional_metadata + # here; any other key (a schema_version, for instance) fails the parse. _write_json( - component_dir / "metadata.json", + component_dir / _COMPONENT_FILENAME, { - "schema_version": _METADATA_SCHEMA_VERSION, "component_name": component_name, "variants": variants_payload, }, @@ -738,33 +848,33 @@ def _genai_provider_name(ep: str) -> str: return ep[: -len("ExecutionProvider")].lower() if ep.endswith("ExecutionProvider") else ep -def _write_genai_config_overlay(variant_dir: Path, component_role: str, v: VariantSpec) -> None: - """Emit a per-variant ``genai_config_overlay.json`` (RFC 7386 merge patch). +def _write_variant_genai_config( + variant_dir: Path, component_role: str, v: VariantSpec, shared: "SharedConfigAssets" +) -> None: + """Write a complete, self-contained ``genai_config.json`` into a variant dir. - Per-variant runtime fields flow through a JSON Merge Patch applied on top - of the package's base ``configs/genai_config.json``. The base has every - role's ``filename`` / ``session_options`` / ``pipeline`` stripped (see - ``_strip_variant_specific``); this overlay restores them. + ORT-GenAI loads ``/genai_config.json`` and knows + nothing about a package-level base config or an RFC 7386 merge patch, so + each variant must carry a full config. - When the variant was built per-role (``v.role_name`` is set, the default - CLI path) the overlay lifts only that role's body — each role gets its - own component / variant directory, so each overlay scopes to exactly the - role it represents. The model-level scalars in + The config is built as ``merge(base, per-variant fields)``. ``base`` is the + package's shared ``genai_config.json`` with every variant-specific key + stripped (see ``_strip_variant_specific``); the per-variant fields are + lifted from the variant's own source config and restore ``filename``, + ``session_options``, ``pipeline`` and the model-level scalars in ``_VARIANT_LEVEL_MODEL_KEYS`` (``context_length``, ``eos_token_id``, - ``pad_token_id``, ``bos_token_id``, ``type``) are written into the - overlay of only the source's primary role (``_pick_primary_role``). - Writing them on every per-role overlay would corrupt the merged config - because GenAI's overlay parser appends arrays rather than replacing - them — ``eos_token_id`` is commonly a list, and a three-role VLM - overlay set would triple every entry. + ``pad_token_id``, ``bos_token_id``, ``type``) — values that legitimately + differ across variants, e.g. an NPU build capping ``context_length`` at + 4224 while CPU/CUDA use the full 131072. - When the variant carries every role at once (no ``role_name``, legacy - multi-role-per-component callers) every role's per-variant body is - lifted into the overlay together with the variant-level scalars. + The merge replaces values rather than appending, so list-valued fields + such as ``eos_token_id`` and ``pipeline`` keep the variant's exact value. - Pipeline-shaped roles (multi-stage exports, e.g. QNN) are covered by - the same lift: ``pipeline`` is in the strip set so the base loses it, - the overlay restores it. + When the variant carries every role at once (no ``role_name`` — the CLI + path) every role's per-variant body is lifted together, producing one + complete config for the single component. When a direct caller built the + variant per-role (``v.role_name`` set) only that role's body is lifted; + other roles live in their own components. Direct ``write_model_package`` callers that don't pass ``source_genai`` fall back to the legacy ``inference_settings``-driven shape so existing @@ -777,22 +887,22 @@ def _write_genai_config_overlay(variant_dir: Path, component_role: str, v: Varia if isinstance(src_model, dict): if v.role_name: - # Preferred path: per-role variant. Only lift this role's body. - # Other roles end up in their own components (one component - # per role), each with their own overlay. + # Direct-caller path: per-role variant. Only lift this role's + # body. Other roles end up in their own components (one + # component per role), each with their own overlay. role_body = src_model.get(v.role_name) if isinstance(role_body, dict): role_patch = _lift_role_overlay_body(role_body, v.onnx_rel_paths) if role_patch: model_patch[v.role_name] = role_patch else: - # Legacy multi-role-per-component path: lift every role's - # per-variant fields. Used by direct ``write_model_package`` - # callers / tests that predate per-role components. + # Multi-role path: this variant holds every role of the source + # model, so lift each role's per-variant fields together into + # one complete config. for role_name, role_body in src_model.items(): if not isinstance(role_body, dict): continue - role_patch = _lift_role_overlay_body(role_body) + role_patch = _lift_role_overlay_body(role_body, v.onnx_rel_paths_by_role.get(role_name)) if role_patch: model_patch[role_name] = role_patch else: @@ -834,24 +944,44 @@ def _write_genai_config_overlay(variant_dir: Path, component_role: str, v: Varia # legitimately differ across variants (e.g. NPU runtime caps # ``context_length`` at 4224 while CPU/CUDA use the full 131072; # pad_token_id can differ when one exporter uses the EOS as PAD and - # another uses the sentinel). Without this lift the merged config + # another uses the sentinel). Without this lift the variant config # would silently use whichever variant happened to win the base - # selection. For per-role variants only the primary role's overlay - # carries these so the same scalar isn't append-merged once per - # component (critical for list-valued ``eos_token_id``). + # selection. Every variant carries them: the merge below replaces + # values, so a list-valued ``eos_token_id`` keeps this variant's + # exact value instead of accumulating entries. if isinstance(src_model, dict): - is_primary = (not v.role_name) or (v.role_name == _pick_primary_role(src_genai)) - if is_primary: - for k in _VARIANT_LEVEL_MODEL_KEYS: - if k in src_model: - # Deep-copy via JSON round-trip so we never share refs with - # the caller's dict; arrays in particular must be - # independent because GenAI's overlay parser treats arrays - # as append-merge. - model_patch[k] = json.loads(json.dumps(src_model[k])) + for k in _VARIANT_LEVEL_MODEL_KEYS: + if k in src_model: + model_patch[k] = src_model[k] + + # Model-level scalars the variant did not declare fall back to the base + # source's values. Without this a variant carrying no source config would + # lose them entirely: the base strips them, and nothing would restore them. + config = _json_merge(shared.base_genai, {"model": dict(shared.model_level_defaults)}) + config = _json_merge(config, {"model": model_patch}) + + # Point at the shared tokenizer asset. ORT resolves ``sha256:`` to the + # content-addressed directory; without this the tokenizer would be looked + # up next to genai_config.json, where it deliberately does not live. + if shared.tokenizer_asset_uri: + config.setdefault("model", {})["tokenizer_dir"] = shared.tokenizer_asset_uri + + _write_json(variant_dir / "genai_config.json", config) + + +def _json_merge(base: Any, patch: Any) -> Any: + """Recursively merge ``patch`` onto ``base``, returning a fresh deep copy. - overlay = {"model": model_patch} - _write_json(variant_dir / "genai_config_overlay.json", overlay) + Objects merge key-by-key; every other value (including lists) is replaced + outright. Deep-copied via a JSON round-trip so the result never shares + references with either input. + """ + if isinstance(base, dict) and isinstance(patch, dict): + merged = dict(base) + for k, v in patch.items(): + merged[k] = _json_merge(merged[k], v) if k in merged else v + return json.loads(json.dumps(merged)) + return json.loads(json.dumps(patch)) def _lift_role_overlay_body(role_body: dict, onnx_rel_paths: Optional[list[str]] = None) -> dict: @@ -862,24 +992,20 @@ def _lift_role_overlay_body(role_body: dict, onnx_rel_paths: Optional[list[str]] EP knobs). All three are stripped from the base genai_config; this helper recovers them as the role's overlay patch. - Filenames are normalised to their basenames. In the per-role-component - layout each role gets its own variant directory under - ``models///`` and the writer places the ONNX(s) there - flat — the original source-side ``decoder/`` / ``vision_encoder/`` / - ``embedding/`` subdirectory prefixes (Mobius VLM convention) are no - longer needed to disambiguate sibling roles inside one variant dir. - When ``onnx_rel_paths`` is supplied (preferred), the role's - ``filename`` is replaced by the writer-known package-relative path so - the overlay matches the on-disk layout exactly even if the source - diverged. Pipeline-stage filenames are likewise rewritten to their - basenames so the per-stage references resolve inside the flat variant - directory. + When ``onnx_rel_paths`` is supplied (the CLI path), each filename is + replaced by the writer-known package-relative path so the overlay + matches the on-disk layout exactly. Those paths preserve the source's + directory structure (``vision_encoder/model.onnx``), which is what keeps + sibling roles from colliding inside the shared variant directory. + Without them the filename falls back to its basename, matching the flat + layout ``_variant_artifacts`` writes for direct callers that supply only + ``onnx_files``. Every filename — top-level role and pipeline stages alike — is - validated as a safe relative path before basename normalisation; - absolute paths or upward traversal raise rather than silently - propagate into a generated overlay. Pipeline and session_options are - deep-copied to avoid aliasing with the caller's dict. + validated as a safe relative path before normalisation; absolute paths + or upward traversal raise rather than silently propagate into a + generated overlay. Pipeline and session_options are deep-copied to + avoid aliasing with the caller's dict. """ patch: dict[str, Any] = {} pipeline = role_body.get("pipeline") @@ -961,15 +1087,13 @@ def _strip_variant_specific( ) -> Any: """Recursively drop variant-specific keys from a genai_config-shaped dict. - ``filename`` and ``session_options`` are intrinsically variant-specific and - must not live in the package's base ``configs/genai_config.json``; per-variant - ``genai_config_overlay.json`` files patch them back in. ``pipeline`` is - also stripped because GenAI's overlay parser appends arrays rather than - replacing them — a pipeline present in both base and overlay would - duplicate every stage on merge. The same logic applies to per-variant - model-level scalars listed in ``_VARIANT_LEVEL_MODEL_KEYS`` (e.g. - ``context_length`` differs between NPU and GPU variants of the same - model). Returns a deep copy. + ``filename``, ``session_options`` and ``pipeline`` are intrinsically + variant-specific and must not leak from one variant's config into another, + so they are stripped from the shared base and re-applied per variant from + that variant's own source config. The same applies to the model-level + scalars listed in ``_VARIANT_LEVEL_MODEL_KEYS`` (e.g. ``context_length`` + differs between NPU and GPU variants of the same model). Returns a deep + copy. """ if isinstance(node, dict): return {k: _strip_variant_specific(v, keys) for k, v in node.items() if k not in keys} @@ -1037,108 +1161,121 @@ def _write_manifest( "schema_version": _MANIFEST_SCHEMA_VERSION, "package_name": package_name, "package_version": package_version, - "components": components, - "configs_dir": _CONFIGS_DIR, + # ORT requires an object mapping component name -> component location. + # Pointing at the directory makes ORT read ``component.json`` from it. + "components": {name: f"{_MODELS_DIR}/{name}" for name in components}, } if producer_info: - # Olive-specific provenance under a namespaced key so future schema - # evolution can't collide with it. - manifest["producer"] = producer_info + # The ORT schema rejects unknown top-level manifest keys, so Olive + # provenance is namespaced inside the free-form additional_metadata. + manifest["additional_metadata"] = {"producer": producer_info} _write_json(output_dir / "manifest.json", manifest) # --------------------------------------------------------------------------- -# configs/ handling +# shared asset handling # --------------------------------------------------------------------------- -def _copy_config_files( - output_dir: Path, - config_files: dict[str, Path], - role_to_component: Optional[dict[str, str]] = None, -) -> None: - configs_dir = output_dir / _CONFIGS_DIR - configs_dir.mkdir(parents=True, exist_ok=True) - configs_root = configs_dir.resolve() +def _write_shared_assets(output_dir: Path, config_files: dict[str, Path]) -> "SharedConfigAssets": + """Materialize package-level config material and describe it to the writer. + + ``genai_config.json`` is not copied verbatim: it becomes the base every + variant's complete config is merged from, with variant-specific keys + stripped. Processor configs are handed back for per-variant copying because + ORT-GenAI resolves those relative to ``genai_config.json``. Everything else + (tokenizer assets) is written once into a content-addressed + ``shared_assets/sha256-/`` directory that variants reference by URI. + """ + shared = SharedConfigAssets() + if not config_files: + return shared + + staged: dict[str, Path] = {} for name, src in config_files.items(): if "/" in name or "\\" in name or name in ("", ".", ".."): logger.warning("Skipping config file with unsafe name %r.", name) continue src_path = Path(src) - dest = configs_dir / name - # Belt-and-suspenders: even with the name check above, refuse a dest - # that doesn't land directly under configs/. - if dest.resolve().parent != configs_root: - logger.warning("Skipping config file %r: resolved path escapes configs/.", name) + if not src_path.exists(): + logger.warning("Config source %s does not exist; skipping.", src_path) continue - if dest.exists(): - if not _paths_equal(src_path, dest): - logger.warning( - "configs/%s already present and differs from %s; keeping the existing copy. " - "Per-variant config differences belong in genai_config_overlay.json, " - "not in the shared configs/ directory.", - name, - src_path, - ) + if name == "genai_config.json": + if src_path.is_file(): + try: + with src_path.open(encoding="utf-8") as fh: + base = json.load(fh) + shared.base_genai = _strip_variant_specific(base) + base_model = base.get("model") + if isinstance(base_model, dict): + shared.model_level_defaults = { + k: base_model[k] for k in _VARIANT_LEVEL_MODEL_KEYS if k in base_model + } + except Exception: + logger.debug("Failed to read base genai_config from %s.", src_path, exc_info=True) continue - if name == "genai_config.json" and src_path.is_file(): - # Strip variant-specific keys from the base genai_config and inject - # ``model..component`` markers so ORT-GenAI can resolve each - # role to a package component (and apply the right per-variant - # overlay). Each variant's genai_config_overlay.json patches the - # stripped keys back in. - try: - with src_path.open(encoding="utf-8") as fh: - base_genai = json.load(fh) - stripped = _strip_variant_specific(base_genai) - if role_to_component: - _inject_role_components(stripped, role_to_component) - _write_json(dest, stripped) - continue - except Exception: - logger.debug( - "Failed to strip variant-specific keys from %s; falling back to verbatim copy.", - src_path, - exc_info=True, - ) + if name in _VARIANT_LOCAL_CONFIG_NAMES: + shared.variant_local_files[name] = src_path + continue + staged[name] = src_path + + if not staged: + return shared + + # Content-address the asset so two packages shipping the same tokenizer + # resolve to the same URI and can share one on-disk copy once installed. + asset_root = output_dir / _SHARED_ASSETS_DIR + tmp_dir = asset_root / "_staging" + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + tmp_dir.mkdir(parents=True) + for name, src_path in staged.items(): + dest = tmp_dir / name if src_path.is_dir(): shutil.copytree(str(src_path), str(dest)) - elif src_path.is_file(): - shutil.copy2(str(src_path), str(dest)) else: - logger.warning("Config source %s does not exist; skipping.", src_path) + shutil.copy2(str(src_path), str(dest)) + digest = _compute_directory_hash(tmp_dir) + asset_dir = asset_root / f"sha256-{digest}" + if asset_dir.exists(): + shutil.rmtree(tmp_dir) + else: + tmp_dir.rename(asset_dir) + + # ``model.tokenizer_dir`` is only meaningful when the asset actually holds a + # tokenizer. ORT-GenAI delegates tokenizer loading to onnxruntime-extensions, + # which unconditionally opens ``/tokenizer_config.json`` and + # hard-fails when it is missing; every other file it reads is optional or + # named by that config (``tokenizer.json``, ``chat_template.jinja``, + # ``tokenizer_module.json``, an arbitrary ``tiktoken_file``). So the presence + # of ``tokenizer_config.json`` is the one reliable signal — pointing + # ``tokenizer_dir`` at a shared asset without it would only misdirect the + # error message for a package that has no tokenizer either way. + if _TOKENIZER_SENTINEL in staged: + shared.tokenizer_asset_uri = f"sha256:{digest}" + elif staged: + logger.warning( + "Shared assets contain no %s; leaving model.tokenizer_dir unset. Files staged: %s.", + _TOKENIZER_SENTINEL, + ", ".join(sorted(staged)), + ) + return shared -def _inject_role_components(genai: dict, role_to_component: dict[str, str]) -> None: - """Inject ``model..component = `` markers in-place. - ORT-GenAI's model-package variant selection requires every role block in - the base ``configs/genai_config.json`` to declare which package component - serves it. Olive-generated source ``genai_config.json`` typically lacks - these markers because the source is a flat-directory build, not a package. +def _compute_directory_hash(source_dir: Path) -> str: + r"""Return the canonical shared-asset digest for ``source_dir``. + + Mirrors ORT's ``ModelPackage_ComputeDirectoryHash``: hash every regular + file, build ``" \n"`` lines sorted by path, + and hash that manifest text. Hashing names as well as contents means + renaming a file inside the asset changes its URI. """ - model_block = genai.get("model") - if not isinstance(model_block, dict): - return - for role, component in role_to_component.items(): - role_block = model_block.get(role) - if isinstance(role_block, dict): - role_block["component"] = component - - -def _paths_equal(a: Path, b: Path) -> bool: - """Return True if a and b have identical content (file or directory).""" - if a.is_file() and b.is_file(): - if a.stat().st_size != b.stat().st_size: - return False - return _sha256_file(a) == _sha256_file(b) - if a.is_dir() and b.is_dir(): - a_entries = sorted(p.name for p in a.iterdir()) - b_entries = sorted(p.name for p in b.iterdir()) - if a_entries != b_entries: - return False - return all(_paths_equal(a / name, b / name) for name in a_entries) - return False + files = sorted( + (p for p in source_dir.rglob("*") if p.is_file()), key=lambda p: p.relative_to(source_dir).as_posix() + ) + manifest = "".join(f"{_sha256_file(p)} {p.relative_to(source_dir).as_posix()}\n" for p in files) + return hashlib.sha256(manifest.encode("utf-8")).hexdigest() # --------------------------------------------------------------------------- @@ -1338,29 +1475,6 @@ def _load_source_genai(source_path: Path) -> Optional[dict]: return None -def _pick_primary_role(source_genai: Optional[dict]) -> Optional[str]: - """Pick the genai_config role that names the model's primary component. - - A genai_config's ``model`` block keys mix per-role objects (``decoder``, - ``embedding``, ...) with model-level scalars (``vocab_size``, - ``context_length``, ...). The primary role is the first key whose value - is an object carrying either a ``filename`` (flat variant) or a - ``pipeline`` (multi-stage variant). Returns ``None`` when no such role - is found (e.g. genai_config missing or malformed). - """ - if not isinstance(source_genai, dict): - return None - model_block = source_genai.get("model") - if not isinstance(model_block, dict): - return None - for role, body in model_block.items(): - if not isinstance(body, dict): - continue - if "pipeline" in body or "filename" in body: - return role - return None - - def _model_artifact_dirs(source_genai: Optional[dict]) -> set[str]: """Return source-root subdirectory names that hold ONNX model artifacts. @@ -1419,13 +1533,13 @@ def _collect_artifacts_per_role(source_path: Path, source_genai: Optional[dict]) The role's value is the ordered list of artifacts that belong to it — one for a flat role, one per stage for a pipeline role. - Every artifact's ``package_rel_path`` is the source filename's basename: - in the per-role-component layout, each role gets its own variant - directory under ``models///``, so files don't need - subdirectory prefixes to disambiguate from sibling roles' ONNXes - (they no longer share a directory). The source filename's subdirectory - is used only to locate the file on disk in the source — it does not - propagate into the package. + Every artifact's ``package_rel_path`` preserves the source's declared + relative location (e.g. ``vision_encoder/model.onnx``). All roles share + one variant directory, so basenames alone would collide whenever a + producer names every role's graph ``model.onnx`` (the Mobius VLM + convention). Keeping the declared path is inherently collision-free — + the files already coexisted under one source directory — and lets each + role's ``filename`` stay exactly as the source declared it. Filenames are validated as safe relative paths; absolute or upward-traversing entries raise ``ValueError``. A source with no @@ -1445,7 +1559,7 @@ def _validated_artifact(filename: str, kind: str, role: str) -> OnnxArtifact: f"Source {source_path} role {role!r} {kind} {filename!r} is not a safe " "relative path (absolute paths and '..' segments are rejected)." ) - return OnnxArtifact(source_path=source_path / filename, package_rel_path=Path(filename).name) + return OnnxArtifact(source_path=source_path / filename, package_rel_path=filename) for role_name, role_body in model_block.items(): if not isinstance(role_body, dict): @@ -1516,6 +1630,44 @@ def _resolve_ep_for_role(source_genai: Optional[dict], role_name: str) -> str: return "CPUExecutionProvider" +def _resolve_ep_for_variant( + source_path: Path, + source_genai: Optional[dict], + artifacts_by_role: dict[str, list[OnnxArtifact]], +) -> str: + """Pick the single ORT EP for a variant holding every role of a model. + + A variant declares one ``ep`` in the manifest, and ORT-GenAI builds all + of a model's sessions against one device: a role whose provider resolves + to a different non-CPU device makes the model fail to load with "Running + a model with multiple providers is not supported". CPU roles are exempt — + ORT-GenAI registers CPU implicitly and a CPU role never claims the + device — so a package mixing, say, a CPU vision encoder with a QNN + decoder is legitimate and takes the decoder's EP. + + Raises when two roles demand different non-CPU EPs, since the resulting + package would be unloadable. + """ + cpu_ep = "CPUExecutionProvider" + eps_by_role: dict[str, str] = {} + for role_name in artifacts_by_role: + ep = _resolve_ep_for_role(source_genai, role_name) + if ep != cpu_ep: + eps_by_role[role_name] = ep + + distinct = sorted(set(eps_by_role.values())) + if len(distinct) > 1: + detail = ", ".join(f"{role}={ep}" for role, ep in sorted(eps_by_role.items())) + raise ValueError( + f"Source {source_path} declares more than one non-CPU execution provider across its " + f"genai_config roles ({detail}). All roles of a model share one package variant and " + "ORT-GenAI runs them on a single device, so they must target the same execution " + "provider (roles left on CPU are fine). Split the roles into separate per-EP sources " + "or align their session_options.provider_options." + ) + return distinct[0] if distinct else cpu_ep + + def _select_base_config_source( targets: list[tuple[str, Path, dict]], ) -> tuple[str, Path, dict]: diff --git a/olive/cli/run.py b/olive/cli/run.py index a0a91e8115..0861a50b79 100644 --- a/olive/cli/run.py +++ b/olive/cli/run.py @@ -55,16 +55,22 @@ def register_subcommand(parser: ArgumentParser): @action def run(self): + from copy import deepcopy + from pathlib import Path + from olive.common.config_utils import load_config_file from olive.workflows import run as olive_run # allow the run_config to be a dict already (for api use) - run_config = self.args.run_config - if not isinstance(run_config, dict): - run_config = load_config_file(run_config) + run_config_input = self.args.run_config + run_config = ( + deepcopy(run_config_input) if isinstance(run_config_input, dict) else load_config_file(run_config_input) + ) + config_overrides = {} if input_model_config := get_input_model_config(self.args, required=False): print("Replacing input model config in run config") run_config["input_model"] = input_model_config + config_overrides["input_model"] = input_model_config elif self.args.test not in (None, False): input_model = run_config.get("input_model") if not isinstance(input_model, dict) or input_model.get("type", "").lower() != "hfmodel": @@ -84,6 +90,19 @@ def run(self): run_config.get("engine", {}).pop(rc_key, None) # add value to run config directly run_config[rc_key] = arg_value + config_overrides[rc_key] = arg_value + + recipe_telemetry_metadata = { + "recipe_command": "WorkflowRun", + "recipe_source": "config_dict" if isinstance(run_config_input, dict) else "config_file", + "recipe_format": "dict" + if isinstance(run_config_input, dict) + else Path(run_config_input).suffix.lstrip(".").lower() or "unknown", + "execution_mode": "list_required_packages" if self.args.list_required_packages else "run", + "package_config_provided": bool(self.args.package_config), + } + if config_overrides: + recipe_telemetry_metadata["config_overrides"] = config_overrides output_path = run_config.get("output_dir") or run_config.get("engine", {}).get("output_dir") workflow_output = olive_run( @@ -91,6 +110,8 @@ def run(self): list_required_packages=self.args.list_required_packages, tempdir=self.args.tempdir, package_config=self.args.package_config, + recipe_telemetry_metadata=recipe_telemetry_metadata, + emit_error_telemetry=False, ) if self.args.test not in (None, False): mark_test_output_path(output_path) diff --git a/olive/common/hf/io_config/io_resolver.py b/olive/common/hf/io_config/io_resolver.py index 91882fb72e..13a585aed2 100644 --- a/olive/common/hf/io_config/io_resolver.py +++ b/olive/common/hf/io_config/io_resolver.py @@ -113,10 +113,30 @@ def get_aliases() -> dict[str, list[str]]: def _get_nested_attr(obj, attr_path: str): - """Get nested attribute using dot notation (e.g., 'vision_config.image_size').""" - parts = attr_path.split(".") - for part in parts: - obj = getattr(obj, part, None) + """Get nested attribute using dot notation (e.g., 'vision_config.image_size'). + + Attribute access is guarded with a broad ``except Exception`` (rather than relying on + ``getattr``'s default, which only swallows ``AttributeError``) because transformers 5.x + heterogeneous / per-layer configs raise their own exception types on ambiguous access. + For example Gemma4-family configs raise + ``transformers.integrations.heterogeneity.configuration_utils.AmbiguousGlobalPerLayerAttributeError`` + when a per-layer attribute such as ``head_dim`` is read off the top-level config without a + layer index. That exception type is not importable across every transformers version Olive + supports, so it is caught structurally instead of by type: an attribute Olive cannot read + unambiguously is simply "not resolvable here", and the caller falls back to the next alias. + + Each path segment may also land on a plain ``dict`` node instead of a ``PretrainedConfig``. + ``ModelWrapper`` accepts a raw (serialized) config ``dict`` and converts it with the base + ``PretrainedConfig.from_dict``, which does not know how to reconstruct model-specific nested + sub-configs (e.g. ``text_config``) into config objects -- they simply stay plain dicts. + ``getattr`` would silently return ``None`` for every attribute of such a dict, so dict nodes + are looked up with ``.get`` instead. + """ + for part in attr_path.split("."): + try: + obj = obj.get(part, None) if isinstance(obj, dict) else getattr(obj, part, None) + except Exception: # pylint: disable=broad-except + return None if obj is None: return None return obj @@ -132,6 +152,12 @@ def resolve_alias(config, name: str, aliases: dict[str, list[str]] | None = None Supports nested attribute paths (e.g., 'vision_config.image_size'). + Resolution order: the canonical ``name`` on the config itself wins, and the aliases are + only consulted as fallbacks. This matters for composite (VL) configs that carry *both* a + top-level value and a nested ``text_config.`` alias with a different value -- the + model's own top-level attribute is authoritative and must not be silently overridden by + the nested fallback. + Args: config: The config object to query. name: The canonical name to look up (e.g., 'hidden_size', 'num_layers'). @@ -144,17 +170,18 @@ def resolve_alias(config, name: str, aliases: dict[str, list[str]] | None = None if aliases is None: aliases = get_aliases() - # Check aliases first + # Direct lookup of the canonical name takes precedence over any alias. + value = _get_nested_attr(config, name) + if value is not None and _is_valid_config_value(value): + return value + + # Fall back to aliases, in declaration order. if name in aliases: for attr in aliases[name]: value = _get_nested_attr(config, attr) - if _is_valid_config_value(value) and value is not None: + if value is not None and _is_valid_config_value(value): return value - # Try direct lookup - value = _get_nested_attr(config, name) - if _is_valid_config_value(value): - return value return None diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py index 07d18b318a..fb893a72ab 100644 --- a/olive/common/hf/quant.py +++ b/olive/common/hf/quant.py @@ -2,6 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +# pylint: disable=attribute-defined-outside-init,protected-access from __future__ import annotations import logging @@ -15,6 +16,27 @@ logger = logging.getLogger(__name__) +# ONNX Runtime's ``com.microsoft::MatMulNBits`` and ``com.microsoft::GatherBlockQuantized`` +# kernels both hard-enforce that the quantization block size is a power of 2 and at least +# 16 (see ``onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc``'s +# ``ORT_ENFORCE`` and the MatMulNBits op contract). Violating it only surfaces as an opaque +# native session-initialization crash, so validate at module-construction time instead. +_MIN_ONNX_BLOCK_SIZE = 16 + + +def _validate_onnx_block_size(block_size: int, op_name: str, fallback_dim_name: str, fallback_dim: int) -> None: + """Raise a clear ``ValueError`` if ``block_size`` is not a valid ONNX Runtime block size.""" + if block_size >= _MIN_ONNX_BLOCK_SIZE and (block_size & (block_size - 1)) == 0: + return + raise ValueError( + f"com.microsoft::{op_name} requires block_size (group_size) to be a power of 2 and >= " + f"{_MIN_ONNX_BLOCK_SIZE}, got {block_size}. A non-positive ``group_size`` falls back to " + f"the full {fallback_dim_name} ({fallback_dim}), which is not a valid block size here. " + f"Re-quantize with an explicit power-of-2 ``group_size`` (e.g. 32, 64, 128) that divides " + f"{fallback_dim_name}={fallback_dim} instead of relying on the per-tensor/per-channel " + "fallback." + ) + def get_quantization_info(model: torch.nn.Module): """Get the quantization info from the model.""" @@ -118,7 +140,11 @@ class QuantLinearTorchFunction(torch.autograd.Function): # pylint: disable=W0223,W0221 @staticmethod def symbolic(g, x, qweight, scales, qzeros, g_idx, bits, group_size, in_features, out_features, dynamo): - tensor_args = [x, qweight, scales, qzeros] + tensor_args = [x, qweight, scales] + if qzeros is not None: + tensor_args.append(qzeros) + elif g_idx is not None: + raise ValueError("MatMulNBits with g_idx requires zero_points; symmetric quantization is not supported.") if g_idx is not None: tensor_args.append(g_idx) attrs = { @@ -160,7 +186,13 @@ def forward( if torch.onnx.is_in_onnx_export(): if dynamo and hasattr(torch.onnx, "ops"): # torch.onnx.ops was introduced in 2.8 - tensor_args = [x, qweight, scales, qzeros] + tensor_args = [x, qweight, scales] + if qzeros is not None: + tensor_args.append(qzeros) + elif g_idx is not None: + raise ValueError( + "MatMulNBits with g_idx requires zero_points; symmetric quantization is not supported." + ) if g_idx is not None: tensor_args.append(g_idx) attrs = { @@ -202,12 +234,14 @@ def __init__( bits: int = 4, dtype: torch.dtype = torch.float32, dynamo: bool = False, + has_qzeros: bool = True, ): super().__init__() self.in_features = in_features self.out_features = out_features self.group_size = group_size if group_size > 0 else in_features + _validate_onnx_block_size(self.group_size, "MatMulNBits", "in_features", in_features) assert bits in [4, 8], "Only 4 and 8 bits are supported for QuantLinearNbit" self.bits = bits @@ -220,10 +254,16 @@ def __init__( dtype=torch.uint8, ), ) - self.register_buffer( - "qzeros", - torch.zeros(out_features, math.ceil(n_blocks_per_col * self.bits / 8), dtype=torch.uint8), - ) + if has_qzeros: + self.register_buffer( + "qzeros", + torch.zeros(out_features, math.ceil(n_blocks_per_col * self.bits / 8), dtype=torch.uint8), + ) + else: + # When qzeros is omitted the contrib op interprets the + # zero point as the unsigned mid value (e.g. 8 for 4-bit), + # which matches Olive's symmetric quantization convention. + self.qzeros = None self.register_buffer("scales", torch.zeros((out_features, n_blocks_per_col), dtype=dtype)) if g_idx: self.register_buffer( @@ -255,8 +295,11 @@ def forward(self, x): def pack(self, iweight, izeros, scales, g_idx=None): """Pack int8 weight and zeros to int4 and int8 respectively. - iweight, izeros and scales must have out_features as the first dimension - iweight, izeros must be uint8 tensors with each holding 4/8 bit values + ``iweight`` and ``scales`` must have ``out_features`` as the + first dimension. ``iweight`` must be a uint8 tensor with each + value holding ``bits`` bits. ``izeros`` may be ``None`` for + symmetric quantization (the contrib op then treats the zero + point as the unsigned mid value). """ # pylint: disable=W0201 # shapes for packing @@ -277,12 +320,15 @@ def pack(self, iweight, izeros, scales, g_idx=None): iweight = (iweight[:, 0::2] & 0xF) | ((iweight[:, 1::2] & 0xF) << 4) self.qweight = iweight.reshape(n, k_blocks, blob_size).contiguous() - # pad to make the K dimension even - izeros = torch.nn.functional.pad(izeros, (0, izeros.shape[-1] & 1), value=0) - # pack the zeros - if bits == 4: - izeros = (izeros[:, 0::2] & 0xF) | ((izeros[:, 1::2] & 0xF) << 4) - self.qzeros = izeros.contiguous() + if izeros is not None: + # pad to make the K dimension even + izeros = torch.nn.functional.pad(izeros, (0, izeros.shape[-1] & 1), value=0) + # pack the zeros + if bits == 4: + izeros = (izeros[:, 0::2] & 0xF) | ((izeros[:, 1::2] & 0xF) << 4) + self.qzeros = izeros.contiguous() + else: + self.qzeros = None self.scales = scales.contiguous() @@ -293,7 +339,7 @@ def from_tensors( cls, iweight: torch.Tensor, scales: torch.Tensor, - izeros: torch.Tensor, + izeros: torch.Tensor | None, group_size: int, bits: int = 4, g_idx: torch.Tensor | None = None, @@ -303,6 +349,8 @@ def from_tensors( """Create a QuantLinearNbit instance from the given tensors. Weight is expected to be in in_features x out_features layout and unsigned. + ``izeros`` may be ``None`` for symmetric quantization — the + contrib op then assumes a midq zero point. """ new_qlinear = cls( group_size, @@ -313,12 +361,245 @@ def from_tensors( bits=bits, dtype=scales.dtype, dynamo=dynamo, + has_qzeros=izeros is not None, + ) + new_qlinear.pack( + iweight.to(torch.uint8).t(), + izeros.to(torch.uint8).t() if izeros is not None else None, + scales.t(), + g_idx, ) - new_qlinear.pack(iweight.to(torch.uint8).t(), izeros.to(torch.uint8).t(), scales.t(), g_idx) if bias is not None: new_qlinear.bias = bias.clone() return new_qlinear + @classmethod + def from_quant_tensor( + cls, + qt, + bias: torch.Tensor | None = None, + dynamo: bool = False, + ) -> QuantLinearNbit: + """Create a ``QuantLinearNbit`` from an Olive 2D ``QuantTensor``. + + Both layouts pack the quantization axis (the input dim) as uint8 + with the same in-byte order, so the buffer **bytes** are + bit-identical and we only need to reshape: + + * ``QuantTensor.qweight`` ``(out, in / pack_factor)`` -> + ``QuantLinearNbit.qweight`` ``(out, n_blocks, blob_size)`` where + ``n_blocks * blob_size == in / pack_factor``. + * ``scales`` already match in shape (``(out, n_blocks)``). + * ``qzeros`` already match in shape; for symmetric weights + ``QuantTensor.qzeros`` is ``None`` and the contrib op + interprets the missing input as a midq zero point, so the + export module simply omits the buffer. + """ + from olive.common.quant.tensor import QuantTensor + + if not isinstance(qt, QuantTensor) or qt.dim() != 2: + raise ValueError("QuantLinearNbit.from_quant_tensor requires a 2D QuantTensor") + if qt.bits == 2: + raise ValueError( + "2-bit QuantTensor cannot be exported to ONNX; QuantLinearNbit only supports " + "4-bit and 8-bit packing (matching the com.microsoft.MatMulNBits contrib op). " + "2-bit quantization is only usable for the PyTorch checkpoint path." + ) + + out_features, in_features = qt.shape + group_size = qt.group_size if qt.group_size > 0 else in_features + new = cls( + group_size=group_size, + in_features=in_features, + out_features=out_features, + g_idx=False, + bias=bias is not None, + bits=qt.bits, + dtype=qt.scales.dtype, + dynamo=dynamo, + has_qzeros=qt.qzeros is not None, + ) + + # bit-identical layouts -> reshape (no unpack/repack) + new.qweight = qt.qweight.detach().clone().reshape(new.qweight.shape).contiguous() + new.scales = qt.scales.detach().clone().reshape(new.scales.shape).contiguous() + if qt.qzeros is not None: + new.qzeros = qt.qzeros.detach().clone().reshape(new.qzeros.shape).contiguous() + + if bias is not None: + new.bias = bias.detach().clone() + return new + + +class QuantEmbeddingTorchFunction(torch.autograd.Function): + """Export a quantized embedding lookup as ``com.microsoft::GatherBlockQuantized``.""" + + # pylint: disable=W0223,W0221 + @staticmethod + def symbolic(g, x, qweight, scales, qzeros, bits, group_size, embedding_dim, dynamo): + tensor_args = [qweight, x, scales] + if qzeros is not None: + tensor_args.append(qzeros) + attrs = {"bits_i": bits, "block_size_i": group_size} + + output = g.op( + "com.microsoft::GatherBlockQuantized", + *tensor_args, + outputs=1, + **attrs, + ) + input_shape = x.type().varyingSizes() + if input_shape is not None and hasattr(x.type(), "with_sizes"): + output_type = scales.type().with_sizes([*input_shape, embedding_dim]) + output.setType(output_type) + return output + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, + bits: int, + group_size: int, + embedding_dim: int, + dynamo: bool = False, + ): + if torch.onnx.is_in_onnx_export(): + if dynamo and hasattr(torch.onnx, "ops"): + tensor_args = [qweight, x, scales] + if qzeros is not None: + tensor_args.append(qzeros) + attrs = {"bits": bits, "block_size": group_size} + return torch.onnx.ops.symbolic( + "com.microsoft::GatherBlockQuantized", + tensor_args, + attrs=attrs, + dtype=scales.dtype, + shape=[*x.shape, embedding_dim], + version=1, + ) + if dynamo: + raise NotImplementedError("torch dynamo export for quantized embedding requires torch 2.8 or higher.") + return torch.zeros((*x.shape, embedding_dim), dtype=scales.dtype, device=x.device) + raise NotImplementedError("QuantEmbeddingTorchFunction forward is only implemented for onnx export") + + +class QuantEmbeddingNbit(torch.nn.Module): + """Quantized embedding layer exported as ``com.microsoft::GatherBlockQuantized``. + + Buffer layout matches Olive's :class:`QuantTensor` 2D representation + exactly (packed uint8 along the last dim), so swapping a host + ``nn.Embedding`` with a ``QuantTensor`` weight for this module is a + direct buffer move. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + group_size: int, + bits: int = 4, + symmetric: bool = True, + padding_idx: int | None = None, + dtype: torch.dtype = torch.float32, + device: torch.device | None = None, + dynamo: bool = False, + ): + super().__init__() + if bits not in (2, 4, 8): + raise ValueError(f"QuantEmbeddingNbit only supports 2/4/8 bits, got {bits}") + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + self.bits = bits + self.symmetric = symmetric + self.padding_idx = padding_idx + self.dynamo = dynamo + + self.group_size = group_size if group_size > 0 else embedding_dim + _validate_onnx_block_size(self.group_size, "GatherBlockQuantized", "embedding_dim", embedding_dim) + pack_factor = 8 // bits + n_groups = math.ceil(embedding_dim / self.group_size) + + self.register_buffer( + "qweight", + torch.zeros( + (num_embeddings, math.ceil(embedding_dim / pack_factor)), + dtype=torch.uint8, + device=device, + ), + ) + self.register_buffer( + "scales", + torch.zeros((num_embeddings, n_groups), dtype=dtype, device=device), + ) + if symmetric: + self.qzeros = None + else: + self.register_buffer( + "qzeros", + torch.zeros( + (num_embeddings, math.ceil(n_groups / pack_factor)), + dtype=torch.uint8, + device=device, + ), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return QuantEmbeddingTorchFunction.apply( + x, + self.qweight, + self.scales, + self.qzeros, + self.bits, + self.group_size, + self.embedding_dim, + self.dynamo, + ) + + @classmethod + def from_quant_tensor( + cls, + qt, + padding_idx: int | None = None, + dynamo: bool = False, + ) -> QuantEmbeddingNbit: + """Create a ``QuantEmbeddingNbit`` from an Olive 2D ``QuantTensor``. + + The QuantTensor 2D buffer layout is bit-identical to this + module's layout, so the buffers are reused directly via + ``detach().clone().reshape(...)``. The explicit ``reshape`` to the + freshly-constructed module's buffer shapes mirrors + :meth:`QuantLinearNbit.from_quant_tensor` and is what makes an incompatible + ``group_size`` (e.g. a per-tensor ``group_size == 0`` QuantTensor, whose + ``scales`` are ``(1, 1)`` rather than ``(num_embeddings, n_groups)``) fail + loudly and immediately here, instead of silently installing a wrong-shaped + buffer that only blows up much later inside the exported graph. + """ + from olive.common.quant.tensor import QuantTensor + + if not isinstance(qt, QuantTensor) or qt.dim() != 2: + raise ValueError("QuantEmbeddingNbit.from_quant_tensor requires a 2D QuantTensor") + + num_embeddings, embedding_dim = qt.shape + new = cls( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + group_size=qt.group_size if qt.group_size > 0 else embedding_dim, + bits=qt.bits, + symmetric=qt.symmetric, + padding_idx=padding_idx, + dtype=qt.scales.dtype, + device=qt.qweight.device, + dynamo=dynamo, + ) + new.qweight = qt.qweight.detach().clone().reshape(new.qweight.shape).contiguous() + new.scales = qt.scales.detach().clone().reshape(new.scales.shape).contiguous() + if qt.qzeros is not None: + new.qzeros = qt.qzeros.detach().clone().reshape(new.qzeros.shape).contiguous() + return new + def make_auto_awq_qlinearnbit(qlinear, dynamo): if qlinear.w_bit not in [4, 8]: @@ -374,6 +655,13 @@ def make_auto_gptq_qlinearnbit(qlinear, dynamo): def make_export_compatible_quant(model: torch.nn.Module, dynamo: bool) -> torch.nn.Module: """Make the model export compatible by replacing the quantized linear layers with 4-bit versions.""" + # Olive-native path: replace nn.Linear / nn.Embedding whose weight is a + # ``QuantTensor`` parameter with the ONNX-exportable wrappers + # ``QuantLinearNbit`` / ``QuantEmbeddingNbit`` defined in this module. + # Done first so subsequent passes that cast the model dtype don't + # dequantize through QuantTensor. + _replace_olive_quant_tensor_modules(model, dynamo) + model, modified = _replace_qlinear_modules( model, dynamo, EXPORT_QLINEAR_MAPPING, "Making export compatible quantized model" ) @@ -381,3 +669,47 @@ def make_export_compatible_quant(model: torch.nn.Module, dynamo: bool) -> torch. # set quantization method to None, gptq doesn't allow dtype casting model.quantization_method = None return model + + +def _replace_olive_quant_tensor_modules(model: torch.nn.Module, dynamo: bool) -> None: + """Swap host modules whose weight is a ``QuantTensor`` with export-compatible wrappers. + + Targets ``nn.Linear`` / ``nn.Embedding`` and replaces them with + :class:`QuantLinearNbit` / :class:`QuantEmbeddingNbit`. + """ + from olive.common.quant.tensor import QuantTensor + + targets: list[tuple[str, torch.nn.Module]] = [] + for name, module in model.named_modules(): + if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)): + continue + weight = module._parameters.get("weight") + if weight is None or not isinstance(weight.data, QuantTensor): + continue + targets.append((name, module)) + + if not targets: + return + + for name, module in targets: + qt: QuantTensor = module.weight.data # type: ignore[assignment] + if isinstance(module, torch.nn.Embedding): + new = QuantEmbeddingNbit.from_quant_tensor( + qt, + padding_idx=module.padding_idx, + dynamo=dynamo, + ) + else: + new = QuantLinearNbit.from_quant_tensor( + qt, + bias=module.bias.detach().clone() if module.bias is not None else None, + dynamo=dynamo, + ) + set_attr(model, name, new) + + # After all replacements, the original ``quantization_method`` / + # ``quantization_config`` set by ``OliveHfQuantizer`` no longer + # describes the runtime modules. Clear it so downstream passes + # (e.g. dtype casts) treat the model as plain torch. + if hasattr(model, "quantization_method"): + model.quantization_method = None diff --git a/olive/common/hf/utils.py b/olive/common/hf/utils.py index 1e8687c62e..51e76ce43b 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-initialize 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/hf/wrapper.py b/olive/common/hf/wrapper.py index 8bde44ef58..dd37dc3b80 100644 --- a/olive/common/hf/wrapper.py +++ b/olive/common/hf/wrapper.py @@ -120,7 +120,20 @@ class LayerWrapper: "opt": ["out_proj"], "qwen": ["c_proj"], } - MLP = {"default": "mlp", "lfm2": "feed_forward", "opt": ""} + # ``granitemoe``/``jamba`` keep their (MoE or dense) feed-forward block under a + # non-``mlp`` attribute: + # * ``GraniteMoeDecoderLayer.block_sparse_moe = GraniteMoeMoE(config)`` + # * ``Jamba{Attention,Mamba}DecoderLayer.feed_forward = JambaSparseMoeBlock(...)`` + # or ``JambaMLP(...)`` -- Jamba interleaves MoE and dense layers, so the dense + # ones simply resolve to a block without ``.experts``/``.router`` and + # ``get_experts()``/``get_router()`` return ``None``. + MLP = { + "default": "mlp", + "granitemoe": "block_sparse_moe", + "jamba": "feed_forward", + "lfm2": "feed_forward", + "opt": "", + } MLP_INPUTS = { "default": ["gate_proj", "up_proj"], "bloom": ["dense_h_to_4h"], @@ -138,6 +151,65 @@ class LayerWrapper: "opt": ["fc2"], "qwen": ["c_proj"], } + # MoE-block conventions. These are resolved relative to ``self.mlp`` + # (i.e., the layer's MLP attribute) because every modern HF MoE + # transformer block lives at ``layer.mlp``. + # + # ``EXPERTS`` is the experts sub-module: + # - For fused-3D MoEs (Mixtral, Qwen3-MoE, GPT-OSS) it owns 3D + # ``nn.Parameter`` tensors such as ``gate_up_proj`` of shape + # ``(num_experts, ...)``. + # - For ``ModuleList(Expert)`` MoEs (PhiMoE, DeepSeek-V3, classic + # Mixtral) it is the ``ModuleList`` whose children are the + # per-expert ``nn.Module`` blocks (each with their own + # ``nn.Linear``s). + # + # ``ROUTER`` is the routing module ("gate" in most, "router" in + # GPT-OSS). It is usually an ``nn.Linear`` (or a small custom module + # containing one) and should typically be kept in full precision. + EXPERTS = { + "default": "experts", + } + # + # Router attribute names verified against transformers 5.14.1: + # * ``gate`` -- qwen2_moe, qwen3_moe, mixtral, deepseek_v3, olmoe (the default) + # * ``router`` -- gpt_oss, phimoe, granitemoe, jamba. Note that Jamba's router is a + # bare ``nn.Linear`` (not a wrapped router module), so it would otherwise be swept + # into the ordinary 2D quantization walk; ``iter_quant_targets`` excludes every + # resolved router of an MoE layer by identity. + ROUTER = { + "default": "gate", + "gpt_oss": "router", + "granitemoe": "router", + "jamba": "router", + "phimoe": "router", + } + # ``SHARED_EXPERT_GATE`` is the sigmoid gate that scales a layer's always-active shared + # expert branch (``F.sigmoid(self.shared_expert_gate(x)) * shared_expert_output``), used + # alongside a normal top-k ``ROUTER`` by qwen2_moe, qwen3_5_moe, qwen3_next, and + # qwen3_omni_moe. Like the router, it directly controls how much the shared expert + # contributes and is a single-row ``nn.Linear`` (out_features=1) -- quantizing it to a + # handful of bits is both undesirable (routing-like signal) and disproportionately lossy + # (one output row). Resolved (when present) purely so ``iter_quant_targets`` can exclude + # it, the same way it excludes ``ROUTER``. + SHARED_EXPERT_GATE = { + "default": "shared_expert_gate", + } + # ``MAMBA`` is a decoder layer's state-space-model (SSM) sub-module, when present -- + # e.g. Jamba interleaves ``JambaMambaDecoderLayer`` (has ``.mamba``) with + # ``JambaAttentionDecoderLayer`` (has ``.self_attn`` instead). A Mamba block's + # ``nn.Linear`` projections (``in_proj``/``x_proj``/``dt_proj``/``out_proj``) feed a + # state-space recursion rather than a plain matmul, so they were never intentionally + # supported by the generic 2D quantization walk -- resolved (when present) purely so + # ``iter_quant_targets`` can exclude them, the same way it excludes routers. + # + # ``qwen3_5_moe``/``qwen3_5_moe_text`` interleave full attention with GatedDeltaNet + # linear-attention layers under ``.linear_attn`` instead of ``.mamba``. + MAMBA = { + "default": "mamba", + "qwen3_5_moe": "linear_attn", + "qwen3_5_moe_text": "linear_attn", + } def __init__(self, layer: nn.Module, model_type: str): # TODO(jambayk): use _layer and property to get the layer? @@ -157,13 +229,39 @@ def get_first_layer_norm(self, return_name: bool = True): def get_second_layer_norm(self, return_name: bool = True): return get_submodules(self.layer, self.SECOND_LAYER_NORM, self.model_type, return_name=return_name) - def get_attention_inputs(self, return_name: bool = True): + def get_attention_inputs(self, return_name: bool = True, partial_ok: bool = False): + """Return the attention input projections of this layer. + + Args: + return_name: Whether to also return the resolved module names. + partial_ok: When ``False`` (the default) every projection named in + ``ATTENTION_INPUTS`` must exist, otherwise an ``AttributeError`` is raised. + This keeps the returned list positional, which callers such as + ``olive.passes.pytorch.rotate`` rely on (they identify ``v_proj`` by index). + When ``True`` missing projections are dropped from the returned list, which + is only correct for callers that treat the list as an unordered set -- + e.g. QKV-group normalization, which needs to tolerate architectures with a + non-QKV attention (DeepSeek-V3's MLA exposes ``q_proj`` / + ``kv_a_proj_with_mqa`` / ``kv_b_proj``, so ``k_proj``/``v_proj`` are absent). + Those extra projections are still quantized by the generic ``nn.Linear`` + walk, and QKV normalization is a no-op for a group of fewer than two members. + + """ if self.attn is None: return ([], []) if return_name else [] attention_inputs, names = get_submodules( - self.attn, self.ATTENTION_INPUTS, self.model_type, return_name=True, return_name_prefix=f"{self.attn_name}." + self.attn, + self.ATTENTION_INPUTS, + self.model_type, + return_name=True, + return_name_prefix=f"{self.attn_name}.", + fail_on_not_found=not partial_ok, ) - if isinstance(attention_inputs[0], UnpackedQKV): + if partial_ok: + keep = [i for i, module in enumerate(attention_inputs) if module is not None] + attention_inputs = [attention_inputs[i] for i in keep] + names = [names[i] for i in keep] + if attention_inputs and isinstance(attention_inputs[0], UnpackedQKV): names = [f"{names[0]}.{part}" for part in ["q_proj", "k_proj", "v_proj"]] attention_inputs = [attention_inputs[0].q_proj, attention_inputs[0].k_proj, attention_inputs[0].v_proj] return attention_inputs if not return_name else (attention_inputs, names) @@ -179,15 +277,100 @@ def get_attention_outputs(self, return_name: bool = True): return_name_prefix=f"{self.attn_name}.", ) - def get_mlp_inputs(self, return_name: bool = True): - return get_submodules( - self.mlp, self.MLP_INPUTS, self.model_type, return_name=return_name, return_name_prefix=f"{self.mlp_name}." - ) + def get_mlp_inputs(self, return_name: bool = True, partial_ok: bool = False): + return self._get_mlp_projections(self.MLP_INPUTS, return_name, partial_ok) - def get_mlp_outputs(self, return_name: bool = True): - return get_submodules( - self.mlp, self.MLP_OUTPUTS, self.model_type, return_name=return_name, return_name_prefix=f"{self.mlp_name}." + def get_mlp_outputs(self, return_name: bool = True, partial_ok: bool = False): + return self._get_mlp_projections(self.MLP_OUTPUTS, return_name, partial_ok) + + def _get_mlp_projections(self, mapping: dict, return_name: bool, partial_ok: bool): + """Resolve the MLP projections named in ``mapping``. + + By default every mapped projection must exist. ``partial_ok=True`` drops missing + projections for MoE-aware callers that intentionally handle layers without a single + dense MLP path. + """ + modules, names = get_submodules( + self.mlp, + mapping, + self.model_type, + return_name=True, + return_name_prefix=f"{self.mlp_name}.", + fail_on_not_found=not partial_ok, + ) + if partial_ok: + keep = [i for i, module in enumerate(modules) if module is not None] + modules = [modules[i] for i in keep] + names = [names[i] for i in keep] + return modules if not return_name else (modules, names) + + def get_experts(self, return_name: bool = True): + """Return the experts sub-module of this layer (or ``None`` if not MoE). + + The experts sub-module is the parent of every per-expert weight + (fused 3D ``nn.Parameter``s, or a ``ModuleList`` of per-expert + ``nn.Module``s). The caller can use the returned module to (a) + collect ids of every ``nn.Module`` under the experts subtree + (for the ``moe=False`` skip set) and (b) iterate the + ``nn.Parameter``s to quantize when ``moe=True``. + + Returns ``None`` (and an empty name when ``return_name=True``) + for layers without an experts sub-module. + """ + if self.mlp is None: + return (None, "") if return_name else None + module = get_submodules(self.mlp, self.EXPERTS, self.model_type, return_name=False, fail_on_not_found=False) + if module is None: + return (None, "") if return_name else None + name = f"{self.mlp_name}.{self.EXPERTS.get(self.model_type, self.EXPERTS['default'])}" + return (module, name) if return_name else module + + def get_router(self, return_name: bool = True): + """Return the router sub-module of this layer (or ``None`` if not MoE). + + Routers are typically small modules (e.g., a single + ``nn.Linear``) and should usually be kept in full precision. + Olive does not quantize routers automatically — this accessor is + used by callers that want to skip the router via + ``modules_to_not_convert``. + """ + if self.mlp is None: + return (None, "") if return_name else None + module = get_submodules(self.mlp, self.ROUTER, self.model_type, return_name=False, fail_on_not_found=False) + if module is None: + return (None, "") if return_name else None + name = f"{self.mlp_name}.{self.ROUTER.get(self.model_type, self.ROUTER['default'])}" + return (module, name) if return_name else module + + def get_shared_expert_gate(self, return_name: bool = True): + """Return this layer's shared-expert gate sub-module (or ``None`` if not present). + + See ``SHARED_EXPERT_GATE``'s docstring for why this is excluded from quantization the + same way ``ROUTER`` is. + """ + if self.mlp is None: + return (None, "") if return_name else None + module = get_submodules( + self.mlp, self.SHARED_EXPERT_GATE, self.model_type, return_name=False, fail_on_not_found=False ) + if module is None: + return (None, "") if return_name else None + name = f"{self.mlp_name}.{self.SHARED_EXPERT_GATE.get(self.model_type, self.SHARED_EXPERT_GATE['default'])}" + return (module, name) if return_name else module + + def get_mamba(self, return_name: bool = True): + """Return this layer's Mamba/SSM sub-module (or ``None`` for a non-Mamba layer). + + Resolved relative to ``self.layer`` (not ``self.mlp``, unlike ``EXPERTS``/``ROUTER``): + a Mamba block is a sibling of the MLP/attention block, not nested inside either. + Used by ``iter_quant_targets`` to exclude the block's projections from the generic + 2D quantization walk -- see ``MAMBA``'s docstring for why. + """ + module = get_submodules(self.layer, self.MAMBA, self.model_type, return_name=False, fail_on_not_found=False) + if module is None: + return (None, "") if return_name else None + name = self.MAMBA.get(self.model_type, self.MAMBA["default"]) + return (module, name) if return_name else module class ModelWrapper: @@ -209,6 +392,10 @@ class ModelWrapper: "gptj": ["transformer.wte"], "opt": ["model.decoder.embed_tokens", "model.decoder.embed_positions"], "qwen": ["transformer.wte"], + # VL checkpoint: decoder lives under ``model.language_model`` alongside + # ``model.visual``, rather than directly under ``model`` -- see ``LAYERS`` below. + # Flat text-only ``qwen3_5_moe`` configs are routed to ``qwen3_5_moe_text`` instead. + "qwen3_5_moe": ["model.language_model.embed_tokens"], } # in newer transformers versions, there is one rotary embedding per model ROTARY_EMBEDDING = { @@ -216,6 +403,7 @@ class ModelWrapper: "falcon": "transformer.rotary_emb", "gpt_neox": "gpt_neox.rotary_emb", "qwen": "transformer.rotary_emb", + "qwen3_5_moe": "model.language_model.rotary_emb", } LM_HEAD = {"default": "lm_head"} PRE_HEAD_LAYERNORM = { @@ -223,6 +411,7 @@ class ModelWrapper: "gpt2": "transformer.ln_f", "lfm2": "model.embedding_norm", "qwen": "transformer.ln_f", + "qwen3_5_moe": "model.language_model.norm", } LAYERS = { "default": "model.layers", @@ -233,11 +422,28 @@ class ModelWrapper: "gptj": "transformer.h", "opt": "model.decoder.layers", "qwen": "transformer.h", + # ``Qwen3_5MoeForConditionalGeneration`` (VL) nests the decoder under + # ``model.language_model`` next to ``model.visual`` (the vision tower). Flat text-only + # ``qwen3_5_moe`` checkpoints report the same ``model_type`` but use the default + # ``model.layers`` layout -- they are routed to ``qwen3_5_moe_text`` by + # ``_resolve_model_type`` (see ``TEXT_ONLY_MODEL_TYPES``), so this entry only applies + # to configs that actually carry a ``vision_config``. + "qwen3_5_moe": "model.language_model.layers", } + # Some ``model_type``s are shared by a composite (vision-language) checkpoint and a flat + # text-only checkpoint. ``qwen3_5_moe`` is one: ``Qwen3_5MoeForConditionalGeneration`` + # nests the decoder under ``model.language_model`` (next to ``model.visual``), while + # text-only ``Qwen3_5MoeForCausalLM`` checkpoints keep the flat ``model.layers`` layout + # even though their config still reports ``model_type == "qwen3_5_moe"``. Mobius + # disambiguates the two exactly this way -- ``config.vision_config is None`` means + # text-only -- so a text-only config is normalized to the ``*_text`` model_type, which has + # no entries in the module-path mappings above and therefore uses the flat "default" + # paths. Maps ``VL model_type -> text-only model_type``. + TEXT_ONLY_MODEL_TYPES = {"qwen3_5_moe": "qwen3_5_moe_text"} def __init__(self, config: Union[PretrainedConfig, dict]): self.config = config if isinstance(config, PretrainedConfig) else PretrainedConfig.from_dict(config) - self.model_type = getattr(self.config, "model_type", None) + self.model_type = self._resolve_model_type(self.config) # model attributes (using unified aliases from defaults.yaml) self.hidden_size = resolve_alias(self.config, "hidden_size") @@ -251,6 +457,25 @@ def __init__(self, config: Union[PretrainedConfig, dict]): self._model = None self._layer_wrappers = None + @classmethod + def _resolve_model_type(cls, config: PretrainedConfig) -> Union[str, None]: + """Return the model_type used to look up module paths in the mappings above. + + Normalizes a flat text-only checkpoint whose ``model_type`` is shared with a composite + VL architecture to the corresponding text-only ``model_type`` -- see + ``TEXT_ONLY_MODEL_TYPES``. + """ + model_type = getattr(config, "model_type", None) + if model_type in cls.TEXT_ONLY_MODEL_TYPES and getattr(config, "vision_config", None) is None: + text_only_model_type = cls.TEXT_ONLY_MODEL_TYPES[model_type] + logger.debug( + "Config for model_type %r has no vision_config; treating it as the text-only %r layout.", + model_type, + text_only_model_type, + ) + return text_only_model_type + return model_type + @property def model(self) -> "PreTrainedModel": if self._model is None: diff --git a/olive/common/mobius_utils.py b/olive/common/mobius_utils.py index 40aa2957be..8c4e5c13b0 100644 --- a/olive/common/mobius_utils.py +++ b/olive/common/mobius_utils.py @@ -8,7 +8,7 @@ ``decoder`` / ``vision_encoder`` / ``embedding``), how each maps back to a submodule, and the role of each component. This adapter lets Olive consume that plan without re-implementing architecture-specific logic. -``mobius-ai`` is imported lazily so Olive keeps working when it is not installed; only the code paths +``mobius-onnx`` is imported lazily so Olive keeps working when it is not installed; only the code paths that actually need a component plan for a Hugging Face model require it. """ @@ -100,15 +100,15 @@ def inspect_components( (no separable components). Raises: - ImportError: If ``mobius-ai`` is not installed. + ImportError: If ``mobius-onnx`` is not installed. """ try: import mobius except ImportError as exc: raise ImportError( - "mobius-ai is required to resolve model components for a Hugging Face model. " - "Install with: pip install mobius-ai" + "mobius-onnx is required to resolve model components for a Hugging Face model. " + "Install with: pip install mobius-onnx" ) from exc raw_components = mobius.inspect_components( diff --git a/olive/common/onnx_io.py b/olive/common/onnx_io.py index 5a47c9058a..5c10668a3b 100644 --- a/olive/common/onnx_io.py +++ b/olive/common/onnx_io.py @@ -2,7 +2,30 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +import json import re +from pathlib import Path +from typing import Optional, Union + + +def get_genai_decoder_config(model_path: Union[str, Path]) -> Optional[dict]: + """Return the decoder section of the genai_config.json saved next to an ONNX model. + + onnxruntime-genai exports some KV cache dimensions symbolically (for example ``kv_cache_dim``) + so that a single ONNX model serves both plain and quantized KV caches. The concrete values are + only available in the genai_config.json written alongside the model. + + :param model_path: Path to the ONNX model file or to the directory containing it. + :return: The ``model.decoder`` section of genai_config.json, or None if there is no such file. + """ + model_path = Path(model_path) + model_dir = model_path.parent if model_path.suffix else model_path + genai_config_path = model_dir / "genai_config.json" + if not genai_config_path.is_file(): + return None + + with genai_config_path.open() as config_file: + return json.load(config_file).get("model", {}).get("decoder") def get_io_config(model_path: str) -> dict: diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py index c1e4c5065b..0f45e33735 100644 --- a/olive/common/quant/hf_utils.py +++ b/olive/common/quant/hf_utils.py @@ -2,23 +2,33 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +# pylint: disable=protected-access from __future__ import annotations +import math from dataclasses import dataclass from typing import TYPE_CHECKING, Callable +import torch import torch.nn as nn from transformers.quantizers.base import HfQuantizer from transformers.utils.quantization_config import QuantizationConfigMixin +from olive.common.quant.patterns import _compiled, is_regex_pattern, match_override +from olive.common.quant.state_dict import ( + bind_quant_tensor_to_buffers, + buffer_names, + install_quant_tensor_param, + refresh_quant_tensor_refs, +) +from olive.common.quant.tensor import QuantTensor +from olive.common.quant.utils import WeightQuantizer from olive.common.utils import StrEnumBase if TYPE_CHECKING: from tqdm.auto import tqdm from transformers import PreTrainedModel - from olive.common.quant.nn import QuantModule - # older transformers expects a StrEnum and accesses .value class OliveHfQuantizationMethod(StrEnumBase): @@ -58,25 +68,47 @@ class OliveHfQuantizationConfig(QuantizationConfigMixin): -1 = per-channel, >0 = groupwise. lm_head: Whether to quantize the language model head. embeds: Whether to quantize the input embeddings. - modules_to_not_convert : List of module names to exclude from quantization. + moe: Whether to quantize MoE expert modules / parameters. + When ``False`` (default), every ``nn.Module`` under each + experts subtree returned by ``LayerWrapper.get_experts()`` + is added to the skip set — this both leaves fused-3D + experts alone *and* fixes the previous silent quantization + of per-expert ``nn.Linear``s in ``ModuleList(Expert)`` + blocks (Mixtral, PhiMoE, Qwen2/3-MoE). + quantize_vision: Whether to quantize a composite vision-language + model's vision tower in this pass. When ``False`` (default), + the vision tower is left in full precision -- the typical + Olive pipeline quantizes it separately downstream (e.g. on + the ONNX side), so leaving it untouched here avoids + double-quantizing it. Set to ``True`` to quantize the vision + tower here too (e.g. when this pass is the only quantization + step for the model). + modules_to_not_convert: List of module name patterns to exclude + from quantization. Plain strings use **substring** matching + (preserving HF semantics); entries prefixed with ``re:`` use + ``re.fullmatch``. overrides: Per-module overrides for quantization parameters. + Keys use **literal equality** matching by default; entries + prefixed with ``re:`` use ``re.fullmatch``. Among matching + keys, the first matching key in insertion (config) order + wins. """ - # pylint: disable - def __init__( + def __init__( # pylint: disable=super-init-not-called self, bits: int, symmetric: bool, group_size: int, lm_head: bool = False, embeds: bool = False, + moe: bool = False, + quantize_vision: bool = False, modules_to_not_convert: list | None = None, overrides: dict | None = None, tie_word_embeddings: bool = False, **kwargs, ): - # pylint: disable=W0231 self.quant_method = OliveHfQuantizationMethod.OLIVE self.bits = bits @@ -84,6 +116,8 @@ def __init__( self.group_size = group_size self.lm_head = lm_head self.embeds = embeds + self.moe = moe + self.quantize_vision = quantize_vision self.modules_to_not_convert = modules_to_not_convert self.overrides = { module_name: OliveHfQuantizationOverrideConfig(**override) @@ -97,6 +131,13 @@ def post_init(self): if self.bits not in [2, 4, 8]: raise ValueError(f"Only 2-bit, 4-bit and 8-bit quantization supported, got {self.bits}") + # Validate all `re:` patterns eagerly (at construction time) rather than lazily on + # first match, so a bad/unsafe regex surfaces immediately even if it never happens + # to match anything during this run. + for pattern in list(self.overrides.keys()) + list(self.modules_to_not_convert or []): + if is_regex_pattern(pattern): + _compiled(pattern) + def to_dict(self) -> dict: """Serialize this instance to a Python dictionary.""" output = super().to_dict() @@ -104,10 +145,14 @@ def to_dict(self) -> dict: overrides = {} for module_name, override in self.overrides.items(): # remove None or default values from the override - cleaned_override = {k: v for k, v in override.__dict__.items() if v is not None and v != output[k]} + cleaned_override = {k: v for k, v in override.__dict__.items() if v is not None and v != output.get(k)} if cleaned_override: overrides[module_name] = cleaned_override - output["overrides"] = sort_layers_by_name(overrides) or None + # NOTE: do NOT sort by layer name here. ``match_override`` resolves + # overlapping matches by first-match-wins in insertion (config) + # order, so re-ordering the dict on serialization would silently + # flip the winner for overlapping patterns on save -> reload. + output["overrides"] = overrides or None else: output["overrides"] = None return output @@ -127,7 +172,9 @@ def get_qlinear_init_args(self, module_name: str) -> dict: "symmetric": self.symmetric, "group_size": self.group_size, } - if override := self.overrides.get(module_name): + best = match_override(module_name, list(self.overrides.keys())) if self.overrides else None + if best is not None: + override = self.overrides[best] init_args.update({k: v for k, v in override.__dict__.items() if v is not None}) return init_args @@ -142,66 +189,82 @@ def sort_key(name: str) -> tuple: class OliveHfQuantizer(HfQuantizer): - """Olive quantizer.""" + """Olive quantizer. + + Layout (see ``olive/common/quant/state_dict.py``): + + * Each quantized weight is installed as ``nn.Parameter(QuantTensor)`` + on the original host module (``nn.Linear``, ``nn.Embedding``, or an + experts module that owns a fused-3D parameter). + * Sibling buffers ``_qweight`` / ``_scales`` / ``_qzeros`` + alias the QuantTensor's inner tensors. These are the only things + written to safetensors; HF's loader fills them via normal dotted + paths. After load we re-bind the QuantTensor inner refs to point + at the freshly-loaded buffer storage. + """ # only support load and inference, no on-the-fly quantization requires_calibration = True - modules_to_not_convert: list[str] | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Populated (best-effort) in ``_process_model_before_weight_loading`` from the + # ``checkpoint_files`` HF passes to ``preprocess_model``. When available, this lets + # ``_process_model_after_weight_loading`` determine ``is_placeholder`` by an exact + # key-membership check against the checkpoint instead of the weaker buffer-identity + # heuristic (identity only detects loaders that *replace* the buffer object, e.g. via + # ``setattr``; it would miss an in-place ``.copy_()``-based loader). ``None`` means + # "couldn't determine the checkpoint's key set" (e.g. no safetensors files), in which + # case the identity heuristic remains the fallback. + self._checkpoint_keys: set[str] | None = None def _process_model_before_weight_loading( - self, model: PreTrainedModel, keep_in_fp32_modules: list[str] | None = None, **kwargs + self, + model: PreTrainedModel, + keep_in_fp32_modules: list[str] | None = None, + checkpoint_files: list[str] | None = None, + **kwargs, ): - from olive.common.quant.nn import QuantEmbedding, QuantLinear - - ids_to_skip = [] - if not self.quantization_config.lm_head: - ids_to_skip.append(id(model.get_output_embeddings())) - if not self.quantization_config.embeds: - ids_to_skip.append(id(model.get_input_embeddings())) - self.modules_to_not_convert = ( - [name for name, module in model.named_modules() if id(module) in ids_to_skip] if ids_to_skip else [] - ) - if self.quantization_config.modules_to_not_convert: - self.modules_to_not_convert.extend(self.quantization_config.modules_to_not_convert) - if keep_in_fp32_modules: - self.modules_to_not_convert.extend(keep_in_fp32_modules) + from olive.common.quant.selection import iter_quant_targets - def should_quantize(module: nn.Module, name: str) -> bool: - """Check if a module should be quantized.""" - return isinstance(module, (nn.Linear, nn.Embedding)) and not any( - key in name for key in self.modules_to_not_convert - ) + self._checkpoint_keys = _read_checkpoint_keys(checkpoint_files) - def create_quantized_module(module: nn.Linear | nn.Embedding, name: str) -> QuantLinear | QuantEmbedding: - """Create a quantized version of a module.""" - common_kwargs = { - **self.quantization_config.get_qlinear_init_args(name), - "device": module.weight.device, - "dtype": module.weight.dtype, - } - if isinstance(module, nn.Embedding): - return QuantEmbedding( - num_embeddings=module.num_embeddings, - embedding_dim=module.embedding_dim, - padding_idx=module.padding_idx, - **common_kwargs, - ) - return QuantLinear( - in_features=module.in_features, - out_features=module.out_features, - bias=module.bias is not None, - **common_kwargs, + skip_patterns: list[str] = [] + if self.quantization_config.modules_to_not_convert: + skip_patterns.extend(self.quantization_config.modules_to_not_convert) + if keep_in_fp32_modules: + skip_patterns.extend(keep_in_fp32_modules) + + for module, pname, full_name in iter_quant_targets( + model, + quantize_lm_head=self.quantization_config.lm_head, + quantize_embeds=self.quantization_config.embeds, + quantize_moe=self.quantization_config.moe, + quantize_vision=getattr(self.quantization_config, "quantize_vision", False), + skip_patterns=skip_patterns, + ): + qargs = self.quantization_config.get_qlinear_init_args(full_name) + param = module._parameters[pname] + qt = _build_placeholder_quant_tensor( + shape=tuple(param.shape), + bits=qargs["bits"], + symmetric=qargs["symmetric"], + group_size=qargs["group_size"], + dtype=param.dtype, + device=param.device, ) - - replace_matching_submodules(model, should_quantize, create_quantized_module) + install_quant_tensor_param(module, pname, qt) if self.quantization_config.tie_word_embeddings: # doing first time so that the weight load doesn't complain about missing weights tie_quant_word_embeddings(model) def _process_model_after_weight_loading(self, model: PreTrainedModel, **kwargs): + # HF's loader assigns freshly-loaded buffer tensors in place, so + # re-bind every QuantTensor parameter to point at the current + # ``_qweight`` / ``_scales`` / ``_qzeros`` buffer storages. + refresh_quant_tensor_refs(model, checkpoint_keys=self._checkpoint_keys) if self.quantization_config.tie_word_embeddings: - # doing again to ensure buffers are tied after loading weights tie_quant_word_embeddings(model) return model @@ -210,11 +273,83 @@ def is_serializable(self, safe_serialization=None) -> bool: @property def is_trainable(self) -> bool: - # TODO(jambayk): investigate what this means (peft, scale+bias, etc.?) - # need to support peft, scale+bias comes for free since everything is in torch return False +def _read_checkpoint_keys(checkpoint_files: list[str] | None) -> set[str] | None: + """Best-effort read of every tensor key present across ``checkpoint_files``. + + Only ``.safetensors`` shards are supported: their header (the tensor-name -> + metadata map) can be read without materializing any tensor data, so this is cheap + even for large checkpoints. Returns ``None`` (meaning "unknown, don't use for + authoritative ``is_placeholder`` decisions") when ``checkpoint_files`` is empty/None + or contains any non-safetensors file (e.g. a legacy pickled ``.bin`` shard), since + there is no equivalently cheap, safe way to list keys for those formats. + """ + if not checkpoint_files: + return None + if any(not str(f).endswith(".safetensors") for f in checkpoint_files): + return None + + try: + from safetensors import safe_open + except ImportError: + return None + + keys: set[str] = set() + try: + for checkpoint_file in checkpoint_files: + with safe_open(checkpoint_file, framework="pt") as f: + keys.update(f.keys()) + except Exception: # pylint: disable=broad-except + # Be conservative: any failure to read falls back to the identity heuristic + # rather than risking an incomplete/incorrect key set being treated as authoritative. + return None + return keys + + +def _build_placeholder_quant_tensor( + *, + shape: tuple[int, ...], + bits: int, + symmetric: bool, + group_size: int, + dtype: torch.dtype, + device: torch.device, +) -> QuantTensor: + """Build a zero-filled ``QuantTensor`` with the correct buffer shapes. + + This is used as a placeholder before ``from_pretrained`` fills in the + real values; the buffer shapes must match what was written at save + time so HF's loader assigns into them without complaint. + """ + quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False) + packing_factor = 8 // bits + qparam_shape = quantizer.get_qparam_shape(shape) + qweight_shape = (*shape[:-1], math.ceil(shape[-1] / packing_factor)) + + qweight = torch.zeros(qweight_shape, dtype=torch.uint8, device=device) + scales = torch.zeros(qparam_shape, dtype=dtype, device=device) + qzeros: torch.Tensor | None + if symmetric: + qzeros = None + else: + qz_shape = (*qparam_shape[:-1], math.ceil(qparam_shape[-1] / packing_factor)) + qzeros = torch.zeros(qz_shape, dtype=torch.uint8, device=device) + + return QuantTensor.from_packed( + qweight=qweight, + scales=scales, + qzeros=qzeros, + bits=bits, + group_size=group_size, + symmetric=symmetric, + shape=shape, + dtype=dtype, + is_placeholder=True, + ) + + def replace_matching_submodules( module: nn.Module, condition: Callable[[nn.Module, str], bool], @@ -265,56 +400,60 @@ def replace_matching_submodules( return module -def _repoint_buffer(src: nn.Module, dst: nn.Module, name: str): - """Repoint a buffer from src to dst. +def tie_quant_word_embeddings(model: PreTrainedModel) -> None: + """Tie the input and output embeddings when both share a quantized weight. - Args: - src: Source module. - dst: Destination module. - name: Name of the buffer to repoint. + Both modules' ``weight`` ``nn.Parameter`` is set to the **same** + ``nn.Parameter(QuantTensor)`` object, and the underlying + ``weight_qweight`` / ``weight_scales`` / ``weight_qzeros`` buffers + are tied (aliased to the input embedding's buffers). This preserves + the standard HF tied-weights semantics for the quantized layout. + Tying is a no-op unless **both** the input and output embeddings + are already backed by compatible ``QuantTensor`` weights with + matching shape and dtype. """ - src_buf = getattr(src, name, None) - - # ensure both are None or both exist - if src_buf is None: - assert getattr(dst, name, None) is None, f"Output embedding has {name} but input does not." + src = model.get_input_embeddings() + dst = model.get_output_embeddings() + if src is None or dst is None: return - # ensure both have the buffer shapes and types match - dst_buf = getattr(dst, name, None) - assert src_buf.shape == dst_buf.shape, ( - f"Cannot tie embeddings: input embedding {name} shape {src_buf.shape} " - f"does not match output embedding shape {dst_buf.shape}." - ) - assert src_buf.dtype == dst_buf.dtype, ( - f"Cannot tie embeddings: input embedding {name} dtype {src_buf.dtype} " - f"does not match output embedding dtype {dst_buf.dtype}." - ) - - # tie the buffers - # pylint: disable=W0212 - dst._buffers[name] = src_buf - dst._non_persistent_buffers_set.add(name) - - -def tie_quant_modules(src: QuantModule, dst: QuantModule): - """Tie the quantization buffers of two QuantModules. - - Args: - src: Source QuantModule. - dst: Destination QuantModule. - - """ - for name in ["qweight", "scales", "qzeros"]: - _repoint_buffer(src, dst, name) - - -def tie_quant_word_embeddings(model: PreTrainedModel): - """Tie the word embeddings and output embeddings if they have the same shape. - - Args: - model: The HuggingFace model to tie embeddings for. + src_param = src._parameters.get("weight") + dst_param = dst._parameters.get("weight") + if ( + src_param is None + or dst_param is None + or not isinstance(src_param.data, QuantTensor) + or not isinstance(dst_param.data, QuantTensor) + ): + return + if src_param.shape != dst_param.shape or src_param.dtype != dst_param.dtype: + return - """ - tie_quant_modules(model.get_input_embeddings(), model.get_output_embeddings()) + qname, sname, zname = buffer_names("weight") + # Tie buffers, marking the destination copy non-persistent so + # ``state_dict()`` / safetensors only emit one set of keys for the + # shared tensors (avoids duplicate qweight/scales on disk). + for n in (qname, sname, zname): + src_buf = src._buffers.get(n) + if src_buf is None: + continue + if n in dst._buffers: + dst._non_persistent_buffers_set.add(n) + dst._buffers[n] = src_buf + + # tie the QuantTensor parameter itself (same Python Parameter object, + # so both modules see the same .data and the same inner tensors). + dst._parameters["weight"] = src_param + + # Re-sync the (now shared) QuantTensor's own inner refs to ``src``'s buffers. + # + # ``refresh_quant_tensor_refs`` is the single source of truth for "which buffers does a + # QuantTensor reference" (it dedupes by object identity and repairs every alias), and + # this call uses that exact same primitive rather than introducing a second rule. It + # exists so that tying is self-consistent *by construction* at any call site: after this + # function returns, ``src._buffers == dst._buffers == src_param.{qweight,scales,qzeros}`` + # regardless of whether the caller happened to run a refresh first. Without it, tying + # would only alias the two ``_buffers`` dicts and could leave the shared parameter + # pointing at buffer objects that neither module hosts any more. + bind_quant_tensor_to_buffers(src_param, src, "weight") diff --git a/olive/common/quant/nn.py b/olive/common/quant/nn.py deleted file mode 100644 index 78c804d052..0000000000 --- a/olive/common/quant/nn.py +++ /dev/null @@ -1,550 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -from __future__ import annotations - -import math -from abc import abstractmethod - -import torch -import torch.nn as nn - -from olive.common.quant.utils import WeightQuantizer, pack_to_uint8, unpack_from_uint8 - - -class QuantModule(nn.Module): - """A base class for quantized modules. - - Only supports 2D weight tensors for now. Quantization axis is assumed to be the last dimension. - - 4-bit and 8-bit quantization - - Symmetric and asymmetric quantization - - Per-channel and groupwise quantization - - For blockwise quantization, it has the following restrictions: - - The size of the last dimension must be divisible by the group size. This is to avoid needing to pad the weights. - Padding is easy but it complicates compatibility between contrib ops and QDQ representations, and weight tieing. - - group size must be >= 16 and power of 2. For compatibility with the contrib ops. - # TODO(jambayk): extend to support more general cases if needed. - """ - - def __init__( - self, - rows: int, - cols: int, - bits: int = 4, - symmetric: bool = True, - group_size: int = -1, - device: torch.device | None = None, - dtype: torch.dtype = torch.float32, - ): - """Initialize QuantLinear layer. - - Args: - rows: Number of rows in the weight matrix - cols: Number of columns in the weight matrix - bits: Number of bits for quantization (4 or 8) - symmetric: Whether to use symmetric quantization - group_size: Quantization group size (-1: per-channel, >0: groupwise) - device: Device to place tensors on - dtype: Data type for scales and bias - - """ - super().__init__() - if bits not in [2, 4, 8]: - raise ValueError(f"Only 2-bit, 4-bit and 8-bit quantization supported, got {bits}") - if group_size != -1 and (group_size < 16 or (group_size & (group_size - 1)) != 0): - raise ValueError("For blockwise quantization, group_size must be >= 16 and power of 2") - if group_size != -1 and cols % group_size != 0: - raise ValueError( - f"For blockwise quantization, cols ({cols}) must be divisible by group_size ({group_size})" - ) - - self.rows = rows - self.cols = cols - self.quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False) - self.device = device - self.dtype = dtype - - packing_factor = 8 // bits - - # using the same layout and packing as auto-gptq - # TODO(jambayk): consider other packing schemes - self.register_buffer( - "qweight", - torch.zeros( - # rows X cols, packed as uint8 along last dim - (self.rows, math.ceil(self.cols / packing_factor)), - dtype=torch.uint8, - device=device, - ), - ) - scale_shape = self.quantizer.get_qparam_shape((self.rows, self.cols)) - self.register_buffer( - "scales", - torch.zeros(scale_shape, dtype=dtype, device=device), - ) - if symmetric: - self.qzeros = None - else: - self.register_buffer( - "qzeros", - torch.zeros( - (scale_shape[0], math.ceil(scale_shape[1] / packing_factor)), - dtype=torch.uint8, - device=device, - ), - ) - - @classmethod - def from_tensors( - cls, - weight: torch.Tensor, - bits: int = 4, - symmetric: bool = True, - group_size: int = -1, - scales: torch.Tensor | None = None, - zero_points: torch.Tensor | None = None, - quantized: bool = False, - **kwargs, - ) -> QuantModule: - """Create a QuantLinear layer from an existing nn.Linear layer. - - Args: - weight: The weight tensor. Expected to be 2D (unsigned, in range [0, 2^bits - 1] if quantized is True). - bits: Number of bits for quantization (4 or 8) - symmetric: Whether to use symmetric quantization - group_size: Quantization group size (-1: per-channel, >0: groupwise) - scales: Optional precomputed scales for quantization - zero_points: Optional precomputed zero points for quantization (unsigned, in range [0, 2^bits - 1]). - quantized: Whether the provided weight is already quantized - kwargs: Additional keyword arguments to pass to the QuantModule constructor - - Returns: - A QuantModule instance with quantized weight - - """ - # pylint: disable=W0201 - if quantized: - if scales is None or zero_points is None: - raise ValueError("scales and/or zero_points missing for quantized weight") - qweight = weight - else: - quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False) - - # compute quantization parameters if not provided - if scales is None: - scales, zero_points = quantizer.find_qparams(weight) - else: - scales = scales.to(weight.device).to(weight.dtype) - zero_points = zero_points.to(weight.device).to(torch.int32) - - # quantize weights - qweight = quantizer.quantize(weight, scales, zero_points) - - qmodule = cls( - bits=bits, - symmetric=symmetric, - group_size=group_size, - device=qweight.device, - dtype=scales.dtype, - **kwargs, - ) - qmodule.qweight = pack_to_uint8(qweight, bits).contiguous() - scale_shape = qmodule.quantizer.get_qparam_shape(qweight.shape) - qmodule.scales = scales.reshape(scale_shape).contiguous() - - # enforce symmetric quantization constraints - if symmetric and not torch.all(zero_points == qmodule.quantizer.midq): - raise ValueError("Zero points must be equal to midq for symmetric quantization") - - if not symmetric: - qmodule.qzeros = pack_to_uint8(zero_points.reshape(scale_shape), bits).contiguous() - - return qmodule - - @torch.no_grad() - def unpack_and_dequantize( - self, qweight: torch.Tensor, scales: torch.Tensor, zero_points: torch.Tensor | None - ) -> torch.Tensor: - """Unpack and dequantize the given quantized weight tensor. - - Returns: - The dequantized weight tensor. - - """ - qweight = unpack_from_uint8(qweight, self.quantizer.bits, (qweight.shape[0], self.cols)) - if zero_points is not None: - zero_points = unpack_from_uint8(zero_points, self.quantizer.bits, scales.shape) - else: - zero_points = torch.full_like( - scales, - self.quantizer.midq, - dtype=torch.int32, - ) - return self.quantizer.dequantize(qweight, scales, zero_points) - - @abstractmethod - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass of the quantized module. - - Args: - x: The input tensor. - - Returns: - The output tensor. - - """ - - -class QuantLinear(QuantModule): - """Quantized Linear layer.""" - - def __init__( - self, - in_features: int, - out_features: int, - bits: int = 4, - symmetric: bool = True, - group_size: int = -1, - bias: bool = True, - device: torch.device | None = None, - dtype: torch.dtype = torch.float32, - ): - """Initialize QuantLinear layer. - - Args: - in_features: Size of input features - out_features: Size of output features - bits: Number of bits for quantization (4 or 8) - symmetric: Whether to use symmetric quantization - group_size: Quantization group size (-1: per-channel, >0: groupwise) - bias: Whether to include bias - device: Device to place tensors on - dtype: Data type for scales and bias - - """ - super().__init__( - rows=out_features, - cols=in_features, - bits=bits, - symmetric=symmetric, - group_size=group_size, - device=device, - dtype=dtype, - ) - - if bias: - self.register_buffer( - "bias", - torch.zeros(out_features, dtype=dtype, device=device), - ) - else: - self.bias = None - - @classmethod - def from_module( - cls, - linear: nn.Linear, - bits: int = 4, - symmetric: bool = True, - group_size: int = -1, - scales: torch.Tensor | None = None, - zero_points: torch.Tensor | None = None, - ) -> QuantLinear: - """Create a QuantLinear layer from an existing nn.Linear layer. - - Args: - linear: The nn.Linear layer to convert - bits: Number of bits for quantization (4 or 8) - symmetric: Whether to use symmetric quantization - group_size: Quantization group size (-1: per-channel, 0: per-tensor, >0: groupwise) - scales: Optional precomputed scales for quantization - zero_points: Optional precomputed zero points for quantization (unsigned, in range [0, 2^bits - 1]). - - Returns: - A QuantLinear instance with quantized weights and scales - - """ - qlinear = cls.from_tensors( - in_features=linear.in_features, - out_features=linear.out_features, - weight=linear.weight, - bits=bits, - symmetric=symmetric, - group_size=group_size, - scales=scales, - zero_points=zero_points, - bias=linear.bias is not None, - ) - if linear.bias is not None: - qlinear.bias = linear.bias.clone() - return qlinear - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass for the quantized linear layer. - - Args: - x: Input tensor of shape (..., in_features) - - Returns: - Output tensor of shape (..., out_features) - - """ - assert x.shape[-1] == self.cols, f"Input shape {x.shape} does not match in_features {self.cols}" - - if torch.onnx.is_in_onnx_export(): - out = QuantLinearFunction.apply( - x, - self.qweight.reshape([*self.scales.shape, self.qweight.shape[-1] // self.scales.shape[-1]]), - self.scales, - self.qzeros, - self.quantizer.bits, - self.quantizer.group_size if self.quantizer.group_size > 0 else self.cols, - self.cols, - self.rows, - ) - return out + self.bias if self.bias is not None else out - - x_dtype = x.dtype - - # unpack weights and zero points - weight = self.unpack_and_dequantize(self.qweight, self.scales, self.qzeros) - return nn.functional.linear(x, weight, bias=self.bias).to(x_dtype) # pylint: disable=not-callable - - def extra_repr(self) -> str: - return ( - f"in_features={self.cols}, out_features={self.rows}, bits={self.quantizer.bits}," - f" symmetric={self.quantizer.symmetric}, group_size={self.quantizer.group_size}," - f" bias={self.bias is not None}" - ) - - -# this is required for torchscript onnx export -class QuantLinearFunction(torch.autograd.Function): - # pylint: disable=W0223,W0221 - @staticmethod - def symbolic(g, x, qweight, scales, qzeros, bits, group_size, in_features, out_features): - tensor_args = [x, qweight, scales] - if qzeros is not None: - tensor_args.append(qzeros) - attrs = { - "K_i": in_features, - "N_i": out_features, - "bits_i": bits, - "block_size_i": group_size, - "accuracy_level_i": 4, - } - - output = g.op( - "com.microsoft::MatMulNBits", - *tensor_args, - outputs=1, - **attrs, - ) - input_shape = x.type().varyingSizes() - if input_shape is not None and hasattr(x.type(), "with_sizes"): - output_type = x.type().with_sizes([*input_shape[:-1], out_features]) - output.setType(output_type) - - return output - - @staticmethod - def forward( - ctx, - x: torch.Tensor, - qweight: torch.Tensor, - scales: torch.Tensor, - qzeros: torch.Tensor, - bits: int, - group_size: int, - in_features: int, - out_features: int, - ): - # there is no clean way to differentiate between torchscript and dynamo export - # using FakeTensor as a proxy for dynamo export - if hasattr(torch.onnx, "ops") and isinstance(x, torch._subclasses.FakeTensor): # pylint: disable=W0212 - tensor_args = [x, qweight, scales] - if qzeros is not None: - tensor_args.append(qzeros) - attrs = { - "K": in_features, - "N": out_features, - "bits": bits, - "block_size": group_size, - "accuracy_level": 4, - } - return torch.onnx.ops.symbolic( - "com.microsoft::MatMulNBits", - tensor_args, - attrs=attrs, - dtype=x.dtype, - shape=[*x.shape[:-1], out_features], - version=1, - ) - return torch.zeros([*x.shape[:-1], out_features], dtype=x.dtype, device=x.device) - - -class QuantEmbedding(QuantModule): - """Quantized Embedding layer.""" - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - bits: int = 4, - symmetric: bool = True, - group_size: int = -1, - device: torch.device | None = None, - dtype: torch.dtype = torch.float32, - padding_idx: int | None = None, - ): - """Initialize QuantEmbedding layer. - - Args: - num_embeddings: Number of embeddings - embedding_dim: Dimension of each embedding - padding_idx: Index of the padding token - bits: Number of bits for quantization (4 or 8) - symmetric: Whether to use symmetric quantization - group_size: Quantization group size (-1: per-channel, >0: groupwise) - device: Device to place tensors on - dtype: Data type for scales - - """ - super().__init__( - rows=num_embeddings, - cols=embedding_dim, - bits=bits, - symmetric=symmetric, - group_size=group_size, - device=device, - dtype=dtype, - ) - self.padding_idx = padding_idx - - @classmethod - def from_module( - cls, - embedding: nn.Embedding, - bits: int = 4, - symmetric: bool = True, - group_size: int = -1, - scales: torch.Tensor | None = None, - zero_points: torch.Tensor | None = None, - ) -> QuantEmbedding: - """Create a QuantEmbedding layer from an existing nn.Embedding layer. - - Args: - embedding: The nn.Embedding layer to convert - bits: Number of bits for quantization (4 or 8) - symmetric: Whether to use symmetric quantization - group_size: Quantization group size (-1: per-channel, >0: groupwise) - scales: Optional precomputed scales for quantization - zero_points: Optional precomputed zero points for quantization (unsigned, in range [0, 2^bits - 1]). - - Returns: - A QuantEmbedding instance with quantized weights and scales - - """ - return cls.from_tensors( - num_embeddings=embedding.num_embeddings, - embedding_dim=embedding.embedding_dim, - weight=embedding.weight, - bits=bits, - symmetric=symmetric, - group_size=group_size, - scales=scales, - zero_points=zero_points, - padding_idx=embedding.padding_idx, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass for the quantized embedding layer. - - Args: - x: Input tensor of shape (...,) - - """ - if torch.onnx.is_in_onnx_export(): - return QuantEmbeddingFunction.apply( - x, - self.qweight, - self.scales, - self.qzeros, - self.quantizer.bits, - self.quantizer.group_size if self.quantizer.group_size > 0 else self.cols, - self.cols, - ) - - # reshape input to 1D so that look up results are 2D - x_shape = x.shape - x = x.reshape(-1) - - # look up quantized weights and qparams - qweight = nn.functional.embedding(x, self.qweight, padding_idx=self.padding_idx) - scales = nn.functional.embedding(x, self.scales, padding_idx=self.padding_idx) - if self.qzeros is not None: - qzeros = nn.functional.embedding(x, self.qzeros, padding_idx=self.padding_idx) - else: - qzeros = None - - # unpack and dequantize - return self.unpack_and_dequantize(qweight, scales, qzeros).reshape(*x_shape, -1) - - def extra_repr(self) -> str: - return ( - f"{self.rows}, {self.cols}, bits={self.quantizer.bits}," - f" symmetric={self.quantizer.symmetric}, group_size={self.quantizer.group_size}," - f" padding_idx={self.padding_idx}" - ) - - -# this is required for torchscript onnx export -class QuantEmbeddingFunction(torch.autograd.Function): - # pylint: disable=W0223,W0221 - @staticmethod - def symbolic(g, x, qweight, scales, qzeros, bits, group_size, embedding_dim): - tensor_args = [qweight, x, scales] - if qzeros is not None: - tensor_args.append(qzeros) - attrs = {"bits_i": bits, "block_size_i": group_size} - - output = g.op( - "com.microsoft::GatherBlockQuantized", - *tensor_args, - outputs=1, - **attrs, - ) - input_shape = x.type().varyingSizes() - if input_shape is not None and hasattr(x.type(), "with_sizes"): - output_type = scales.type().with_sizes([*input_shape, embedding_dim]) - output.setType(output_type) - - return output - - @staticmethod - def forward( - ctx, - x: torch.Tensor, - qweight: torch.Tensor, - scales: torch.Tensor, - qzeros: torch.Tensor, - bits: int, - group_size: int, - embedding_dim: int, - ): - if hasattr(torch.onnx, "ops") and isinstance(x, torch._subclasses.FakeTensor): # pylint: disable=W0212 - tensor_args = [qweight, x, scales] - if qzeros is not None: - tensor_args.append(qzeros) - attrs = {"bits": bits, "block_size": group_size} - return torch.onnx.ops.symbolic( - "com.microsoft::GatherBlockQuantized", - tensor_args, - attrs=attrs, - dtype=scales.dtype, - shape=[*x.shape, embedding_dim], - version=1, - ) - return torch.zeros((*x.shape, embedding_dim), dtype=scales.dtype, device=x.device) diff --git a/olive/common/quant/patterns.py b/olive/common/quant/patterns.py new file mode 100644 index 0000000000..acbac4468e --- /dev/null +++ b/olive/common/quant/patterns.py @@ -0,0 +1,245 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Pattern matching helpers for Olive quantization module selection. + +`overrides` keys and `modules_to_not_convert` entries use the following +semantics: + +- Plain string keys are *literal* for ``overrides`` (equality match) and + *substring* for ``modules_to_not_convert`` (matches the existing HF + ``modules_to_not_convert`` semantics). +- Keys prefixed with ``re:`` opt into regular-expression matching via + ``re.fullmatch``. + +When multiple ``overrides`` keys match the same target, the **first** +matching key in config (insertion) order wins (``match_override``); this +is the finalized precedence rule — "longest / most specific pattern" is +deliberately not used. + +``re:`` patterns get a best-effort, compile-time UX check for +catastrophic-backtracking *shapes* (``_assert_regex_safe``): over-long +patterns and groups that nest any repetition/alternation (at any nesting +depth) inside an unbounded outer quantifier (e.g. ``(a+)+``, +``(a{1,2})+``, ``((a|aa))+$``) are rejected with a clear error before +matching is attempted. + +This is **not a security boundary**. ``re:`` override/skip patterns come +from the user's own quantization config and run in the user's own +process — there is no trust boundary being crossed, so the scan is +intentionally a blacklist of known-bad shapes rather than a sound static +analysis. It catches common accidental catastrophic-backtracking shapes +early, with an actionable error, before they hang the process; other +exponential shapes (e.g. backreferences, adjacent unbounded quantifiers +like ``a*a*x$``) are not modeled and could still hang a user's own run. + +These helpers are the single source of truth for the matching logic and +are used by both ``OliveHfQuantizationConfig`` and the Olive walker. +""" + +from __future__ import annotations + +import re +from functools import lru_cache + +REGEX_PREFIX = "re:" + +# Safety bound for user-supplied ``re:`` patterns. ``re.fullmatch`` has no timeout and a +# length cap alone does not prevent catastrophic backtracking (short adversarial patterns +# such as ``(a+)+$`` can already hang matching). We reject patterns that combine a group +# quantified with an unbounded quantifier (``*`` / ``+`` / ``{n,}``) with a body that itself +# contains *any* repetition (bounded or unbounded) or *any* alternation at *any* nesting +# depth or a top-level alternation — the classic ReDoS shapes ``(a+)+``, ``(a*)*``, +# ``(a|a)*``, the bounded-body variant ``(a{1,2})+``, and the nested-alternation variant +# ``((a|aa))+$`` — in addition to enforcing a conservative length cap. +# +# This is a best-effort, incomplete blacklist scan, not a soundness guarantee: it catches +# the shapes above early with a clear compile-time error, but does not model every possible +# exponential-backtracking construct (e.g. backreferences, adjacent unbounded quantifiers +# like ``a*a*x$``). ``re:`` patterns are trusted, user-authored config — not adversarial +# input — so this trade-off is intentional; see the module docstring. +_MAX_REGEX_LEN = 200 + + +def is_regex_pattern(pattern: str) -> bool: + """Return True if ``pattern`` opts into regex matching.""" + return isinstance(pattern, str) and pattern.startswith(REGEX_PREFIX) + + +def _skip_char_class(raw: str, i: int) -> int: + """Return the index just past a ``[...]`` character class starting at ``raw[i] == '['``.""" + n = len(raw) + j = i + 1 + if j < n and raw[j] == "^": + j += 1 + if j < n and raw[j] == "]": # a literal ``]`` as the first class member + j += 1 + while j < n and raw[j] != "]": + if raw[j] == "\\": + j += 1 + j += 1 + return j + 1 + + +def _brace_is_unbounded(raw: str, i: int) -> bool: + """Whether the ``{...}`` quantifier starting at ``raw[i] == '{'`` has no finite upper bound.""" + close = raw.find("}", i) + if close == -1: + return False + inner = raw[i + 1 : close] + # ``{n,}`` (no upper bound) is unbounded; ``{n}`` / ``{n,m}`` are bounded. + return inner.endswith(",") or ("," in inner and inner.split(",", 1)[1] == "") + + +def _body_has_repetition_or_alt(body: str) -> bool: + """Whether a group body contains any repetition (bounded or unbounded) or alternation. + + A quantified group (``(...)+``, ``(...)*``, ``(...){n,}``) whose body itself repeats - + whether or not that inner repetition is bounded - can still exhibit exponential + backtracking (e.g. ``(a{1,2})+$`` on a run of "a"s), because the outer unbounded + quantifier can match the same substring in exponentially many ways via the inner + repetition's bounded-but-plural options. Alternation nested at *any* depth is just as + dangerous (e.g. ``((a|aa))+$`` — the ``|`` is inside an inner group, not at the outer + body's top level, but still gives the outer ``+`` exponentially many ways to partition + a matched run). So *any* repetition or alternation anywhere inside a group that is + itself quantified unboundedly is treated as unsafe — the ``depth`` counter below exists + solely to correctly skip over character-class (``[...]``) contents, not to scope where + ``|`` is considered unsafe. + + This is a best-effort blacklist scan (see module docstring) — not a claim that every + accepted pattern is polynomial-time, only that this specific family of shapes is + rejected. + """ + depth = 0 + i = 0 + n = len(body) + while i < n: + c = body[i] + if c == "\\": + i += 2 + continue + if c == "[": + i = _skip_char_class(body, i) + continue + if c == "(": + depth += 1 + elif c == ")": + depth = max(0, depth - 1) + elif c == "|" or c in "*+?" or c == "{": + return True + i += 1 + return False + + +def _assert_regex_safe(raw: str) -> None: + """Reject regex patterns matching known catastrophic-backtracking (ReDoS) shapes. + + Raises ``ValueError`` for over-long patterns and for nested unbounded quantifiers + (repetition or alternation at any depth inside a group that is itself quantified + unboundedly). This is a best-effort, incomplete static check, not a soundness + guarantee — it catches common accidental catastrophic-backtracking shapes early, with + a clear error, before they hang the process. It does not prove every accepted pattern + is safe: other exponential shapes (backreferences, adjacent unbounded quantifiers like + ``a*a*x$``) are not modeled. ``re:`` patterns are trusted, user-authored quantization + config, not adversarial input, so this trade-off is intentional (see module docstring). + + NOTE (#2598 item 5): this scanner does not understand ``(?#...)`` inline-comment syntax. + A ``[`` inside such a comment is misread as the start of a real character class, so + ``_skip_char_class`` skips forward hunting for ``]`` and can swallow an actual dangerous + shape that lies between the comment and that ``]`` -- e.g. ``(?#[)(a{1,2})+$]`` slips + past this check (Python's ``re`` ends the comment at the first unescaped ``)``, so the + regex that actually runs is just ``(a{1,2})+$]``, a known-dangerous shape). This is the + 3rd consecutive bypass of this blacklist-enumeration approach (previous two: `(a{1,2})+$`, + `((a|aa))+$`), which is a decent signal that enumerating "safe shapes" up front is the + wrong shape of fix, not that patching this one variant will finish the job. + + Decision (2026-07-31, tracked in issue #2598): NOT treated as a security vulnerability + under Olive's current trust model -- see the "not a security boundary" note in the + module docstring. Left unpatched for now. A structurally more robust option would be a + runtime timeout around the actual ``re.fullmatch`` call (catches "match takes too long" + regardless of pattern shape, so a new adversarial shape can't bypass it the way this + static scan can be bypassed) -- see the ``item5-redos-post-hoc-verification`` follow-up + note for tradeoffs (``signal.alarm`` is Unix/main-thread only; a thread+timeout wrapper + doesn't actually stop the wasted background work; a linear-time engine like `regex`'s + timeout param or `re2` would be the clean fix but adds a new dependency). + TODO: revisit if Olive's `re:` config path is ever exposed to untrusted input (e.g. a + multi-tenant service or a CI job reading externally-supplied quantization configs). + """ + if len(raw) > _MAX_REGEX_LEN: + raise ValueError( + f"Regex override/skip pattern is too long ({len(raw)} > {_MAX_REGEX_LEN} chars); " + "keep patterns short to bound matching cost." + ) + stack: list[int] = [] + i = 0 + n = len(raw) + while i < n: + c = raw[i] + if c == "\\": + i += 2 + continue + if c == "[": + i = _skip_char_class(raw, i) + continue + if c == "(": + stack.append(i) + elif c == ")": + start = stack.pop() if stack else -1 + nxt = raw[i + 1] if i + 1 < n else "" + quantified_unbounded = nxt in ("*", "+") or (nxt == "{" and _brace_is_unbounded(raw, i + 1)) + if quantified_unbounded and start != -1 and _body_has_repetition_or_alt(raw[start + 1 : i]): + raise ValueError( + f"Regex override/skip pattern {raw!r} contains a nested quantified group " + "(catastrophic-backtracking risk), e.g. `(a+)+` or `(a{1,2})+`. Rewrite the " + "pattern without a repeated/alternated group nested inside an unbounded " + "quantifier." + ) + i += 1 + + +@lru_cache(maxsize=512) +def _compiled(pattern: str) -> re.Pattern: + raw = pattern[len(REGEX_PREFIX) :] + _assert_regex_safe(raw) + return re.compile(raw) + + +def match_override(name: str, patterns) -> str | None: + """Find the override pattern for ``name`` using first-match-wins semantics. + + Patterns are evaluated in insertion order (config order); the **first** matching pattern + wins. Plain string patterns match by **literal equality**; ``re:`` patterns match by + ``re.fullmatch``. Insertion order (not "longest / most specific pattern") is the finalized, + documented precedence rule so that overlapping overrides resolve deterministically. + + Returns the original pattern (with prefix preserved) so the caller can look it up in the + underlying overrides dict, or ``None`` if no pattern matched. + """ + if not patterns: + return None + for pattern in patterns: + if is_regex_pattern(pattern): + if _compiled(pattern).fullmatch(name): + return pattern + elif name == pattern: + return pattern + return None + + +def match_skip(name: str, patterns) -> bool: + """Return True if ``name`` is matched by any skip pattern. + + Plain string patterns use **substring** matching to preserve the + existing HF ``modules_to_not_convert`` semantics. ``re:`` patterns + use ``re.fullmatch``. + """ + if not patterns: + return False + for pattern in patterns: + if is_regex_pattern(pattern): + if _compiled(pattern).fullmatch(name): + return True + elif pattern and pattern in name: + return True + return False diff --git a/olive/common/quant/selection.py b/olive/common/quant/selection.py new file mode 100644 index 0000000000..93ce56f621 --- /dev/null +++ b/olive/common/quant/selection.py @@ -0,0 +1,453 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Quantization target selection. + +Centralizes the logic that walks a model once and decides which +parameters to quantize. Both Olive's HF quantizer (which installs +:class:`QuantTensor` placeholders before weight loading) and the +PyTorch RTN/GPTQ passes (which attach calibration metadata) consume +the same set of targets — only the per-target action differs. + +Every target is a single ``nn.Parameter``, yielded as a +``(module, pname, full_name)`` tuple. ``full_name`` is the key used +for overrides / skip-pattern lookups (``module_name`` for ``"weight"`` +on ``nn.Linear`` / ``nn.Embedding``; ``f"{module_name}.{pname}"`` +otherwise). The selector makes no distinction between 2D linear / +embedding weights and 3D fused-MoE parameters — downstream code reads +the parameter's own shape and lets +:class:`~olive.common.quant.utils.WeightQuantizer` handle any rank +along the last dim. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch.nn as nn + +from olive.common.hf.io_config.io_resolver import resolve_alias +from olive.common.quant.patterns import match_skip + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + from olive.common.hf.wrapper import ModelWrapper + + +QuantTarget = tuple[nn.Module, str, str] +"""``(module, pname, full_name)`` for a single parameter selected for quantization.""" + + +def _collect_experts( + model: nn.Module, + wrapper: ModelWrapper | None, +) -> list[tuple[nn.Module, str]]: + """Return ``(experts_module, dotted_name)`` for every MoE layer.""" + if wrapper is None: + return [] + out: list[tuple[nn.Module, str]] = [] + for lw in wrapper.get_layer_wrappers(): + experts, name = lw.get_experts(return_name=True) + if experts is not None: + out.append((experts, name)) + return out + + +def _collect_moe_routers(wrapper: ModelWrapper | None) -> list[nn.Module]: + """Return the router module of every layer that also resolves an experts subtree. + + Routers decide which experts a token is sent to; quantizing them changes the routing + decisions themselves, so they are kept in full precision. Most architectures wrap the + router in a dedicated module (``MixtralTopKRouter``, ``GraniteMoeTopKRouter``, ...) that + the ``nn.Linear``/``nn.Embedding`` walk never sees, but some -- e.g. Jamba, whose + ``JambaSparseMoeBlock.router`` is a bare ``nn.Linear`` -- would otherwise be swept into + the ordinary 2D walk. Excluding by *resolved module identity* (rather than by name + pattern) covers both shapes. + + Only routers of layers with resolvable experts are excluded, so a dense layer that + happens to own an attribute named ``gate`` is never silently skipped. + """ + if wrapper is None: + return [] + routers: list[nn.Module] = [] + for lw in wrapper.get_layer_wrappers(): + get_router = getattr(lw, "get_router", None) + if get_router is None: + continue + router = get_router(return_name=False) + if router is None or lw.get_experts(return_name=False) is None: + continue + routers.append(router) + return routers + + +def _collect_shared_expert_gates(wrapper: ModelWrapper | None) -> list[nn.Module]: + """Return the shared-expert gate module of every layer that also resolves an experts subtree. + + A shared-expert gate (``qwen2_moe``, ``qwen3_5_moe``, ``qwen3_next``, ``qwen3_omni_moe``) + is a single-row ``nn.Linear`` sigmoid gate that scales an always-active "shared expert" + branch, alongside the normal top-k ``ROUTER``. Like the router, it directly controls + routing-like behavior and is disproportionately lossy to quantize (one output row), so it + is excluded the same way ``ROUTER`` is -- see ``LayerWrapper.SHARED_EXPERT_GATE``. + + Only gates of layers with resolvable experts are excluded, so a dense layer that happens + to own an attribute named ``shared_expert_gate`` is never silently skipped. + """ + if wrapper is None: + return [] + gates: list[nn.Module] = [] + for lw in wrapper.get_layer_wrappers(): + get_gate = getattr(lw, "get_shared_expert_gate", None) + if get_gate is None: + continue + gate = get_gate(return_name=False) + if gate is None or lw.get_experts(return_name=False) is None: + continue + gates.append(gate) + return gates + + +def _collect_mamba_modules(wrapper: ModelWrapper | None) -> list[nn.Module]: + """Return every layer's Mamba/SSM sub-module (state-space model), when present. + + A Mamba block's ``nn.Linear`` projections (``in_proj``/``x_proj``/``dt_proj``/``out_proj``) + feed a state-space recursion rather than a plain matmul -- generically sweeping them into + the ordinary 2D quantization walk was never intentional support, just an oversight of the + walk picking up every ``nn.Linear``/``nn.Embedding`` it can see. Excluded unconditionally + (no ``quantize_mamba`` escape hatch), the same way routers are always kept full precision. + """ + if wrapper is None: + return [] + mamba_modules: list[nn.Module] = [] + for lw in wrapper.get_layer_wrappers(): + get_mamba = getattr(lw, "get_mamba", None) + if get_mamba is None: + continue + mamba = get_mamba(return_name=False) + if mamba is not None: + mamba_modules.append(mamba) + return mamba_modules + + +# Attribute names under which multimodal checkpoints hang their vision tower. Matched against +# a module's own attribute name (the last component of its dotted ``named_modules`` name). +_VISION_TOWER_ATTR_NAMES = ("visual", "vision_tower", "vision_model", "vision_encoder") + + +def _collect_vision_towers(model: nn.Module) -> list[nn.Module]: + """Return the vision-tower sub-modules of a composite vision-language model. + + Olive's PyTorch-side (RTN/GPTQ) quantization targets the *text decoder* only; a VL + checkpoint's vision encoder is quantized separately (int8) on the ONNX side. Without this + exclusion the generic ``named_modules`` walk also sweeps in ``model.visual.*`` (patch + embeddings, attention/MLP projections, merger, ...), which would then be quantized twice — + once to int4 here and once by the later ONNX pass. + + Detection is deliberately conservative: only applied when the model's config declares a + ``vision_config`` (i.e. it really is a composite multimodal model), so a standalone vision + model — whose root module may itself be named ``vision_model`` — is never emptied of + targets. Users can still exclude additional subtrees via ``modules_to_not_convert``. + """ + config = getattr(model, "config", None) + if config is None or getattr(config, "vision_config", None) is None: + return [] + towers: list[nn.Module] = [] + for name, module in model.named_modules(): + # skip the root module (name == "") -- never treat the model itself as a vision tower + if name and name.rsplit(".", 1)[-1] in _VISION_TOWER_ATTR_NAMES: + towers.append(module) + return towers + + +def _layers_missing_experts(wrapper: ModelWrapper | None) -> list[int]: + """Return indices of layers that look structurally MoE but whose experts couldn't be resolved. + + A layer is only counted here when it *has a router* (``LayerWrapper.get_router()`` is not + ``None``) — dense layers legitimately interleaved with MoE layers (e.g. DeepSeek's + ``first_k_dense_replace``) have no router and are exempt, avoiding false positives on + architectures with a mix of dense and MoE layers. + """ + if wrapper is None: + return [] + missing: list[int] = [] + for i, lw in enumerate(wrapper.get_layer_wrappers()): + get_router = getattr(lw, "get_router", None) + if get_router is None: + # Test doubles / minimal wrappers without a get_router accessor cannot signal + # "structurally MoE" — treat as unknown (not missing) rather than raising here. + continue + router = get_router(return_name=False) + if router is None: + continue + experts = lw.get_experts(return_name=False) + if experts is None: + missing.append(i) + return missing + + +# Canonical alias key for the MoE expert count; the candidate attribute paths (flat and +# nested, e.g. DBRX's ``ffn_config.moe_num_experts``) live in +# ``olive/assets/io_configs/defaults.yaml`` under ``aliases.num_experts`` so new +# architectures are added as data, not code. +_MOE_EXPERT_COUNT_ALIAS = "num_experts" + +# Leaf attribute names used by the generic sub-config sweep below. Kept deliberately tight: +# a false positive here turns into a hard refusal to quantize. +_MOE_LEAF_ATTRS = ("num_local_experts", "num_experts", "n_routed_experts", "moe_num_experts") + +# Nested sub-config names to sweep in addition to whatever HF declares in +# ``type(config).sub_configs``. +_EXTRA_SUB_CONFIG_NAMES = ("text_config", "thinker_config", "ffn_config", "decoder_config", "llm_config") + +_MAX_SUB_CONFIG_DEPTH = 3 + + +def _positive_int(value) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def _sub_config_indicates_moe(config, depth: int) -> bool: + """Bounded sweep of nested sub-configs for a positive expert-count attribute. + + Defense-in-depth for nested-MoE architectures not yet enumerated in + ``aliases.num_experts``. Uses HF's own ``sub_configs`` declaration where available. + """ + if depth <= 0 or config is None: + return False + names = list(getattr(type(config), "sub_configs", None) or ()) + names += [n for n in _EXTRA_SUB_CONFIG_NAMES if n not in names] + for name in names: + sub = getattr(config, name, None) + if sub is None or isinstance(sub, (str, int, float, bool, list, tuple)): + continue + # ``sub`` may still be a raw dict if the config's ``post_init`` has not run. + getter = sub.get if isinstance(sub, dict) else (lambda a, _s=sub: getattr(_s, a, None)) + if any(_positive_int(getter(attr)) for attr in _MOE_LEAF_ATTRS): + return True + if not isinstance(sub, dict) and _sub_config_indicates_moe(sub, depth - 1): + return True + return False + + +def _config_indicates_moe(model: nn.Module) -> bool: + """Best-effort detection of an MoE architecture from the model config. + + Returns ``True`` only when a known MoE expert-count attribute is present and positive, + either at the top level or on a nested sub-config (e.g. DBRX's + ``config.ffn_config.moe_num_experts``). Used to fail closed when the experts subtree + cannot be resolved. + """ + config = getattr(model, "config", None) + if config is None: + return False + if _positive_int(resolve_alias(config, _MOE_EXPERT_COUNT_ALIAS)): + return True + return _sub_config_indicates_moe(config, _MAX_SUB_CONFIG_DEPTH) + + +def iter_quant_targets( + model: nn.Module, + *, + quantize_lm_head: bool, + quantize_embeds: bool, + quantize_moe: bool, + quantize_vision: bool = False, + skip_patterns: Iterable[str] = (), + extra_skip_modules: Iterable[nn.Module] = (), + skip_already_quantized: bool = True, +) -> Iterator[QuantTarget]: + """Walk ``model`` once and yield every parameter selected for quantization. + + Yielded parameters are: + + * ``nn.Linear.weight`` and ``nn.Embedding.weight`` (2D), and + * direct ``nn.Parameter`` attributes on each experts module + (typically 3D fused-MoE weights), when ``quantize_moe=True``. + + Selection rules (first matching skip wins): + + * ``extra_skip_modules`` (caller-supplied set, e.g. attention + inputs excluded by GPTQ) skips the module by identity. + * ``quantize_lm_head=False`` skips the output embedding module. + * ``quantize_embeds=False`` skips every ``nn.Embedding`` module. + ``quantize_embeds=True`` targets only ``model.get_input_embeddings()`` + when resolvable (symmetric with ``lm_head``'s precise targeting of + ``get_output_embeddings()``); falls back to every ``nn.Embedding`` + when the accessor is unavailable (e.g. non-HF synthetic fixtures). + * ``quantize_moe=False`` skips every ``nn.Module`` under any + experts subtree — this both leaves fused parameters alone *and* + prevents silently quantizing per-expert ``nn.Linear``s inside + ``ModuleList(Expert)`` blocks. + * the router module of every MoE layer is always skipped (routers + stay in full precision), including bare ``nn.Linear`` routers such + as Jamba's. + * ``quantize_vision=False`` (default) skips every module under a + composite vision-language model's vision tower (``visual`` / + ``vision_tower`` / ``vision_model`` / ``vision_encoder``): the + typical Olive pipeline quantizes the vision tower separately + downstream (e.g. on the ONNX side), so leaving it untouched here + avoids double-quantizing it. Callers that quantize a VL model + end-to-end in this single PyTorch-side pass (no separate downstream + vision quantization step) should set ``quantize_vision=True`` to + include it. + * ``skip_patterns`` matches the parameter's ``full_name`` via the + shared HF-style substring / ``re:``-prefixed regex matcher. + * When ``skip_already_quantized=True`` (default), parameters whose + underlying tensor is already a :class:`QuantTensor` are skipped + (idempotent re-runs). + """ + from olive.common.hf.wrapper import ModelWrapper + from olive.common.quant.tensor import QuantTensor + + try: + wrapper = ModelWrapper.from_model(model) + except Exception: # pylint: disable=broad-except + # Not every model is wrappable (e.g., random test fixtures). + # Without the wrapper we cannot honor MoE / lm_head / embeds + # category flags; fall back to the unfiltered 2D walk. + wrapper = None + + lm_head_module: nn.Module | None = None + if hasattr(model, "get_output_embeddings"): + lm_head_module = model.get_output_embeddings() + + # Precise input-embedding target, symmetric with ``lm_head_module`` above. When + # available, ``quantize_embeds=True`` targets *only* this module (matching the + # documented "input embeddings" contract) rather than every ``nn.Embedding`` (which + # would also sweep in positional / token-type tables). Fallback: when + # ``get_input_embeddings`` is unavailable or returns ``None`` (e.g. synthetic test + # fixtures without a full HF model), retain the broad "all nn.Embedding" behavior. + input_embeds_module: nn.Module | None = None + if hasattr(model, "get_input_embeddings"): + input_embeds_module = model.get_input_embeddings() + + expert_modules = _collect_experts(model, wrapper) + expert_module_ids = {id(m) for m, _ in expert_modules} + + # Fail-closed: if the model/config advertises an MoE architecture but we could not + # resolve any experts subtree, refuse to walk. Silently falling through to the plain + # 2D walk would (a) leave fused expert weights at full precision and (b) — worse — + # quantize every ``nn.Linear`` under an unrecognized ``ModuleList`` experts subtree + # even when ``quantize_moe=False``, reproducing the exact bug the ``moe`` flag fixes. + # Raising here (before the generator yields any target) guarantees no parameter is + # modified before the error surfaces. + if not expert_modules and _config_indicates_moe(model): + raise ValueError( + "Model config indicates a Mixture-of-Experts architecture, but Olive could not " + "locate its experts subtree (LayerWrapper.get_experts() returned nothing for every " + "layer). This architecture is not yet supported by Olive's MoE-aware quantization " + "walk. Refusing to quantize to avoid silently mis-handling the experts. Add the " + "architecture's experts/router names to LayerWrapper.EXPERTS/ROUTER, or exclude the " + "experts explicitly via modules_to_not_convert." + ) + + # Fail-closed, per-layer: even when *some* layers resolve experts, a layer that is + # structurally MoE (it has a resolvable router/gate) but whose experts subtree failed to + # resolve is a discovery failure for that layer, not a legitimately expert-free dense + # layer. Architectures that legitimately interleave dense layers with MoE layers (e.g. + # DeepSeek's ``first_k_dense_replace``) have no router on the dense layers, so they are + # exempt and do not trip this guard. + # + # Only run this check when we already have independent evidence the model is MoE (either + # the config says so, or at least one layer already resolved experts) -- ``get_router()`` + # matches purely on attribute name (default "gate"), so a genuinely dense architecture + # that happens to name an unrelated submodule ``mlp.gate`` (e.g. some gated-activation + # MLP variants) must not trip a "partially supported MoE architecture" refusal on its own. + if expert_modules or _config_indicates_moe(model): + missing_layers = _layers_missing_experts(wrapper) + if missing_layers: + total_layers = len(wrapper.get_layer_wrappers()) if wrapper is not None else 0 + raise ValueError( + "Olive detected a router/gate on " + f"{len(missing_layers)} of {total_layers} decoder layers (indices " + f"{missing_layers}) but could not resolve their experts subtree " + "(LayerWrapper.get_experts() returned nothing). This looks like a partially " + "supported Mixture-of-Experts architecture. Refusing to quantize to avoid " + "silently leaving those layers' experts unquantized (or misclassifying their " + "sub-modules) with moe=True. Add the architecture's experts/router names to " + "LayerWrapper.EXPERTS/ROUTER, or exclude the affected layers explicitly via " + "modules_to_not_convert." + ) + + # ID-based skip set for fast identity checks during the named_modules walk. + skip_ids: set[int] = {id(m) for m in extra_skip_modules} + if not quantize_lm_head and lm_head_module is not None: + skip_ids.add(id(lm_head_module)) + # Routers stay full precision regardless of ``quantize_moe`` -- see + # :func:`_collect_moe_routers`. + for router in _collect_moe_routers(wrapper): + for sub in router.modules(): + skip_ids.add(id(sub)) + # Shared-expert gates stay full precision regardless of ``quantize_moe``, same reasoning + # as routers -- see :func:`_collect_shared_expert_gates`. + for gate in _collect_shared_expert_gates(wrapper): + for sub in gate.modules(): + skip_ids.add(id(sub)) + # Mamba/SSM blocks stay full precision unconditionally -- see :func:`_collect_mamba_modules`. + for mamba in _collect_mamba_modules(wrapper): + for sub in mamba.modules(): + skip_ids.add(id(sub)) + # A composite VL model's vision tower is skipped by default -- see + # :func:`_collect_vision_towers`. Callers that quantize the whole model in one PyTorch-side + # pass (no downstream ONNX vision quantization step) can set ``quantize_vision=True`` to + # opt back in, mirroring ``quantize_moe``'s opt-in/out shape. + if not quantize_vision: + for tower in _collect_vision_towers(model): + for sub in tower.modules(): + skip_ids.add(id(sub)) + if not quantize_moe: + for experts, _ in expert_modules: + for sub in experts.modules(): + skip_ids.add(id(sub)) + + patterns = list(skip_patterns or ()) + + def _is_skipped(module: nn.Module, full_name: str) -> bool: + if id(module) in skip_ids: + return True + return bool(patterns) and match_skip(full_name, patterns) + + def _is_already_quantized(param) -> bool: + return skip_already_quantized and (isinstance(param, QuantTensor) or isinstance(param.data, QuantTensor)) + + for name, module in model.named_modules(): + # nn.Linear / nn.Embedding ``weight`` — legacy override-key + # convention: full_name == module_name. When ``quantize_embeds`` + # is False every ``nn.Embedding`` is skipped (positional / + # token-type / etc.) — this closes the loophole of an + # unintended embedding sneaking through. When ``quantize_embeds`` + # is True and ``model.get_input_embeddings()`` is resolvable, + # only that precise module is targeted (symmetric with + # ``lm_head``'s precise targeting of ``get_output_embeddings()``); + # otherwise (no HF accessor available) every ``nn.Embedding`` is + # targeted, matching the previous broad behavior for non-HF + # synthetic fixtures. + if isinstance(module, (nn.Linear, nn.Embedding)): + if isinstance(module, nn.Embedding): + if not quantize_embeds: + continue + if input_embeds_module is not None and module is not input_embeds_module: + continue + if _is_skipped(module, name): + continue + weight = module.weight + if weight is None or _is_already_quantized(weight): + continue + yield module, "weight", name + continue + + # Fused-MoE pass: direct parameters on experts modules. Only 3D fused expert + # *weight* tensors are quantization targets. Requiring ``dim() == 3`` (rather than + # ``dim() in (2, 3)``) structurally excludes 2D non-weight params such as gpt-oss's + # ``gate_up_proj_bias`` / ``down_proj_bias``, which must stay full precision. + if not quantize_moe or id(module) not in expert_module_ids: + continue + for pname, param in module.named_parameters(recurse=False): + if param is None or param.dim() != 3 or _is_already_quantized(param): + continue + full_name = f"{name}.{pname}" if name else pname + if _is_skipped(module, full_name): + continue + yield module, pname, full_name diff --git a/olive/common/quant/state_dict.py b/olive/common/quant/state_dict.py new file mode 100644 index 0000000000..b84629a7c7 --- /dev/null +++ b/olive/common/quant/state_dict.py @@ -0,0 +1,372 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=protected-access +"""State-dict helpers for Olive's quantized weight representation. + +Olive's quantization layout: + +* The quantized state of a weight named ```` (typically + ``"weight"`` for ``nn.Linear``/``nn.Embedding``, or e.g. + ``"gate_up_proj"`` for a fused-3D MoE expert tensor) is stored as + plain buffers on the host module: + + * ``_qweight`` - packed uint8 tensor (always present) + * ``_scales`` - per-group scales (always present) + * ``_qzeros`` - per-group zero points (asymmetric only) + + This matches the suffix convention so on-disk safetensors keys are + HF-loader friendly: ``model.layers.0.mlp.gate_proj.weight_qweight``, + ``...experts.gate_up_proj_qweight``, etc. + +* At runtime, the host module's original parameter + ``module._parameters[pname]`` becomes + ``nn.Parameter(QuantTensor(...), requires_grad=False)``. The + ``QuantTensor``'s inner ``qweight``/``scales``/``qzeros`` references + alias the same Python tensor objects as the buffers above, so + ``module.`` is a live view over the buffers and the original + forward (e.g. ``F.linear``) dispatches through + ``QuantTensor.__torch_function__``. + +To keep save/load simple we only need a single state-dict save hook +(to suppress the QuantTensor parameter entry, since the buffers already +carry the data) plus a post-load helper that refreshes the QuantTensor +inner references after HF assigns freshly-loaded buffer tensors. +""" + +from __future__ import annotations + +import torch + +_INSTALLED_FLAG = "_olive_quant_state_dict_hook_installed" + +QWEIGHT_SUFFIX = "_qweight" +SCALES_SUFFIX = "_scales" +QZEROS_SUFFIX = "_qzeros" + + +def buffer_names(pname: str) -> tuple[str, str, str]: + """Return the ``(qweight, scales, qzeros)`` buffer names for parameter ``pname``.""" + return f"{pname}{QWEIGHT_SUFFIX}", f"{pname}{SCALES_SUFFIX}", f"{pname}{QZEROS_SUFFIX}" + + +def _save_hook(module: torch.nn.Module, state_dict: dict, prefix: str, local_metadata: dict) -> None: + """Drop ``QuantTensor`` parameter entries from ``state_dict``. + + Inner ``qweight``/``scales``/``qzeros`` tensors are already exposed + as plain buffers on ``module`` and therefore appear in ``state_dict`` + under their own keys; the QuantTensor parameter entry is a redundant + (and non-serialisable) duplicate. + """ + # Local import to avoid a circular dependency at module-import time. + from olive.common.quant.tensor import QuantTensor + + for pname in list(module._parameters): + full_key = f"{prefix}{pname}" + value = state_dict.get(full_key) + if isinstance(value, QuantTensor): + del state_dict[full_key] + + +def install_state_dict_hooks(module: torch.nn.Module) -> None: + """Install Olive's state-dict save hook on ``module`` (idempotent).""" + if getattr(module, _INSTALLED_FLAG, False): + return + module._register_state_dict_hook(_save_hook) + setattr(module, _INSTALLED_FLAG, True) + + +def ensure_state_dict_hooks(model: torch.nn.Module) -> None: + """Install the save hook on every submodule that hosts a QuantTensor parameter. + + Belt-and-suspenders for paths that may install a ``QuantTensor`` + parameter without going through :func:`install_quant_tensor_param` + (e.g. retie helpers, future loaders). The hook is idempotent. + """ + from olive.common.quant.tensor import QuantTensor + + for sub_module in model.modules(): + for param in sub_module._parameters.values(): + if param is None: + continue + if isinstance(param, QuantTensor) or isinstance(getattr(param, "data", None), QuantTensor): + install_state_dict_hooks(sub_module) + break + + +def install_quant_tensor_param( + module: torch.nn.Module, + pname: str, + qt, # QuantTensor +) -> None: + """Install ``qt`` on ``module`` as ```` plus aliased sibling buffers. + + Replaces ``module._parameters[pname]`` with ``nn.Parameter(qt)`` and + registers ``_qweight``/``_scales``/(optionally) ``_qzeros`` + as buffers whose storage is the same as the QuantTensor's inner + tensors. The state-dict save hook is installed as a side effect. + """ + from olive.common.quant.tensor import QuantTensor + + if not isinstance(qt, QuantTensor): + raise TypeError(f"Expected QuantTensor, got {type(qt).__name__}") + + qname, sname, zname = buffer_names(pname) + + # Detach existing buffers/parameters with the same names first so + # ``register_buffer`` / parameter assignment is idempotent. + for n in (qname, sname, zname): + if n in module._buffers: + del module._buffers[n] + + # ``nn.Parameter(qt, requires_grad=False)`` for a tensor subclass + # returns the underlying QuantTensor instance directly (after a + # ``detach()`` that goes through ``_apply_fn_to_data`` and produces + # view-aliased inner tensors). See ``torch.nn.Parameter.__new__`` — + # for a tensor subclass it returns ``data.detach().requires_grad_(...)`` + # which is the QuantTensor itself, not a wrapping Parameter. We + # alias the host module's buffers to *that* instance's inner tensors + # so save / refresh paths consistently read the same storage. + param = torch.nn.Parameter(qt, requires_grad=False) + module._parameters[pname] = param + + module.register_buffer(qname, param.qweight, persistent=True) + module.register_buffer(sname, param.scales, persistent=True) + if param.qzeros is not None: + module.register_buffer(zname, param.qzeros, persistent=True) + + install_state_dict_hooks(module) + + +def get_quant_buffers(module: torch.nn.Module, pname: str): + """Return ``(qweight, scales, qzeros)`` from ``module._buffers`` for ``pname``. + + Returns ``None`` when the module does not carry a complete set of quantized buffers + (``qzeros`` is legitimately absent for symmetric quantization, ``qweight``/``scales`` + are not). + """ + qname, sname, zname = buffer_names(pname) + qweight = module._buffers.get(qname) + scales = module._buffers.get(sname) + if qweight is None or scales is None: + return None + return qweight, scales, module._buffers.get(zname) + + +def bind_quant_tensor_to_buffers(qt, module: torch.nn.Module, pname: str) -> bool: + """Point ``qt``'s inner tensor refs at ``module``'s current ``_*`` buffers. + + Returns ``True`` when at least one of the module's buffer objects differed (by identity) + from what ``qt`` referenced before, ``False`` when nothing changed or when ``module`` has + no quantized buffers for ``pname``. Note this is deliberately *not* the "was this + parameter loaded from a checkpoint" signal — that requires evidence for **every** + mandatory buffer; see :func:`_all_buffers_replaced`. + """ + buffers = get_quant_buffers(module, pname) + if buffers is None: + return False + qweight, scales, qzeros = buffers + changed = qt.qweight is not qweight or qt.scales is not scales or qt.qzeros is not qzeros + qt.qweight = qweight + qt.scales = scales + qt.qzeros = qzeros + return changed + + +def alias_module_buffers_to_quant_tensor(qt, module: torch.nn.Module, pname: str) -> None: + """Re-point ``module``'s ``_*`` buffer dict entries at ``qt``'s inner tensors. + + Used to keep tied/aliased hosting modules (e.g. ``lm_head`` tied to + ``model.embed_tokens``) pointing at the exact same buffer objects the shared + ``QuantTensor`` uses, so ``state_dict()`` / save and live forward computation agree. + """ + qname, sname, zname = buffer_names(pname) + for name, tensor in ((qname, qt.qweight), (sname, qt.scales), (zname, qt.qzeros)): + if name in module._buffers and tensor is not None: + module._buffers[name] = tensor + + +def _full_name(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + +def _collect_quant_hosts(module: torch.nn.Module): + """Group every ``QuantTensor``-hosting site in ``module`` by QuantTensor object identity. + + Returns ``{id(qt): (qt, [(prefix, sub_module, pname), ...])}``. Tied weights install the + very same ``QuantTensor`` object on two modules, so a single entry can have multiple + hosting sites; sites are listed in ``named_modules()`` order. + """ + from olive.common.quant.tensor import QuantTensor + + hosts: dict[int, tuple[QuantTensor, list[tuple[str, torch.nn.Module, str]]]] = {} + for prefix, sub_module in module.named_modules(): + for pname, param in list(sub_module._parameters.items()): + if param is None: + continue + # ``param`` itself is the QuantTensor instance stored on the module + # (``nn.Parameter(qt)`` for a tensor subclass returns the underlying + # QuantTensor — see ``torch.nn.Parameter.__new__``). + qt = param if isinstance(param, QuantTensor) else param.data + if not isinstance(qt, QuantTensor): + continue + if get_quant_buffers(sub_module, pname) is None: + continue + hosts.setdefault(id(qt), (qt, []))[1].append((prefix, sub_module, pname)) + return hosts + + +def _mandatory_buffer_names(qt, pname: str) -> tuple[str, ...]: + """Return the buffer names that must be loaded for ``pname`` to count as fully loaded. + + ``qweight`` and ``scales`` always exist. ``qzeros`` only exists for asymmetric + quantization — ``qt.qzeros`` (not ``qt.symmetric``) is authoritative, because it is + what ``install_quant_tensor_param`` keys the buffer registration off. Requiring a + ``_qzeros`` key/swap for a symmetric parameter that legitimately has no such buffer + would fail-close on a perfectly complete checkpoint. + """ + qname, sname, zname = buffer_names(pname) + return (qname, sname, zname) if qt.qzeros is not None else (qname, sname) + + +def _all_keys_in_checkpoint(qt, prefix: str, pname: str, checkpoint_keys: set[str]) -> bool: + """Whether *every* mandatory buffer of this site has a key in the checkpoint manifest. + + Checking only ``_qweight`` would call a truncated/corrupt checkpoint that carries + e.g. ``_scales`` but not ``_qweight`` "loaded", clearing ``is_placeholder`` over + zero-filled weights. + """ + return all(_full_name(prefix, name) in checkpoint_keys for name in _mandatory_buffer_names(qt, pname)) + + +def _all_buffers_replaced(qt, sub_module: torch.nn.Module, pname: str) -> bool: + """Whether *every* mandatory buffer object of this site differs (by identity) from ``qt``'s. + + Only a real ``load_state_dict(..., assign=True)`` replaces the buffer objects in + ``module._buffers``; a parameter that received nothing keeps the exact placeholder + buffer objects ``install_quant_tensor_param`` registered. Requiring *all* of them to + have been replaced (rather than any one) is what keeps a partial load — e.g. only + ``_scales`` swapped while ``_qweight`` is still the all-zero placeholder — from being + miscounted as a full load. + """ + buffers = get_quant_buffers(sub_module, pname) + if buffers is None: + return False + qweight, scales, qzeros = buffers + if qt.qweight is qweight or qt.scales is scales: + return False + # A missing ``_qzeros`` buffer is not evidence of a load either, hence the ``is None`` arm. + return not (qt.qzeros is not None and (qzeros is None or qt.qzeros is qzeros)) + + +def _select_source_site(qt, sites, checkpoint_keys: set[str] | None) -> tuple[int, bool]: + """Pick which hosting site's buffers hold the freshly-loaded checkpoint data. + + Returns ``(index_into_sites, was_loaded)``. ``was_loaded`` is ``True`` only when some + site has evidence that *all* of its mandatory buffers (see + :func:`_mandatory_buffer_names`) received real data — from the checkpoint's key + manifest or, failing that, from buffer-object identity. Partial evidence (only some of + the buffers) counts as not loaded, so the caller fails closed instead of clearing + ``is_placeholder`` over placeholder storage. + """ + if checkpoint_keys is not None: + # Exact key-membership check against the checkpoint's own manifest: authoritative + # regardless of whether the loader replaced the buffer object (``setattr``) or + # mutated it in place (``.copy_()``). For tied weights only the source module's + # keys are persisted, so this also picks the correct module out of the alias group. + for index, (prefix, _, pname) in enumerate(sites): + if _all_keys_in_checkpoint(qt, prefix, pname, checkpoint_keys): + return index, True + + # Buffer-object identity. This is the only available signal when the checkpoint's key + # set is unknown, and it is also a necessary *second* signal when the manifest is known: + # ``checkpoint_keys`` holds the raw on-disk key names, but HF can remap them while + # loading (e.g. ``save_original_format=True`` writes MoE experts as legacy per-expert + # ``experts.{i}.w1.weight_qweight`` keys that the loader fuses back into + # ``experts.gate_up_proj_qweight``), so a missing raw key does *not* by itself prove the + # parameter went unloaded. + for index, (_, sub_module, pname) in enumerate(sites): + if _all_buffers_replaced(qt, sub_module, pname): + return index, True + return 0, False + + +def refresh_quant_tensor_refs(module: torch.nn.Module, checkpoint_keys: set[str] | None = None) -> None: + """Re-point each ``QuantTensor`` parameter at the module's current buffers. + + HF's loader assigns freshly-loaded buffer tensors via + ``module.load_state_dict({name: tensor}, assign=True)``, which + replaces the buffer object in ``module._buffers``. Any + ``QuantTensor`` parameter installed earlier would still reference + the old (placeholder) storage, so we walk the parameters and re-bind + each one's inner tensors to the current buffers. + + Each **unique** ``QuantTensor`` object is processed exactly once, keyed by object + identity. Tied weights (``lm_head`` tied to ``model.embed_tokens``) host the *same* + ``QuantTensor`` object on two modules: rebinding once per hosting module would be + last-write-wins over ``named_modules()`` order, and the alias module's ``_buffers`` + entries are a one-time snapshot taken at tie time (HF's loader replaces the *source* + module's buffer objects, not the alias's), so the losing write could silently bind the + shared tensor to stale placeholder storage. This function is therefore the single + source of truth for "which buffers does a QuantTensor actually reference": after it + runs, every aliasing hosting module's ``_buffers`` entries are re-pointed at the same + objects the shared ``QuantTensor`` references, so ``state_dict()`` / save and live + forward computation agree without needing a separate repair pass. + + This is a no-op for modules with no QuantTensor parameters. + + Args: + module: root module to walk. + checkpoint_keys: when provided, the full set of tensor keys actually present + in the checkpoint (e.g. read from the safetensors file headers). It is the + preferred signal for picking the source site and for ``is_placeholder``: the + full dotted name of *every* mandatory buffer of a parameter + (``_qweight``/``_scales``, plus ``_qzeros`` for asymmetric quantization) is + checked for membership directly, independent of *how* the loader wrote the + value (``setattr`` replacement or an in-place ``.copy_()``). When ``None`` + (checkpoint format/files unknown), only the buffer-identity heuristic is used. + + Raises: + RuntimeError: when ``checkpoint_keys`` is provided (so the checkpoint manifest is + known) and a quantized parameter is still a placeholder that neither signal + shows as *fully* loaded — i.e. the manifest is missing at least one of its + mandatory buffer keys **and** at least one of its buffer objects was never + replaced by the loader. Without this the model would silently hold (fully or + partially) zero-filled placeholder weights. Fail-closed only applies when the + manifest is known; with ``checkpoint_keys=None`` the permissive identity + heuristic is preserved and nothing is raised — a partially loaded parameter + simply keeps ``is_placeholder=True``. + + """ + missing: list[str] = [] + for qt, sites in _collect_quant_hosts(module).values(): + index, was_loaded = _select_source_site(qt, sites, checkpoint_keys) + if not was_loaded and checkpoint_keys is not None and qt.is_placeholder: + prefix, _, pname = sites[0] + missing.append(_full_name(prefix, pname)) + continue + + src_prefix, src_module, src_pname = sites[index] + bind_quant_tensor_to_buffers(qt, src_module, src_pname) + for other_prefix, other_module, other_pname in sites: + if other_prefix == src_prefix and other_pname == src_pname: + continue + alias_module_buffers_to_quant_tensor(qt, other_module, other_pname) + + if was_loaded: + # Real checkpoint data is now bound to this parameter — clear the + # placeholder lifecycle flag so in-place initializers on it raise instead + # of silently no-oping (the no-op is only valid before real data is + # loaded). + qt.is_placeholder = False + + if missing: + raise RuntimeError( + "Quantized checkpoint is missing weights for: " + + ", ".join(sorted(missing)) + + ". The model's quantization_config declares these parameters as quantized, but the " + "checkpoint does not contain the complete set of `_qweight` / `_scales` " + "(/ `_qzeros`) keys for them, so they would silently load as (partially) zero-filled " + "placeholders. The checkpoint is incomplete or corrupt." + ) diff --git a/olive/common/quant/tensor.py b/olive/common/quant/tensor.py new file mode 100644 index 0000000000..5f45df8de3 --- /dev/null +++ b/olive/common/quant/tensor.py @@ -0,0 +1,762 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=protected-access,super-init-not-called,redefined-builtin,not-callable +"""QuantTensor — wrapper ``torch.Tensor`` subclass for weight-quantized parameters. + +It stores quantization buffers (``qweight``, ``scales``, ``qzeros``) but presents the +shape / dtype / device of the dequantized full-precision weight. + +Design notes: + +* The class is a **wrapper** subclass (``_make_wrapper_subclass``) — it + carries no real storage of its own, so the dense FP weight is never + materialized in memory; only the packed buffers are allocated. +* ``F.linear`` and ``F.embedding`` are dispatched via + ``__torch_function__``: + - Eager: unpack + dequantize on the fly and forward to the dense op. + - Under ``torch.onnx.is_in_onnx_export()``: raise — Olive's ONNX + conversion pass swaps any ``nn.Linear`` / ``nn.Embedding`` whose + weight is a ``QuantTensor`` for the existing exportable + ``QuantLinearNbit`` / ``QuantEmbeddingNbit`` ``nn.Module``s (see + ``olive/common/hf/quant.py``) *before* the tracer ever inspects + the parameter. This keeps the legacy ``com.microsoft::MatMulNBits`` / + ``com.microsoft::GatherBlockQuantized`` symbolic emission intact. +* All other ops (including ``model.to(dtype/device)``, ``.detach()``, + ``.contiguous()``, ``.clone()``) are routed through ``_apply_fn_to_data`` + via ``__torch_dispatch__`` so the inner buffers move with the wrapper. +* For 3D fused MoE experts (``(num_experts, out, in)``) the same buffers + carry an additional leading dim. ``__getitem__`` / ``index_select`` on + the leading dim return a 2D ``QuantTensor`` (so per-expert + ``F.linear(current_state, weight[expert_idx])`` continues to dispatch + through the same code path). +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch +import torch.nn.functional as F + +from olive.common.quant.utils import ( + WeightQuantizer, + pack_to_uint8, + unpack_from_uint8, +) + +__all__ = ["QuantTensor", "implements"] + + +_TORCH_FN_TABLE: dict[Callable, Callable] = {} + +# In-place random/constant weight-initializer ops that ``PreTrainedModel._initialize_weights`` +# (and similar module-init code paths) may call on a freshly-installed *placeholder* +# QuantTensor parameter before its real buffers are filled in from a checkpoint. Their +# numeric effect is immediately discarded once the checkpoint's ``_qweight`` / +# ``_scales`` / ``_qzeros`` buffers are loaded, so treating them as no-ops (rather than +# raising, or worse, silently dequantizing a 3D expert tensor just to throw the result away) +# is both safe and necessary for HF model loading to succeed for MoE / fused-3D targets. +# The no-op only applies while ``QuantTensor.is_placeholder`` is True (see +# ``__torch_dispatch__`` below); on a real (already-quantized) QuantTensor these ops raise +# instead, since silently no-oping there would discard real data. +_NOOP_INIT_OPS: set[Callable] = { + torch.ops.aten.normal_.default, + torch.ops.aten.uniform_.default, + torch.ops.aten.zero_.default, + torch.ops.aten.fill_.Scalar, + torch.ops.aten.fill_.Tensor, +} + + +def implements(*torch_fns: Callable) -> Callable[[Callable], Callable]: + """Register a torch-function override for ``QuantTensor``.""" + + def decorator(fn: Callable) -> Callable: + for torch_fn in torch_fns: + _TORCH_FN_TABLE[torch_fn] = fn + return fn + + return decorator + + +def _midq(bits: int) -> int: + return 1 << (bits - 1) + + +# Central message for the ONNX-export rejection of 3D fused-MoE QuantTensors. Olive's MoE +# quantization is storage-only; ONNX export of the experts is delegated to Mobius / ORT +# GenAI ModelBuilder. This must be raised for *every* unsupported 3D op that would otherwise +# reach a dequantizing fallback (not only matmul/bmm), so ONNX never silently emits a dense +# expert graph. +_MOE_ONNX_EXPORT_MSG = ( + "Olive's MoE quantization is storage-only: ONNX export of 3D fused-expert QuantTensors " + "is not supported. Non-MoE parts export through the MatMulNBits / GatherBlockQuantized " + "path; use Mobius / ORT GenAI ModelBuilder to emit com.microsoft.QMoE (or a per-expert " + "MatMulNBits loop) for the experts." +) + + +def _zero_points_or_default(weight: QuantTensor) -> torch.Tensor: + """Unpack zero_points or return a tensor full of the symmetric mid-q value.""" + if weight.qzeros is not None: + return unpack_from_uint8(weight.qzeros, weight.bits, tuple(weight.scales.shape)).to(torch.int32) + return torch.full(weight.scales.shape, _midq(weight.bits), dtype=torch.int32, device=weight.scales.device) + + +def _dequantize(weight: QuantTensor) -> torch.Tensor: + """Unpack + dequantize ``weight`` into a dense tensor of ``weight.dtype``.""" + if weight.dim() not in (2, 3): + raise NotImplementedError(f"QuantTensor only supports 2D / 3D layouts, got {weight.dim()}D") + + quantizer = WeightQuantizer( + bits=weight.bits, symmetric=weight.symmetric, group_size=weight.group_size, signed=False + ) + qw = unpack_from_uint8(weight.qweight, weight.bits, tuple(weight.shape)) + zp = _zero_points_or_default(weight) + return quantizer.dequantize(qw, weight.scales, zp).to(weight.dtype) + + +class QuantTensor(torch.Tensor): + """A weight-quantized tensor. + + Holds: + qweight: ``torch.uint8`` packed quantized values along the last + dim. Shape ``(*, math.ceil(in_features * bits / 8))``. + scales: per-group scales, dtype matches the dequantized dtype. + qzeros: ``torch.uint8`` packed zero-points, or ``None`` for + symmetric quantization. + + Attributes (non-tensor): + bits, group_size, symmetric. + + The shape / dtype / device exposed via the wrapper subclass are + those of the **dequantized** weight, so the host ``nn.Linear`` / + ``nn.Embedding`` continues to see the right metadata. + """ + + qweight: torch.Tensor + scales: torch.Tensor + qzeros: torch.Tensor | None + bits: int + group_size: int + symmetric: bool + is_placeholder: bool + + @staticmethod + def __new__( + cls, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor | None, + bits: int, + group_size: int, + symmetric: bool, + shape: torch.Size | tuple[int, ...], + dtype: torch.dtype, + is_placeholder: bool = False, + ) -> QuantTensor: + return torch.Tensor._make_wrapper_subclass( # type: ignore[attr-defined] + cls, + tuple(shape), + dtype=dtype, + device=qweight.device, + requires_grad=False, + ) + + def __init__( + self, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor | None, + bits: int, + group_size: int, + symmetric: bool, + shape: torch.Size | tuple[int, ...], + dtype: torch.dtype, + is_placeholder: bool = False, + ) -> None: + self.qweight = qweight + self.scales = scales + self.qzeros = qzeros + self.bits = int(bits) + self.group_size = int(group_size) + self.symmetric = bool(symmetric) + # Lifecycle marker. True only for the zero-filled parameter installed by + # ``OliveHfQuantizer._process_model_before_weight_loading`` *before* the checkpoint's + # ``_qweight`` / ``_scales`` / ``_qzeros`` buffers are loaded. It is cleared by + # ``refresh_quant_tensor_refs`` once real data is bound. In-place weight initializers + # (``nn.init.normal_`` & friends, called by HF's ``PreTrainedModel._initialize_weights``) + # are no-ops *only* while this is True; on a real quantized tensor they raise. + self.is_placeholder = bool(is_placeholder) + + # ------------------------------------------------------------------ + # Construction helpers + # ------------------------------------------------------------------ + + @classmethod + def from_float( + cls, + weight: torch.Tensor, + bits: int = 4, + symmetric: bool = True, + group_size: int = -1, + scales: torch.Tensor | None = None, + zero_points: torch.Tensor | None = None, + ) -> QuantTensor: + """Quantize a 2D or 3D FP weight tensor and produce a ``QuantTensor``. + + Quantization is along the last dim — for 3D fused MoE weights of + shape ``(num_experts, out, in)`` each ``(out, in)`` slice gets + its own per-group scales / zero-points along ``in``, with no + explicit leading-dim loop. + """ + if weight.dim() not in (2, 3): + raise ValueError(f"QuantTensor only supports 2D and 3D weights, got shape {tuple(weight.shape)}") + + quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False) + qparam_shape = quantizer.get_qparam_shape(tuple(weight.shape)) + + if scales is None or zero_points is None: + scales, zero_points = quantizer.find_qparams(weight) + else: + scales = scales.to(weight.device).to(weight.dtype).reshape(qparam_shape) + zero_points = zero_points.to(weight.device).to(torch.int32).reshape(qparam_shape) + + qweight_int = quantizer.quantize(weight, scales, zero_points) + qweight_packed = pack_to_uint8(qweight_int, bits).contiguous() + scales_packed = scales.reshape(qparam_shape).contiguous() + if symmetric: + if not torch.all(zero_points == quantizer.midq): + raise ValueError("Zero points must equal midq for symmetric quantization") + qzeros_packed = None + else: + qzeros_packed = pack_to_uint8(zero_points.reshape(qparam_shape), bits).contiguous() + + return cls( + qweight=qweight_packed, + scales=scales_packed, + qzeros=qzeros_packed, + bits=bits, + group_size=group_size, + symmetric=symmetric, + shape=tuple(weight.shape), + dtype=scales_packed.dtype, + ) + + @classmethod + def from_packed( + cls, + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor | None, + bits: int, + group_size: int, + symmetric: bool, + shape: tuple[int, ...], + dtype: torch.dtype | None = None, + is_placeholder: bool = False, + ) -> QuantTensor: + """Reconstruct a ``QuantTensor`` from already-packed buffers.""" + return cls( + qweight=qweight, + scales=scales, + qzeros=qzeros, + bits=bits, + group_size=group_size, + symmetric=symmetric, + shape=shape, + dtype=dtype if dtype is not None else scales.dtype, + is_placeholder=is_placeholder, + ) + + # ------------------------------------------------------------------ + # Dequantization + # ------------------------------------------------------------------ + + def to_dense(self) -> torch.Tensor: + """Unpack + dequantize into a dense FP tensor of ``self.dtype``.""" + return _dequantize(self) + + # ------------------------------------------------------------------ + # Flatten / Unflatten for torch.compile and friends + # ------------------------------------------------------------------ + + def __tensor_flatten__(self): + names = ["qweight", "scales"] + if self.qzeros is not None: + names.append("qzeros") + meta = { + "bits": self.bits, + "group_size": self.group_size, + "symmetric": self.symmetric, + "shape": tuple(self.shape), + "dtype": self.dtype, + "has_qzeros": self.qzeros is not None, + "is_placeholder": self.is_placeholder, + } + return names, meta + + @classmethod + def __tensor_unflatten__(cls, inner_tensors, meta, outer_size, outer_stride): + return cls( + qweight=inner_tensors["qweight"], + scales=inner_tensors["scales"], + qzeros=inner_tensors["qzeros"] if meta["has_qzeros"] else None, + bits=meta["bits"], + group_size=meta["group_size"], + symmetric=meta["symmetric"], + shape=meta["shape"], + dtype=meta["dtype"], + is_placeholder=meta.get("is_placeholder", False), + ) + + # ------------------------------------------------------------------ + # _apply_fn_to_data — propagate per-tensor transforms (.to, detach…) + # through every inner buffer + # ------------------------------------------------------------------ + + def _apply_fn_to_data(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> QuantTensor: + new_qweight = fn(self.qweight) + new_scales = fn(self.scales) + new_qzeros = fn(self.qzeros) if self.qzeros is not None else None + return QuantTensor( + qweight=new_qweight, + scales=new_scales, + qzeros=new_qzeros, + bits=self.bits, + group_size=self.group_size, + symmetric=self.symmetric, + shape=tuple(self.shape), + dtype=new_scales.dtype, + is_placeholder=self.is_placeholder, + ) + + # ------------------------------------------------------------------ + # Dispatch + # ------------------------------------------------------------------ + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + handler = _TORCH_FN_TABLE.get(func) + if handler is not None: + return handler(*args, **kwargs) + # Fall through to __torch_dispatch__ for everything else. + return super().__torch_function__(func, types, args, kwargs) + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + aten = torch.ops.aten + + if func in (aten.detach.default, aten.clone.default, aten.alias.default, aten.contiguous.default): + self_ = args[0] + extra_args = args[1:] + return self_._apply_fn_to_data(lambda x: func(x, *extra_args, **kwargs)) + + if func is aten._to_copy.default: + self_ = args[0] + dtype = kwargs.get("dtype") + device = kwargs.get("device") + + def _move(x: torch.Tensor) -> torch.Tensor: + copy_kwargs: dict[str, Any] = {} + if device is not None: + copy_kwargs["device"] = device + # only scales are real-dtype; keep qweight/qzeros as uint8 + if dtype is not None and x.is_floating_point(): + copy_kwargs["dtype"] = dtype + return func(x, **copy_kwargs) if copy_kwargs else x + + return self_._apply_fn_to_data(_move) + + if func is aten.copy_.default: + self_ = args[0] + src = args[1] + if not isinstance(src, QuantTensor): + raise TypeError(f"Cannot copy_ a non-QuantTensor source into a QuantTensor (got {type(src)})") + self_.qweight.copy_(src.qweight) + self_.scales.copy_(src.scales) + if self_.qzeros is not None and src.qzeros is not None: + self_.qzeros.copy_(src.qzeros) + # Mirror the source's placeholder state: copying real data into ``self_`` + # makes it real too, so a later in-place initializer must raise instead of + # silently no-oping (the no-op is only valid while no real data has landed). + # Conversely, copying *from* a placeholder leaves ``self_`` a placeholder + # (its content is still throwaway dummy data). + self_.is_placeholder = src.is_placeholder + return self_ + + if func in _NOOP_INIT_OPS and isinstance(args[0], QuantTensor): + self_ = args[0] + if self_.is_placeholder: + # In-place random/constant weight initializers (e.g. ``nn.init.normal_``) + # can reach here via HF's ``PreTrainedModel._initialize_weights`` -- it + # unconditionally (re-)initializes every parameter that isn't yet in the + # checkpoint's key set, which includes our *placeholder* QuantTensor param + # before ``from_pretrained`` fills the real ``_qweight`` / ``_scales`` + # / ``_qzeros`` buffers from the checkpoint (see ``OliveHfQuantizer``). The + # placeholder's numeric content is immediately overwritten by that buffer + # load, so these initializers are safe (and required) no-ops here rather + # than a real 3D op that would need to dequantize/reify the expert tensor. + return self_ + # A real (non-placeholder) QuantTensor already holds real quantized data: an + # in-place initializer here would silently discard it (the packed buffers + # cannot represent an arbitrary in-place mutation), so raise instead of + # no-oping. + raise RuntimeError( + f"In-place initializer {func} is not supported on a quantized QuantTensor " + f"(shape={tuple(self_.shape)}, bits={self_.bits}). Quantized storage cannot " + "represent an arbitrary in-place mutation; the call would be silently " + "discarded. Re-quantize from a dense tensor instead " + "(``QuantTensor.from_float(...)``), or mutate the packed buffers " + "(``.qweight`` / ``.scales`` / ``.qzeros``) directly." + ) + + # Fallback: dequantize any QuantTensor args and re-dispatch. + new_args = [_maybe_dense(a) for a in args] + new_kwargs = {k: _maybe_dense(v) for k, v in kwargs.items()} + return func(*new_args, **new_kwargs) + + # Friendlier repr — full dequant would defeat the purpose. + def __repr__(self) -> str: # pragma: no cover - trivial + return ( + f"QuantTensor(shape={tuple(self.shape)}, dtype={self.dtype}, device={self.device}, " + f"bits={self.bits}, group_size={self.group_size}, symmetric={self.symmetric})" + ) + + +_MOE_EAGER_OOM_MSG = ( + "Unsupported eager op on a 3D fused-MoE QuantTensor would require fully dequantizing " + "the expert tensor (OOM risk) and is refused. Slice the expert dim first (e.g. " + "`weight[expert_ids]`), or call `.to_dense()` explicitly outside a memory-sensitive path." +) + + +def _maybe_dense(x: Any) -> Any: + if isinstance(x, QuantTensor): + if x.dim() >= 3: + # An unregistered op reaching this generic fallback would otherwise fully + # dequantize the (potentially huge) fused-3D expert tensor. Refuse in both + # eager mode (OOM risk) and under ONNX export (storage-only contract). + raise RuntimeError(_MOE_ONNX_EXPORT_MSG if torch.onnx.is_in_onnx_export() else _MOE_EAGER_OOM_MSG) + return x.to_dense() + return x + + +# ---------------------------------------------------------------------- +# Torch-function overrides +# ---------------------------------------------------------------------- + + +@implements(F.linear) +def _linear(input: torch.Tensor, weight: QuantTensor, bias: torch.Tensor | None = None) -> torch.Tensor: # noqa: A002 + if torch.onnx.is_in_onnx_export(): + raise RuntimeError( + "Olive QuantTensor cannot be traced by torch.onnx.export directly. " + "Use olive.common.hf.quant.make_export_compatible_quant(model, dynamo=...) " + "before exporting, which replaces nn.Linear modules backed by a " + "QuantTensor with an exportable QuantLinearNbit nn.Module." + ) + if weight.dim() != 2: + raise RuntimeError( + "F.linear expects a 2D weight; got a " + f"{weight.dim()}D QuantTensor. For 3D fused MoE experts, slice the leading " + "dim first (e.g. `weight[expert_idx]`)." + ) + dense = weight.to_dense().to(input.dtype) + return F.linear(input, dense, bias) + + +@implements(F.embedding) +def _embedding( + input: torch.Tensor, # noqa: A002 + weight: QuantTensor, + padding_idx: int | None = None, + max_norm: float | None = None, + norm_type: float = 2.0, + scale_grad_by_freq: bool = False, + sparse: bool = False, +) -> torch.Tensor: + if torch.onnx.is_in_onnx_export(): + raise RuntimeError( + "Olive QuantTensor cannot be traced by torch.onnx.export directly. " + "Use olive.common.hf.quant.make_export_compatible_quant(model, dynamo=...) " + "before exporting, which replaces nn.Embedding modules backed by a " + "QuantTensor with an exportable QuantEmbeddingNbit nn.Module." + ) + if weight.dim() != 2: + raise RuntimeError(f"F.embedding expects a 2D weight; got a {weight.dim()}D QuantTensor.") + dense = weight.to_dense() + return F.embedding(input, dense, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse) + + +def _quant_metadata(t: QuantTensor) -> tuple: + """Non-tensor identity of a ``QuantTensor`` (everything except the buffer contents).""" + return (t.bits, t.group_size, t.symmetric, tuple(t.shape), t.dtype, t.qzeros is not None) + + +@implements(torch.equal) +def _equal(a, b) -> bool: + """Structural equality that never materializes the dense weight. + + ``transformers>=5``'s ``PreTrainedModel.tie_weights`` calls ``torch.equal`` on the two + tied word-embedding parameters *during* ``from_pretrained`` (in + ``_finalize_model_loading``, i.e. **before** the quantizer's + ``_process_model_after_weight_loading`` hook runs). At that point the placeholder + ``QuantTensor``'s inner buffers may still live on the ``meta`` device, and the generic + ``_maybe_dense`` fallback would dequantize them -- which hard-fails with + ``NotImplementedError: aten::equal ... with Meta tensors``. Comparing the packed buffers + directly (with an object-identity fast path for the tied case) avoids both the crash and + a needless full dequantization of a potentially huge weight. + """ + if a is b: + return True + if isinstance(a, QuantTensor) and isinstance(b, QuantTensor): + if _quant_metadata(a) != _quant_metadata(b): + return False + pairs = [(a.qweight, b.qweight), (a.scales, b.scales)] + if a.qzeros is not None: + pairs.append((a.qzeros, b.qzeros)) + for x, y in pairs: + if x is y: + continue + if x.device.type == "meta" or y.device.type == "meta": + # ``meta`` tensors carry no data, so only storage identity is knowable and + # two distinct meta buffers cannot be proven equal. + return False + if not torch.equal(x, y): + return False + return True + # Mixed QuantTensor / dense comparison: fall back to the dense path. + return torch.equal(_maybe_dense(a), _maybe_dense(b)) + + +@implements(torch.matmul, torch.Tensor.matmul) +def _matmul(a, b): + if torch.onnx.is_in_onnx_export() and (isinstance(a, QuantTensor) or isinstance(b, QuantTensor)): + raise RuntimeError( + "ONNX export of matmul on a QuantTensor is not supported in Olive. " + "Olive's MoE quantization is storage-only; use Mobius to emit " + "com.microsoft.QMoE or a per-expert MatMulNBits loop." + ) + return torch.matmul(_maybe_dense(a), _maybe_dense(b)) + + +@implements(torch.bmm) +def _bmm(a, b): + if torch.onnx.is_in_onnx_export() and (isinstance(a, QuantTensor) or isinstance(b, QuantTensor)): + raise RuntimeError( + "ONNX export of bmm on a QuantTensor is not supported in Olive. " + "Olive's MoE quantization is storage-only; use Mobius to emit " + "com.microsoft.QMoE or a per-expert MatMulNBits loop." + ) + return torch.bmm(_maybe_dense(a), _maybe_dense(b)) + + +@implements(torch.Tensor.__getitem__) +def _getitem(self: QuantTensor, idx): + # Selecting along the leading (expert) dim of a 3D QuantTensor keeps the result on the + # quantized fast path — for scalar ``weight[int]`` selection, a 0-D/1-D integer tensor + # or list of ints (e.g. a flattened ``(tokens,)`` expert-id gather), a boolean/uint8 + # mask, a slice, and tuple-form leading-only indexing ``weight[expert_ids, :, :]``. + # Indexing the packed buffers directly avoids dequantizing the *entire* expert tensor + # (the OOM risk that would otherwise defeat MoE quantization). + # + # Rank >= 2 integer-tensor indices (e.g. an un-flattened ``(tokens, k)`` top-k routing + # tensor) are deliberately NOT accepted here, even though the packed-buffer arithmetic + # below would happily produce a >3D QuantTensor "view" for them (see #2598 item 4): + # every dense-consuming op (``_dequantize``, ``F.linear``/``F.embedding``, and the + # generic ``_maybe_dense`` OOM guard used by ``matmul``/``bmm``/every unregistered op) + # refuses any QuantTensor with rank > 3, so such a result would be a dead end — it + # could never be dequantized or fed into any op. There is currently no caller of this + # rank in the repo, and no validated design for how the OOM guard should distinguish a + # small already-gathered batch (safe to dequantize) from the full multi-GB expert + # tensor (unsafe) once rank stops being a reliable proxy for size. Callers needing + # multi-dim batched expert selection should flatten to 1-D first (e.g. + # ``weight[expert_ids.flatten()]``), consume the result, then reshape the *dense + # output* back to the original batch shape — the standard "flatten batch dims / do the + # op / unflatten" pattern. Revisit this once a real consumer (e.g. a vectorized + # GPTQ-MoE forward/calibration path) exists to validate the OOM-guard redesign against. + if self.dim() == 3 and _indexes_leading_dim_only(idx, int(self.shape[0])): + new_qweight = self.qweight[idx] + new_scales = self.scales[idx] + new_qzeros = self.qzeros[idx] if self.qzeros is not None else None + # Leading batch dim of the selection (``()`` for int, ``(K,)`` for a 0-D/1-D + # tensor/list/mask index) — always rank <= 1 now that ``_indexes_leading_dim_only`` + # rejects rank >= 2 integer tensors (#2598 item 4). + leading = tuple(new_qweight.shape[:-2]) + new_shape = (*leading, *tuple(self.shape[1:])) + # Belt-and-braces: catches any residual arithmetic slip in ``_indexes_leading_dim_only`` + # (or a future extension of it) before it produces a QuantTensor whose ``.shape`` + # metadata disagrees with the actual packed-buffer ranks. + if tuple(new_scales.shape[:-2]) != leading: + raise RuntimeError( + f"Internal error: QuantTensor index {idx!r} produced inconsistent buffer ranks " + f"(qweight {tuple(new_qweight.shape)} vs scales {tuple(new_scales.shape)})." + ) + return QuantTensor( + qweight=new_qweight, + scales=new_scales, + qzeros=new_qzeros, + bits=self.bits, + group_size=self.group_size, + symmetric=self.symmetric, + shape=new_shape, + dtype=self.dtype, + is_placeholder=self.is_placeholder, + ) + if self.dim() >= 3: + # Any other 3D indexing pattern would require fully dequantizing the expert tensor + # (this also covers rank >= 2 integer-tensor indices, e.g. an un-flattened + # ``(tokens, k)`` top-k routing tensor -- see #2598 item 4 / the note above). + if torch.onnx.is_in_onnx_export(): + raise RuntimeError(_MOE_ONNX_EXPORT_MSG) + raise RuntimeError( + f"Unsupported indexing pattern {idx!r} on a {self.dim()}D fused-MoE QuantTensor. " + "Only leading-dim (expert) selection preserves quantized storage (boolean masks " + "must be 1-D and match the expert dim; integer tensor/list indices must be 0-D " + "or 1-D -- flatten a multi-dim batch of expert ids first, e.g. " + "`weight[expert_ids.flatten()]`, and reshape the *dense output* back afterward); " + "other indexing would require fully dequantizing the expert tensor and is " + "refused to avoid silent memory blow-up. Slice the expert dim first (e.g. " + "`weight[expert_ids]`)." + ) + return self.to_dense()[idx] + + +def _is_bool_index(idx: Any) -> bool: + """Whether ``idx`` is (or contains, for a Python list) a boolean mask. + + ``torch.uint8`` tensors are included here on purpose: PyTorch's indexing + semantics treat a ``uint8`` tensor index as a *legacy boolean mask* + (``nonzero()``-style selection, with a deprecation warning), not as an + integer gather index — unlike every other unsigned/narrow integer dtype + (``int8``/``int16``/``uint16``/``uint32``/``uint64``), which torch simply + refuses to index with. Without this, a 1-D uint8 index of the wrong + length, or a 0-D uint8 scalar (e.g. ``expert_idx[0]`` in real MoE routing + code), would fall through to the generic "any non-float/complex tensor is + a safe gather index" branch below and silently produce a shape-inserting + result instead of raising or correctly selecting. + """ + if isinstance(idx, torch.Tensor): + return idx.dtype in (torch.bool, torch.uint8) + if isinstance(idx, list): + return any(isinstance(e, bool) for e in idx) + return isinstance(idx, bool) + + +def _indexes_leading_dim_only(idx: Any, leading_dim: int) -> bool: + """Whether ``idx`` selects along a single (leading) dimension only. + + Safe forms — each consumes *exactly one* dimension, so ``_getitem``'s + ``(*index_batch_dims, *self.shape[1:])`` shape arithmetic stays valid: + + * ``int`` / ``slice`` / list of ``int`` / a **0-D or 1-D integer** tensor *except* + ``uint8`` (e.g. a flattened ``(tokens,)`` expert-id gather); + * a **1-D boolean (or legacy ``uint8``) mask** whose length equals ``leading_dim``. + + Rejected: rank >= 2 integer tensors (see the module-level note in ``_getitem`` on why + un-flattened multi-dim batch indices like a ``(tokens, k)`` top-k routing tensor are + deliberately not supported here — #2598 item 4), rank != 1 boolean/``uint8`` masks and + length-mismatched masks (they consume more or fewer than one dim — the source of the + metadata/data shape mismatch), ``bool`` scalars (which *add* a dim), and lists + containing ``bool`` (torch treats them as masks). ``uint8`` tensors are treated as + masks (never as gather indices) because that is PyTorch's own legacy indexing + semantics for that dtype — see ``_is_bool_index``. + + Tuples are accepted only when the head is a valid leading index and every remaining + element is a full slice (``:``). + """ + if isinstance(idx, tuple): + if not idx or len(idx) > 3: + return False + head, *rest = idx + if not _indexes_leading_dim_only(head, leading_dim): + return False + return all(isinstance(r, slice) and r == slice(None) for r in rest) + if _is_bool_index(idx): + if isinstance(idx, torch.Tensor): + return idx.dim() == 1 and idx.shape[0] == leading_dim + # Python list-of-bools mask. + return isinstance(idx, list) and all(isinstance(e, bool) for e in idx) and len(idx) == leading_dim + if isinstance(idx, slice): + return True + if isinstance(idx, int): # ``bool`` already handled above + return True + if isinstance(idx, list): + return all(isinstance(e, int) and not isinstance(e, bool) for e in idx) + # Rank >= 2 integer tensors are rejected here on purpose -- see the docstring above and + # the note in ``_getitem`` (#2598 item 4): a rank-2+ leading index would produce a + # QuantTensor with rank > 3, which every dense-consuming op in this module refuses. + return isinstance(idx, torch.Tensor) and not idx.is_floating_point() and not idx.is_complex() and idx.dim() <= 1 + + +@implements(torch.Tensor.to) +def _to(self: QuantTensor, *args, **kwargs): + # Use torch's own _parse_to to robustly resolve the (device, dtype, + # non_blocking, convert_to_format) tuple — covers every signature + # including nn.Module.to's ``t.to(None, dtype, non_blocking)``. + device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs) # type: ignore[attr-defined] + + if device is None and dtype is None: + return self + + def _move(x: torch.Tensor) -> torch.Tensor: + move_kwargs: dict[str, Any] = {} + if device is not None: + move_kwargs["device"] = device + if dtype is not None and x.is_floating_point(): + move_kwargs["dtype"] = dtype + return x.to(**move_kwargs) if move_kwargs else x + + return self._apply_fn_to_data(_move) + + +# ---------------------------------------------------------------------- +# Movement / view ops +# ---------------------------------------------------------------------- +# Storage-only quantization intercepts ``F.linear`` / ``F.embedding`` (and integer +# expert selection). Shape-movement ops (transpose / reshape / view / permute / …) are +# *not* implemented against the packed layout: letting them fall through to the default +# tensor-subclass machinery silently produces a malformed ``QuantTensor`` (constructed via +# ``__new__`` without the packed metadata), which then fails deep inside an unrelated op. +# We reject them centrally with a clear, actionable error instead — and under ONNX export +# surface the same MoE-export guidance as every other unsupported 3D op. +_UNSUPPORTED_MOVEMENT_FNS = ( + torch.transpose, + torch.Tensor.transpose, + torch.t, + torch.Tensor.t, + torch.permute, + torch.Tensor.permute, + torch.swapaxes, + torch.Tensor.swapaxes, + torch.movedim, + torch.Tensor.movedim, + torch.reshape, + torch.Tensor.reshape, + torch.Tensor.view, + torch.flatten, + torch.Tensor.flatten, + torch.Tensor.expand, + torch.squeeze, + torch.Tensor.squeeze, + torch.unsqueeze, + torch.Tensor.unsqueeze, +) + + +@implements(*_UNSUPPORTED_MOVEMENT_FNS) +def _unsupported_movement(*args, **kwargs): + self = next((a for a in args if isinstance(a, QuantTensor)), None) + if self is not None and self.dim() >= 3 and torch.onnx.is_in_onnx_export(): + raise RuntimeError(_MOE_ONNX_EXPORT_MSG) + raise RuntimeError( + "Shape-movement / view ops (transpose, reshape, view, permute, flatten, expand, …) are " + "not supported on an Olive QuantTensor: quantization is storage-only and these ops would " + "require fully dequantizing the packed weight. Access the dense weight explicitly via " + "``.to_dense()`` if you really need to reshape it (e.g. outside a memory-sensitive path)." + ) diff --git a/olive/common/quant/utils.py b/olive/common/quant/utils.py index 8163e2b4fa..289d4e4e90 100644 --- a/olive/common/quant/utils.py +++ b/olive/common/quant/utils.py @@ -9,7 +9,22 @@ # TODO(jambayk): consider supporting transposed weights, useful for onnx weights class WeightQuantizer: - """Class to quantize weight tensors.""" + """Class to quantize weight tensors. + + Operates on N-D tensors and always quantizes along the **last** + dimension. For a tensor with shape ``(*leading_dims, last)``: + + * per-tensor (``group_size=0``): a single scalar scale shared across + all elements; ``scales`` has shape ``(1,) * ndim``. + * per-channel (``group_size=-1``): one scale per leading-index; + ``scales`` has shape ``(*leading_dims, 1)``. + * groupwise (``group_size>0``): ``last`` must be divisible by + ``group_size``; ``scales`` has shape ``(*leading_dims, last // group_size)``. + + The same code path handles the 2D ``(out, in)`` case used by + ``nn.Linear``/``nn.Embedding`` and the 3D fused-MoE + ``(num_experts, out, in)`` case — no leading-dim loop required. + """ def __init__(self, bits: int = 4, symmetric: bool = True, group_size: int = 0, signed: bool = False): """Initialize the quantizer with parameters. @@ -21,7 +36,7 @@ def __init__(self, bits: int = 4, symmetric: bool = True, group_size: int = 0, s signed: Whether to use signed quantization (default is False, meaning unsigned) """ - assert bits in [2, 4, 8], "Only 4-bit and 8-bit quantization supported" + assert bits in [2, 4, 8], "Only 2-bit, 4-bit and 8-bit quantization supported" self.bits = bits self.symmetric = symmetric self.group_size = group_size @@ -30,63 +45,64 @@ def __init__(self, bits: int = 4, symmetric: bool = True, group_size: int = 0, s self.maxq, self.minq = get_maxq_minq(self.bits, self.signed) self.midq = (self.maxq + self.minq + 1) // 2 - def get_num_groups(self, shape: tuple[int, int]) -> int: - """Get the number of groups for quantization based on the input shape and group_size. + def get_num_groups(self, shape: tuple[int, ...]) -> int: + """Get the number of groups along the last dim. Args: - shape: The shape (out_features, in_features) of the tensor to quantize + shape: The shape of the tensor to quantize. Returns: - The number of groups for quantization + The number of groups along the last dim. """ if self.group_size == 0: raise ValueError("group_size must be greater than 0 for groupwise quantization") - group_size = self.group_size if self.group_size > 0 else shape[1] - assert shape[1] % group_size == 0, f"in_features {shape[1]} must be divisible by group_size {group_size}" - return shape[1] // group_size + last = shape[-1] + group_size = self.group_size if self.group_size > 0 else last + assert last % group_size == 0, f"last dim {last} must be divisible by group_size {group_size}" + return last // group_size - def get_qparam_shape(self, shape: tuple[int, int]) -> tuple[int, ...]: - """Get the shapes for quantization parameters based on the input shape and group_size. + def get_qparam_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: + """Get the shape for scales / zero-points given the input shape. Args: - shape: The shape (out_features, in_features) of the tensor to quantize + shape: The shape of the tensor to quantize. Returns: - A tuple of shapes for scales and zero points + The shape of the scales / zero-points tensor. """ if self.group_size == 0: - return (1, 1) - return (shape[0], self.get_num_groups(shape)) + return (1,) * len(shape) + return (*shape[:-1], self.get_num_groups(shape)) @torch.no_grad() def find_qparams(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Find quantization parameters (scale and zero point) for the given tensor. Args: - tensor: The tensor to quantize. Expected to be 2D with shape (out_features, in_features). + tensor: The N-D tensor to quantize. Quantization is along the last dim. Returns: - A tuple of (scale, zero_point) + A tuple of (scale, zero_point) with shape :meth:`get_qparam_shape`. """ tensor, _ = self._reshape_tensor(tensor) - # calculate min and max - tmp = torch.zeros(tensor.shape[0:-1], device=tensor.device, dtype=tensor.dtype) - min_val = torch.minimum(tensor.min(-1)[0], tmp) - max_val = torch.maximum(tensor.max(-1)[0], tmp) + # calculate min and max along the last (group) dim + zero = torch.zeros(tensor.shape[:-1], device=tensor.device, dtype=tensor.dtype) + min_val = torch.minimum(tensor.min(-1)[0], zero) + max_val = torch.maximum(tensor.max(-1)[0], zero) if self.symmetric: max_val = torch.maximum(abs(min_val), max_val) - tmp = min_val < 0 - if torch.any(tmp): - min_val[tmp] = -max_val[tmp] + neg = min_val < 0 + if torch.any(neg): + min_val[neg] = -max_val[neg] - tmp = (min_val == 0) & (max_val == 0) - min_val[tmp] = -1 - max_val[tmp] = 1 + zero_pair = (min_val == 0) & (max_val == 0) + min_val[zero_pair] = -1 + max_val[zero_pair] = 1 scales = (max_val - min_val) / (self.maxq - self.minq) if self.symmetric: @@ -102,12 +118,12 @@ def quantize(self, tensor: torch.Tensor, scales: torch.Tensor, zero_points: torc """Quantize the given tensor using the provided scales and zero points. Args: - tensor: The tensor to quantize. Expected to be 2D with shape (out_features, in_features). - scales: The scales for quantization. - zero_points: The zero points for quantization. + tensor: The N-D tensor to quantize. + scales: The scales with shape :meth:`get_qparam_shape`. + zero_points: The zero points with shape :meth:`get_qparam_shape`. Returns: - The quantized tensor. + The quantized tensor with the original input shape. """ tensor, shape = self._reshape_tensor(tensor) @@ -120,20 +136,19 @@ def quantize(self, tensor: torch.Tensor, scales: torch.Tensor, zero_points: torc @torch.no_grad() def dequantize(self, q_tensor: torch.Tensor, scales: torch.Tensor, zero_points: torch.Tensor) -> torch.Tensor: - """Dequantize the given quantized tensor using the provided scales and zero points. + """Dequantize the given quantized tensor. Args: - q_tensor: The quantized tensor to dequantize. Expected to be 2D with shape (out_features, in_features). - scales: The scales for dequantization. - zero_points: The zero points for dequantization. + q_tensor: The quantized N-D tensor. + scales: The scales with shape :meth:`get_qparam_shape`. + zero_points: The zero points with shape :meth:`get_qparam_shape`. Returns: - The dequantized tensor. + The dequantized tensor with the original input shape. """ q_tensor, shape = self._reshape_tensor(q_tensor) - # apply dequantization # both q_tensor and zero_points should be int32, so no need to worry about overflow tensor = (q_tensor - zero_points.unsqueeze(-1)) * scales.unsqueeze(-1) return tensor.reshape(shape) @@ -144,7 +159,7 @@ def fake_quantize( """Fake quantize the given tensor using the provided scales and zero points. Args: - tensor: The tensor to quantize. Expected to be 2D with shape (out_features, in_features). + tensor: The N-D tensor to quantize. scales: The scales for quantization. If None, scales will be computed from the tensor. zero_points: The zero points for quantization. If None, zero points will be computed from the tensor. @@ -158,20 +173,16 @@ def fake_quantize( return self.dequantize(q_tensor, scales, zero_points) def _reshape_tensor(self, tensor: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]: - """Reshape the tensor based on the group size. - - Args: - tensor: The tensor to reshape. - - Returns: - The reshaped tensor and its original shape. + """Reshape so the last dim becomes ``(num_groups, group_size)``. + Returns the reshaped tensor and its original shape. """ - shape = tensor.shape + shape = tuple(tensor.shape) if self.group_size == 0: - tensor = tensor.reshape(1, 1, -1) + # per-tensor: collapse everything into a single (1, ..., 1, prod) tensor + tensor = tensor.reshape(*((1,) * (len(shape) - 1)), 1, -1) else: - tensor = tensor.reshape(shape[0], self.get_num_groups(shape), -1) + tensor = tensor.reshape(*shape[:-1], self.get_num_groups(shape), -1) return tensor, shape @@ -198,14 +209,17 @@ def get_maxq_minq(bits: int, signed: bool) -> tuple[int, int]: @torch.no_grad() def pack_to_uint8(tensor: torch.Tensor, bits: int) -> torch.Tensor: - """Pack 2/4/8 bit tensor into uint8 tensor on the last dimension. + """Pack 2/4/8 bit values into uint8 along the last dimension. + + Works for tensors of any rank — only the last dim is packed; all + leading dims are preserved. Args: - tensor: The 2D input tensor. Values are expected to be unsigned in the range [0, 2^bits - 1] + tensor: The input tensor. Values are expected to be unsigned in the range ``[0, 2^bits - 1]``. bits: Number of bits (2, 4 or 8) Returns: - A tensor of uint8 values with packed data + A tensor of uint8 values with packed data along the last dim. """ assert bits in [2, 4, 8], "Only 2-bit, 4-bit and 8-bit quantization supported" @@ -216,44 +230,47 @@ def pack_to_uint8(tensor: torch.Tensor, bits: int) -> torch.Tensor: packing_factor = 8 // bits - # padd if necessary to ensure tensor shape is divisible by packing_factor + # pad the last dim so it is divisible by ``packing_factor`` num_padding = (-tensor.shape[-1]) % packing_factor if num_padding > 0: - pad = (0, num_padding, 0, 0) - tensor = torch.nn.functional.pad(tensor, pad, mode="constant", value=0) + tensor = torch.nn.functional.pad(tensor, (0, num_padding), mode="constant", value=0) packed_size = tensor.shape[-1] // packing_factor tensor = tensor.to(torch.uint8) packed_tensor = torch.zeros( - (tensor.shape[0], packed_size), + (*tensor.shape[:-1], packed_size), dtype=torch.uint8, device=tensor.device, ) for i in range(packing_factor): - packed_tensor |= tensor[:, i::packing_factor] << bits * i + packed_tensor |= tensor[..., i::packing_factor] << bits * i return packed_tensor @torch.no_grad() -def unpack_from_uint8(packed_tensor: torch.Tensor, bits: int, shape: tuple[int, int]) -> torch.Tensor: - """Unpack uint8 tensor into 2/4/8 bit tensor on the last dimension. +def unpack_from_uint8(packed_tensor: torch.Tensor, bits: int, shape: tuple[int, ...]) -> torch.Tensor: + """Unpack a uint8-packed tensor into 2/4/8 bit values along the last dimension. + + Works for tensors of any rank — only the last dim is unpacked; all + leading dims must match ``shape[:-1]``. Args: - packed_tensor: The 2D input tensor with packed uint8 values + packed_tensor: The packed uint8 tensor. bits: Number of bits (2, 4 or 8) - shape: The original shape of the tensor before packing + shape: The original shape of the tensor before packing. Returns: - A tensor of int32 values with unpacked data + A tensor of int32 values with unpacked data. """ assert packed_tensor.dtype == torch.uint8, "Input tensor must be of dtype uint8" maxq, _ = get_maxq_minq(bits, signed=False) - wf = torch.arange(0, 8, bits, device=packed_tensor.device, dtype=torch.uint8).unsqueeze(0) - - unpacked_tensor = torch.bitwise_right_shift(packed_tensor.unsqueeze(2), wf.unsqueeze(0)) - unpacked_tensor = unpacked_tensor.reshape(shape[0], -1) - unpacked_tensor = unpacked_tensor[:, : shape[1]] + wf = torch.arange(0, 8, bits, device=packed_tensor.device, dtype=torch.uint8) + # (..., packed_size, packing_factor) + unpacked_tensor = torch.bitwise_right_shift(packed_tensor.unsqueeze(-1), wf) + # collapse last two dims and trim padding back to original last-dim size + unpacked_tensor = unpacked_tensor.reshape(*packed_tensor.shape[:-1], -1) + unpacked_tensor = unpacked_tensor[..., : shape[-1]] return torch.bitwise_and(unpacked_tensor, maxq).to(torch.int32) diff --git a/olive/common/utils.py b/olive/common/utils.py index 5f39996f54..805836d230 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): @@ -481,8 +496,12 @@ def hardlink_copy_file(src, dst, *, follow_symlinks=True): dst.unlink() try: - os.link(src, dst, follow_symlinks=follow_symlinks) - except OSError as e: + try: + os.link(src, dst, follow_symlinks=follow_symlinks) + except NotImplementedError: + # Python 3.14 on Windows does not support the follow_symlinks argument. + os.link(src, dst) + except (NotImplementedError, OSError) as e: # for instance, hardlinking across filesystems is not supported logger.debug("Linking failed with %s. Copying.", e) shutil.copy2(src, dst, follow_symlinks=follow_symlinks) diff --git a/olive/data/component/dataloader.py b/olive/data/component/dataloader.py index d0be7acf0c..ca9a449034 100644 --- a/olive/data/component/dataloader.py +++ b/olive/data/component/dataloader.py @@ -8,7 +8,7 @@ import torch -from olive.common.onnx_io import get_kv_info +from olive.common.onnx_io import get_genai_decoder_config, get_kv_info from olive.common.utils import format_data from olive.data.registry import Registry from olive.logging import get_verbosity @@ -133,6 +133,13 @@ def __init__(self, dataloader, model_path: str, io_config: dict): self.position_ids = "position_ids" in self.io_config["input_names"] self.past_seq_len = "past_seq_len" in self.io_config["input_names"] self.kv_info = get_kv_info(self.io_config) + if self.kv_info and not all(isinstance(self.kv_info[key], int) for key in ("num_kv_heads", "head_size")): + # onnxruntime-genai exports the KV cache head size symbolically, so the concrete + # values have to be read back from the genai_config.json next to the model. + decoder_config = get_genai_decoder_config(model_path) + if decoder_config: + self.kv_info["num_kv_heads"] = decoder_config["num_key_value_heads"] + self.kv_info["head_size"] = decoder_config["head_size"] self.has_gqa = False for node in onnx.load(self.model_path, load_external_data=False).graph.node: if node.op_type == "GroupQueryAttention": diff --git a/olive/data/component/text_generation.py b/olive/data/component/text_generation.py index f53debeafc..1520219176 100644 --- a/olive/data/component/text_generation.py +++ b/olive/data/component/text_generation.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +import logging from pathlib import Path from random import Random from typing import Callable, Optional, Union @@ -16,6 +17,8 @@ from olive.data.component.dataset import ClassificationDataset from olive.data.constants import IGNORE_INDEX +logger = logging.getLogger(__name__) + class TextGenStrategy(StrEnumBase): """Strategy for tokenizing a dataset.""" @@ -253,8 +256,10 @@ def text_gen_pre_process(dataset, tokenizer, all_kwargs): joined_input_ids += input_ids + joiner_tokens end_loc = 0 # position of unused token in joined_input_ids - # '- args.max_seq_len ' is used to make sure we don't get a sequence that is too short - for begin_loc in range(0, len(joined_input_ids) - args.max_seq_len, step): + # '- args.max_seq_len' is used to make sure we don't get a sequence that is too short + # '+ 1' since range is exclusive at the end: a block starting exactly at + # len(joined_input_ids) - max_seq_len is still a full length block + for begin_loc in range(0, len(joined_input_ids) - args.max_seq_len + 1, step): # end_loc is the beginning of the next sequence end_loc = begin_loc + args.max_seq_len # get the input sequence @@ -358,11 +363,24 @@ def text_gen_pre_process(dataset, tokenizer, all_kwargs): # found a good sample break resamples += 1 - if not encodings: - # could not find a good sample after resampling + if encodings is None or (args.drop_short_sequences and encodings.input_ids.shape[1] < args.max_seq_len): + # could not find a good sample after resampling. `encodings` holds the last rejected + # row once the retries are exhausted, and a non-empty BatchEncoding is truthy, so the + # length has to be re-checked explicitly to honor drop_short_sequences. continue append_text_gen_input_ids(tokenized_inputs, encodings.input_ids[0], encodings.attention_mask[0]) + # every strategy can stop early and silently deliver fewer samples than requested, for reasons that + # depend on the strategy. Warn so that the shortfall is visible to the user. + num_delivered = len(tokenized_inputs["input_ids"]) + if args.max_samples is not None and num_delivered < args.max_samples: + logger.warning( + "Only %d samples were generated but max_samples=%d was requested. %s", + num_delivered, + args.max_samples, + get_sample_shortfall_hint(args, total_examples), + ) + if not args.use_attention_mask: # remove attention_mask tokenized_inputs.pop("attention_mask") @@ -375,6 +393,57 @@ def text_gen_pre_process(dataset, tokenizer, all_kwargs): return ClassificationDataset(hf_dataset, "labels", max_samples=args.max_samples) +def get_sample_shortfall_hint(args: TextGenParams, total_examples: int) -> str: + """Explain why a strategy delivered fewer samples than max_samples and how to get more. + + The cause and the effective remedies are different for each strategy, so the hint is strategy specific. + """ + if args.strategy == TextGenStrategy.JOIN_RANDOM: + # samples are drawn from random start positions and may overlap, so the corpus only ever needs + # max_seq_len tokens in total, i.e. max_samples does not change the requirement + return ( + f"The '{args.strategy}' strategy samples random start positions, so it only needs" + f" max_seq_len={args.max_seq_len} tokens after a start position and max_samples does not increase that" + f" requirement. All random_retries={args.random_retries} attempts for the missing samples started too" + " close to the end of the split. Use a larger dataset split, lower max_seq_len, or raise random_retries." + ) + + if "join" in args.strategy: + # JOIN and JOIN_SLIDING_WINDOW walk the joined corpus in blocks of max_seq_len taken every `step` tokens + step = args.stride if args.strategy == TextGenStrategy.JOIN_SLIDING_WINDOW else args.max_seq_len + required = args.max_seq_len + (args.max_samples - 1) * step + remedy = "Use a larger dataset split, or lower max_samples/max_seq_len." + if args.strategy == TextGenStrategy.JOIN_SLIDING_WINDOW: + # step == stride here, so a larger stride *increases* the required corpus length; lowering + # stride (more overlap between windows) is what actually reduces it. + remedy = "Use a larger dataset split, lower stride (more overlap), or lower max_samples/max_seq_len." + return ( + f"The '{args.strategy}' strategy takes a block of max_seq_len={args.max_seq_len} tokens every" + f" step={step} tokens of the joined corpus, so it needs at least max_seq_len + (max_samples - 1) * step" + f" = {required} tokens and the split provides fewer. {remedy}" + ) + + if args.strategy == TextGenStrategy.LINE_BY_LINE_RANDOM: + return ( + f"The '{args.strategy}' strategy produces at most one sample per sampled row and no acceptable row was" + f" found for the missing samples within random_retries={args.random_retries} attempts" + f" (drop_short_sequences={args.drop_short_sequences} discards rows shorter than" + f" max_seq_len={args.max_seq_len}). Use a larger dataset split, lower max_seq_len, or raise" + " random_retries." + ) + + # LINE_BY_LINE + dropped_note = ( + f" and drop_short_sequences=True discards rows shorter than max_seq_len={args.max_seq_len}" + if args.drop_short_sequences + else "" + ) + return ( + f"The '{args.strategy}' strategy produces at most one sample per non-empty row, the split provides" + f" {total_examples} non-empty rows{dropped_note}. Use a larger dataset split, or lower max_samples." + ) + + def get_text( example: dict[str, str], chat_template: Optional[Union[bool, str]] = None, diff --git a/olive/model/handler/mixin/hf.py b/olive/model/handler/mixin/hf.py index a5e51c162b..9c55350798 100644 --- a/olive/model/handler/mixin/hf.py +++ b/olive/model/handler/mixin/hf.py @@ -3,6 +3,8 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import logging +import shutil +import tempfile from pathlib import Path from typing import TYPE_CHECKING, Any, Optional, Union @@ -117,10 +119,83 @@ def save_metadata(self, output_dir: str, exclude_load_keys: Optional[list[str]] tokenizer_filepaths = save_tokenizer(self.get_hf_tokenizer(), output_dir, **kwargs) saved_filepaths.extend([fp for fp in tokenizer_filepaths if Path(fp).exists()]) + # save processor / image processor; per-file, don't overwrite anything that already exists + # (see ``_copy_missing_files``). This writes preprocessor_config.json (and any image + # processor files) so downstream tools that load from this output_dir (e.g. mobius's + # AutoProcessor.from_pretrained) get the model's real preprocessing config instead of + # silently falling back to defaults. Only applicable to multimodal models (e.g. VL + # checkpoints); text-only models have no processor and are already covered by the + # tokenizer save above. + # + # Note: unlike the tokenizer save above, this is not gated on a single sentinel file + # (e.g. "does preprocessor_config.json already exist?") -- a processor can emit several + # files (preprocessor_config.json, chat_template.json, ...), and an earlier step may have + # saved only some of them. Always calling ``_save_processor`` lets ``_copy_missing_files`` + # fill in whichever files are still missing, file by file. + saved_filepaths.extend(self._save_processor(output_dir, exclude_load_keys=exclude_load_keys, **kwargs)) + logger.debug("Save metadata files to %s: %s", output_dir, saved_filepaths) return saved_filepaths + def _save_processor(self, output_dir: Path, exclude_load_keys: Optional[list[str]] = None, **kwargs) -> list[str]: + """Save the model's processor files (preprocessor_config.json, ...) to output_dir. + + Never overwrites a file that already exists in ``output_dir``: + ``ProcessorMixin.save_pretrained`` also re-saves the processor's tokenizer, which would + clobber a tokenizer that an earlier step intentionally customized and saved. The + processor is therefore saved to a temporary directory first and only its new files are + copied over. + + :param output_dir: output directory to save the processor files in + :param exclude_load_keys: list of keys to exclude from load_kwargs + :param kwargs: additional keyword arguments to pass to `save_pretrained` method + :return: list of file paths that were written + """ + from transformers import AutoProcessor + from transformers.tokenization_utils_base import PreTrainedTokenizerBase + + try: + processor = AutoProcessor.from_pretrained( + self.model_name_or_path, **self.get_load_kwargs(exclude_load_keys=exclude_load_keys) + ) + except Exception as e: # pylint: disable=broad-except + # unexpected: loading failed for a reason other than "this model has no processor" + # (network / auth / incompatible config). Surface it -- VL models genuinely need + # preprocessor_config.json downstream. + logger.warning("Failed to load processor for %r, no processor files saved: %s", self.model_name_or_path, e) + return [] + + if isinstance(processor, PreTrainedTokenizerBase): + # expected for text-only models: AutoProcessor falls back to returning the + # tokenizer, which the tokenizer save above already handled. + logger.debug("No processor for %r (AutoProcessor returned a tokenizer).", self.model_name_or_path) + return [] + + try: + with tempfile.TemporaryDirectory(prefix="olive_processor_") as temp_dir: + processor.save_pretrained(temp_dir, **kwargs) + return self._copy_missing_files(Path(temp_dir), output_dir) + except Exception as e: # pylint: disable=broad-except + logger.warning("Failed to save processor files for %r: %s", self.model_name_or_path, e) + return [] + + @staticmethod + def _copy_missing_files(src_dir: Path, output_dir: Path) -> list[str]: + """Copy files from src_dir into output_dir, keeping any file that already exists there.""" + copied_filepaths = [] + for src_path in sorted(src_dir.rglob("*")): + if not src_path.is_file(): + continue + dst_path = output_dir / src_path.relative_to(src_dir) + if dst_path.exists(): + logger.debug("Keeping existing %s instead of overwriting it.", dst_path) + continue + dst_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(src_path, dst_path) + copied_filepaths.append(str(dst_path)) + return copied_filepaths + def get_hf_io_config(self) -> Optional[dict[str, Any]]: """Get Io config for the model.""" return get_model_io_config( diff --git a/olive/passes/onnx/_genai_speech_worker.py b/olive/passes/onnx/_genai_speech_worker.py index 44bae1800c..c3e84fcfa0 100644 --- a/olive/passes/onnx/_genai_speech_worker.py +++ b/olive/passes/onnx/_genai_speech_worker.py @@ -284,7 +284,19 @@ def main(argv=None) -> int: Usage: ``python _genai_speech_worker.py ``. The request JSON provides ``genai_model_path``, ``wav_path``, ``max_new_tokens`` and ``first_n``. """ + # Enable faulthandler so native crashes (segfault, bus error, etc.) dump a Python-level + # traceback to stderr before the process dies. Without this a signal-killed subprocess + # produces no output and the parent has no clue where the crash happened. + import faulthandler + + faulthandler.enable() + logging.basicConfig(level=logging.WARNING) + # Flush stderr/stdout line-by-line so that any output written before a native crash + # is available to the parent process (buffered streams are lost on a hard kill). + sys.stderr.reconfigure(line_buffering=True) + sys.stdout.reconfigure(line_buffering=True) + argv = list(sys.argv[1:] if argv is None else argv) if len(argv) != 2: sys.stderr.write("usage: _genai_speech_worker.py \n") @@ -294,12 +306,19 @@ def main(argv=None) -> int: with open(request_path) as f: request = json.load(f) - result = generate( - genai_model_path=request["genai_model_path"], - wav_path=request["wav_path"], - max_new_tokens=int(request["max_new_tokens"]), - first_n=int(request["first_n"]), - ) + try: + result = generate( + genai_model_path=request["genai_model_path"], + wav_path=request["wav_path"], + max_new_tokens=int(request["max_new_tokens"]), + first_n=int(request["first_n"]), + ) + except Exception: + import traceback + + traceback.print_exc() + return 1 + with open(result_path, "w") as f: json.dump(result, f) return 0 diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py index b2df6c8705..82bb9dc13d 100644 --- a/olive/passes/onnx/common.py +++ b/olive/passes/onnx/common.py @@ -5,6 +5,7 @@ import json import logging import re +from collections.abc import Iterable from copy import deepcopy from pathlib import Path from typing import Any, Callable, Optional, Union @@ -792,40 +793,46 @@ def update_llm_pipeline_genai_config( def update_llm_pipeline_genai_config_gpu( - model: ONNXModelHandler, + model: Union[ONNXModelHandler, CompositeModelHandler], output_model_dir: Union[str, Path], - input_model_path: Union[str, Path], decoder_config_extra: Optional[dict[str, Any]] = None, -) -> ONNXModelHandler: + composite_components: Optional[Iterable[tuple[str, ONNXModelHandler]]] = None, +) -> Union[ONNXModelHandler, CompositeModelHandler]: """Update the LLM pipeline in the model's genai_config.json file. - :param model: The model to update. + :param model: The model (single or composite) to update. + :param output_model_dir: Directory where the updated genai_config.json should be written. :param decoder_config_extra: Extra configuration for the decoder. + :param composite_components: Optional iterable of (component_name, ONNXModelHandler) + used to build a multi-component pipeline. + :return: The same `model` object (with its directory now having updated genai_config.json). """ output_model_dir = Path(output_model_dir) - # update genai_config if it exists + additional_files = model.model_attributes["additional_files"] genai_config_path = None - genai_config_path = Path(input_model_path).parent / "genai_config.json" + for file_path in additional_files: + if Path(file_path).name == "genai_config.json": + genai_config_path = file_path + break - if genai_config_path.exists(): - genai_config_path = str(genai_config_path.resolve()) - else: + if not genai_config_path: return model with open(genai_config_path) as f: genai_config = json.load(f) - # update model_type genai_config["model"]["type"] = "decoder-pipeline" - # Update the provider_options list - provider_option = {"qnn": {"backend_type": "gpu"}} - genai_config["model"]["decoder"]["session_options"]["provider_options"] = [provider_option] + provider_option = {"qnn": {"backend_type": "gpu", "enable_dx12_shared_memory_allocator": "1"}} + decoder = genai_config["model"].setdefault("decoder", {}) + session_opts = decoder.setdefault("session_options", {}) + session_opts["provider_options"] = [provider_option] # update decoder config decoder_config = genai_config["model"]["decoder"] decoder_config.get("sliding_window", {}).pop("slide_inputs", None) + for key, value in (decoder_config_extra or {}).items(): exisiting_value = decoder_config.get(key) if isinstance(exisiting_value, dict): @@ -835,13 +842,39 @@ def update_llm_pipeline_genai_config_gpu( else: decoder_config[key] = value - pipeline_config = {} - component_io_config = model.io_config - pipeline_config["model_onnx"] = { - "filename": Path(model.model_path).name, - "inputs": component_io_config["input_names"], - "outputs": component_io_config["output_names"], - } + # --- Build pipeline_config --- + pipeline_config: dict[str, Any] = {} + + if composite_components is None: + if not isinstance(model, ONNXModelHandler): + handlers = list(model.get_model_components()) + if not handlers: + return model + _, single_handler = handlers[0] + else: + single_handler = model + + component_io_config = single_handler.io_config + component_key = Path(single_handler.model_path).stem + pipeline_config[component_key] = { + "filename": Path(single_handler.model_path).name, + "inputs": component_io_config["input_names"], + "outputs": component_io_config["output_names"], + } + + else: + # Composite case: one entry per component + for comp_name, comp_handler in composite_components: + component_io_config = comp_handler.io_config + pipeline_config[comp_name] = { + "filename": Path(comp_handler.model_path).name, + "inputs": component_io_config["input_names"], + "outputs": component_io_config["output_names"], + } + if comp_name.endswith("decode"): + pipeline_config[comp_name]["run_on_prompt"] = False + else: + pipeline_config[comp_name]["run_on_token_gen"] = False decoder_config["pipeline"] = [pipeline_config] @@ -849,40 +882,65 @@ def update_llm_pipeline_genai_config_gpu( new_genai_config_path = output_model_dir / "genai_config.json" with new_genai_config_path.open("w") as f: json.dump(genai_config, f, indent=4) + additional_files.remove(genai_config_path) + additional_files.append(str(new_genai_config_path)) return model def update_llm_pipeline_genai_config_gpu_ctxbin( - model_path: Union[str, Path], + model: Union[ONNXModelHandler, CompositeModelHandler], + output_model_path: Union[str, Path], ) -> None: - """Update the filename fields in the model's genai_config.json file from 'model' to 'model_ctx'. + """Update the genai_config.json entry for one context binary component. - The genai_config.json file is updated in place in the model's directory. - :param model_path: Path to the model file. + :param model: Source model is used to locate and update genai_config.json. + :param output_model_path: Path to the context binary output file. """ - # Find genai_config in the model's directory - model_dir = Path(model_path).parent - genai_config_path = model_dir / "genai_config.json" + output_model_path = Path(output_model_path) + + # Extract additional_files from model -- same as update_llm_pipeline_genai_config_gpu + additional_files = model.model_attributes["additional_files"] + genai_config_path = None + for file_path in additional_files: + if Path(file_path).name == "genai_config.json": + genai_config_path = file_path + break + + if not genai_config_path: + return - if not genai_config_path.exists(): + ctx_stem = output_model_path.stem + if not ctx_stem.endswith("_ctx"): return + src_stem = ctx_stem[: -len("_ctx")] + src_filename = f"{src_stem}.onnx" + ctx_filename = f"{ctx_stem}.onnx" with open(genai_config_path) as f: genai_config = json.load(f) - # Update decoder filename to 'model_ctx' - if "decoder" in genai_config.get("model", {}): - if "filename" in genai_config["model"]["decoder"]: - genai_config["model"]["decoder"]["filename"] = "model/model_ctx.onnx" + decoder = genai_config.get("model", {}).get("decoder", {}) - # Update filename in pipeline configuration - decoder_config = genai_config["model"]["decoder"] - if "pipeline" in decoder_config and isinstance(decoder_config["pipeline"], list): - for pipeline_item in decoder_config["pipeline"]: - if "model_onnx" in pipeline_item and "filename" in pipeline_item["model_onnx"]: - pipeline_item["model_onnx"]["filename"] = "model/model_ctx.onnx" + # Update top-level decoder.filename if it points to this model + if decoder.get("filename") == src_filename: + decoder["filename"] = ctx_filename - # Save the updated genai_config back to the same location - with genai_config_path.open("w") as f: + # Update the single matching pipeline entry + for pipeline_item in decoder.get("pipeline", []): + if not isinstance(pipeline_item, dict): + continue + for comp_name in list(pipeline_item.keys()): + comp = pipeline_item[comp_name] + if isinstance(comp, dict) and comp.get("filename") == src_filename: + comp["filename"] = ctx_filename + if comp_name == src_stem: + pipeline_item[ctx_stem] = pipeline_item.pop(comp_name) + break # only one entry matches per call + + # Save to output dir and update additional_files pointer. + new_genai_config_path = output_model_path.parent / "genai_config.json" + with new_genai_config_path.open("w") as f: json.dump(genai_config, f, indent=4) + additional_files.remove(genai_config_path) + additional_files.append(str(new_genai_config_path)) diff --git a/olive/passes/onnx/context_binary.py b/olive/passes/onnx/context_binary.py index ba7fc433b8..914e2ea28f 100644 --- a/olive/passes/onnx/context_binary.py +++ b/olive/passes/onnx/context_binary.py @@ -61,6 +61,14 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon default_value=False, description="Whether to disable CPU fallback.", ), + "update_genai_config": PassConfigParam( + type_=bool, + default_value=False, + description=( + "Whether to update the genai_config.json pipeline entry for the generated context binary." + " Only applicable to QNN GPU." + ), + ), } @staticmethod @@ -117,8 +125,20 @@ def _run_single_target( "session_options": config.session_options, "embed_context": config.embed_context, "disable_cpu_fallback": config.disable_cpu_fallback, + "update_genai_config": config.update_genai_config, + "model": model, } + if ( + config.weight_sharing + and self.accelerator_spec.execution_provider == ExecutionProvider.QNNExecutionProvider + and str(self.accelerator_spec.accelerator_type).lower() == "gpu" + ): + assert isinstance(model, CompositeModelHandler), ( + "weight_sharing for QNN GPU is only supported for composite prefill/decode models generated by " + "StaticLLM (model must be a CompositeModelHandler)." + ) + if isinstance(model, ONNXModelHandler): return self._generate_context_binary( model_path=model.model_path, @@ -187,6 +207,11 @@ def process_context_iterator(component_models, llm_pipeline, output_dir): device=self.accelerator_spec.accelerator_type, generate_kwargs=generate_kwargs, weight_sharing=config.weight_sharing, + shared_bin_stem=( + "model_ctx" + if config.weight_sharing and str(self.accelerator_spec.accelerator_type).lower() == "gpu" + else None + ), ) return CompositeModelHandler( list(new_component_models.values()), @@ -203,6 +228,7 @@ def _generate_composite_binaries( output_model_dir: str, generate_kwargs: dict[str, str], weight_sharing: bool = False, + shared_bin_stem: Optional[str] = None, ) -> dict[str, ONNXModelHandler]: """Generate context binary for each model in the composite model. @@ -210,24 +236,37 @@ def _generate_composite_binaries( :param output_model_dir: Directory to save the output model files. :param generate_kwargs: Additional arguments for the context binary generation. :param weight_sharing: Whether to enable weight sharing between the models. + :param shared_bin_stem: weight_sharing bin is generated under this neutral stem instead of its own logical name for QNN GPU. :return: Map of model names to ONNXModelHandler for the generated context binaries. """ + output_model_dir = Path(output_model_dir) new_models = {} for idx, (model_name, model_path) in enumerate(model_paths_map.items()): generate_kwargs = deepcopy(generate_kwargs) + generation_name = model_name if weight_sharing: generate_kwargs["share_ep_contexts"] = True generate_kwargs["stop_share_ep_contexts"] = idx == len(model_paths_map) - 1 generate_kwargs["embed_context"] = False generate_kwargs["ignore_missing_cb_bin"] = idx != len(model_paths_map) - 1 + if idx == 0 and shared_bin_stem: + generation_name = shared_bin_stem - new_models[model_name] = cls._generate_context_binary( + handler = cls._generate_context_binary( model_path=model_path, - output_model_path=Path(output_model_dir) / f"{model_name}.onnx", + output_model_path=output_model_dir / f"{generation_name}.onnx", device=device, + logical_name=model_name, **generate_kwargs, ) + if generation_name != model_name: + renamed_path = output_model_dir / f"{model_name}.onnx" + (output_model_dir / handler.onnx_file_name).replace(renamed_path) + handler = ONNXModelHandler(model_path=output_model_dir, onnx_file_name=renamed_path.name) + + new_models[model_name] = handler + return new_models @staticmethod @@ -243,6 +282,9 @@ def _generate_context_binary( share_ep_contexts: bool = False, stop_share_ep_contexts: bool = False, ignore_missing_cb_bin: bool = False, + update_genai_config: bool = False, + model: Optional[Union[ONNXModelHandler, CompositeModelHandler]] = None, + logical_name: Optional[str] = None, ) -> ONNXModelHandler: """Generate context binary for the model. @@ -255,6 +297,8 @@ def _generate_context_binary( :param share_ep_contexts: Whether to share EP contexts. :param stop_share_ep_contexts: Whether to stop sharing EP contexts. :param ignore_missing_cb_bin: Whether to ignore missing context binary files. + :param update_genai_config: Whether to update the genai_config.json pipeline entry (QNN GPU only). + :param logical_name: Component's logical name used to update genai_config.json. :return: ONNXModelHandler for the generated context binary. """ import onnxruntime as ort @@ -271,7 +315,13 @@ def _generate_context_binary( if execution_provider == ExecutionProvider.QNNExecutionProvider: if str(device).lower() == "gpu": provider_options["backend_path"] = "libQnnGpu.so" if platform.system() == "Linux" else "QnnGpu.dll" - update_llm_pipeline_genai_config_gpu_ctxbin(model_path) + if update_genai_config: + ctxbin_path = ( + Path(output_model_path).with_name(f"{logical_name}.onnx") + if logical_name + else Path(output_model_path) + ) + update_llm_pipeline_genai_config_gpu_ctxbin(model, ctxbin_path) else: provider_options["backend_path"] = "libQnnHtp.so" if platform.system() == "Linux" else "QnnHtp.dll" if share_ep_contexts: diff --git a/olive/passes/onnx/conversion.py b/olive/passes/onnx/conversion.py index ad47cb99fa..9b2575dd6c 100644 --- a/olive/passes/onnx/conversion.py +++ b/olive/passes/onnx/conversion.py @@ -3,11 +3,13 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import collections +import contextlib import functools import inspect import logging import multiprocessing import tempfile +import threading from copy import deepcopy from pathlib import Path from typing import Any, Optional, Union @@ -43,6 +45,16 @@ logger = logging.getLogger(__name__) +class _DynamicLayerPatchState: + def __init__(self): + self.lock = threading.Lock() + self.depth = 0 + self.original_lazy_initialization = None + + +_DYNAMIC_LAYER_PATCH_STATE = _DynamicLayerPatchState() + + def _torch_is_older_than(version_str: str) -> bool: torch_version = version.parse(torch.__version__).release return torch_version < version.parse(version_str).release @@ -105,8 +117,28 @@ def _unflatten_dynamic_cache(values, context: torch.utils._pytree.Context): return cache +def _patched_dynamic_layer_lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor): + self.dtype, self.device = key_states.dtype, key_states.device + + def empty_cache(states: torch.Tensor): + like = torch.narrow(states, dim=-2, start=0, length=0) + if hasattr(states, "fake_mode"): + with states.fake_mode: + return torch.empty_like(like, dtype=states.dtype, device=states.device) + return torch.empty_like(like, dtype=states.dtype, device=states.device) + + self.keys = empty_cache(key_states) + self.values = empty_cache(value_states) + self.is_initialized = True + + +@contextlib.contextmanager def _patch_dynamic_layer_for_export(): - """Patch DynamicLayer.lazy_initialization for torch.export compatibility (transformers >= 5.0). + """Temporarily patch DynamicLayer.lazy_initialization for torch.export compatibility (transformers >= 5.0). + + Use as a context manager; the original method is restored on exit, including when an exception + is raised inside the `with` block. No-op if the installed transformers version doesn't define + DynamicLayer.lazy_initialization. The original uses torch.tensor([]) which creates a 1D empty tensor (shape [0]). torch.export needs consistent tensor ranks, so we use torch.narrow + torch.empty_like @@ -114,23 +146,31 @@ def _patch_dynamic_layer_for_export(): """ from transformers.cache_utils import DynamicLayer - if not hasattr(DynamicLayer, "lazy_initialization"): - return - - def patched_lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor = None): - self.dtype, self.device = key_states.dtype, key_states.device - like = torch.narrow(key_states, dim=-2, start=0, length=0) - if hasattr(key_states, "fake_mode"): - with key_states.fake_mode: - self.keys = torch.empty_like(like, dtype=self.dtype, device=self.device) - self.values = torch.empty_like(like, dtype=self.dtype, device=self.device) + with _DYNAMIC_LAYER_PATCH_STATE.lock: + if not hasattr(DynamicLayer, "lazy_initialization"): + patch_applied = False else: - self.keys = torch.empty_like(like, dtype=self.dtype, device=self.device) - self.values = torch.empty_like(like, dtype=self.dtype, device=self.device) - self.is_initialized = True + patch_applied = True + if _DYNAMIC_LAYER_PATCH_STATE.depth == 0: + _DYNAMIC_LAYER_PATCH_STATE.original_lazy_initialization = DynamicLayer.lazy_initialization + DynamicLayer.lazy_initialization = _patched_dynamic_layer_lazy_initialization + logger.debug("Patched DynamicLayer.lazy_initialization for torch.export compatibility.") + _DYNAMIC_LAYER_PATCH_STATE.depth += 1 + + if not patch_applied: + yield + return - DynamicLayer.lazy_initialization = patched_lazy_initialization - logger.debug("Patched DynamicLayer.lazy_initialization for torch.export compatibility.") + try: + yield + finally: + with _DYNAMIC_LAYER_PATCH_STATE.lock: + _DYNAMIC_LAYER_PATCH_STATE.depth -= 1 + if _DYNAMIC_LAYER_PATCH_STATE.depth == 0: + if getattr(DynamicLayer, "lazy_initialization", None) is _patched_dynamic_layer_lazy_initialization: + DynamicLayer.lazy_initialization = _DYNAMIC_LAYER_PATCH_STATE.original_lazy_initialization + logger.debug("Restored DynamicLayer.lazy_initialization after torch.export.") + _DYNAMIC_LAYER_PATCH_STATE.original_lazy_initialization = None def _convert_past_key_values_to_dynamic_cache(dummy_kwargs: dict, config=None) -> dict: @@ -338,51 +378,53 @@ def _export_pytorch_model( dummy_kwargs = {} dummy_inputs = tuple(dummy_inputs) - # Apply patches for DynamicCache / past_key_values compatibility - if version.parse(transformers.__version__) >= version.parse("5.0"): + transformers_5_or_later = version.parse(transformers.__version__) >= version.parse("5.0") + if transformers_5_or_later: # transformers >= 5.0: DynamicCache refactored to use DynamicLayer - _register_dynamic_cache_export_support() - _patch_dynamic_layer_for_export() - model_config = getattr(pytorch_model, "config", None) - dummy_kwargs = _convert_past_key_values_to_dynamic_cache(dummy_kwargs, config=model_config) - if io_config.dynamic_shapes: - io_config.dynamic_shapes = _convert_dynamic_shapes_for_dynamic_cache(io_config.dynamic_shapes) - else: - # transformers < 5.0: patch forward to convert list <-> DynamicCache - _patch_model_if_necessary(pytorch_model) - - # NOTE: Usually validation is done in io_config.py, but because - # dynamic_shapes has nested complexity, and it can't be validated multiple - # times like others, we validate it here. - io_config.dynamic_shapes, dummy_inputs, dummy_kwargs = _validate_dynamic_shapes( - io_config.dynamic_shapes, dummy_inputs, dummy_kwargs, pytorch_model - ) - # torch.export requires strict type match between inputs and dynamic_shapes; - # _validate_dynamic_shapes may return OrderedDict, so convert back to plain dict - if isinstance(io_config.dynamic_shapes, collections.OrderedDict): - io_config.dynamic_shapes = dict(io_config.dynamic_shapes) - if isinstance(dummy_kwargs, collections.OrderedDict): - dummy_kwargs = dict(dummy_kwargs) - - # When dynamo=True, PyTorch prefers dynamic_shapes over dynamic_axes. - # If dynamic_shapes is None and fallback is enabled, don't pass dynamic_axes - # to avoid conversion errors. The fallback path will handle dynamic axes. - dynamic_axes_for_export = io_config.dynamic_axes if io_config.dynamic_shapes else None - - onnx_program = torch.onnx.export( # pylint: disable=unexpected-keyword-arg,no-value-for-parameter - pytorch_model, - dummy_inputs, - kwargs=dummy_kwargs, - opset_version=config.target_opset, - input_names=io_config.input_names, - output_names=io_config.output_names, - dynamic_axes=dynamic_axes_for_export, - dynamic_shapes=io_config.dynamic_shapes, - dynamo=True, - optimize=config.optimize, - report=logger.isEnabledFor(logging.DEBUG), - ) + + patch_context = _patch_dynamic_layer_for_export() if transformers_5_or_later else contextlib.nullcontext() + with patch_context: + if transformers_5_or_later: + model_config = getattr(pytorch_model, "config", None) + dummy_kwargs = _convert_past_key_values_to_dynamic_cache(dummy_kwargs, config=model_config) + if io_config.dynamic_shapes: + io_config.dynamic_shapes = _convert_dynamic_shapes_for_dynamic_cache(io_config.dynamic_shapes) + else: + # transformers < 5.0: patch forward to convert list <-> DynamicCache + _patch_model_if_necessary(pytorch_model) + + # NOTE: Usually validation is done in io_config.py, but because + # dynamic_shapes has nested complexity, and it can't be validated multiple + # times like others, we validate it here. + io_config.dynamic_shapes, dummy_inputs, dummy_kwargs = _validate_dynamic_shapes( + io_config.dynamic_shapes, dummy_inputs, dummy_kwargs, pytorch_model + ) + # torch.export requires strict type match between inputs and dynamic_shapes; + # _validate_dynamic_shapes may return OrderedDict, so convert back to plain dict + if isinstance(io_config.dynamic_shapes, collections.OrderedDict): + io_config.dynamic_shapes = dict(io_config.dynamic_shapes) + if isinstance(dummy_kwargs, collections.OrderedDict): + dummy_kwargs = dict(dummy_kwargs) + + # When dynamo=True, PyTorch prefers dynamic_shapes over dynamic_axes. + # If dynamic_shapes is None and fallback is enabled, don't pass dynamic_axes + # to avoid conversion errors. The fallback path will handle dynamic axes. + dynamic_axes_for_export = io_config.dynamic_axes if io_config.dynamic_shapes else None + + onnx_program = torch.onnx.export( # pylint: disable=unexpected-keyword-arg,no-value-for-parameter + pytorch_model, + dummy_inputs, + kwargs=dummy_kwargs, + opset_version=config.target_opset, + input_names=io_config.input_names, + output_names=io_config.output_names, + dynamic_axes=dynamic_axes_for_export, + dynamic_shapes=io_config.dynamic_shapes, + dynamo=True, + optimize=config.optimize, + report=logger.isEnabledFor(logging.DEBUG), + ) assert onnx_program is not None model = onnx_program.model # We can run load_to_model on all models: If the model is from dynamo=True, diff --git a/olive/passes/onnx/discrepancy_check.py b/olive/passes/onnx/discrepancy_check.py index 91fa7e008b..abb1c1a5be 100644 --- a/olive/passes/onnx/discrepancy_check.py +++ b/olive/passes/onnx/discrepancy_check.py @@ -15,6 +15,7 @@ import numpy as np import onnx +from olive.common.onnx_io import get_genai_decoder_config from olive.data.config import DataConfig from olive.hardware import AcceleratorSpec from olive.hardware.accelerator import Device @@ -63,7 +64,9 @@ def _infer_shape(dynamic_shape, known_values=None): "total_sequence_length": 8, } if known_values: - default_values.update(known_values) + # Shapes mix symbolic names and concrete ints, so only keep the symbolic entries; + # otherwise the error message below would compare ints against strings. + default_values.update({key: value for key, value in known_values.items() if isinstance(key, str)}) inferred_shape = [] for dim in dynamic_shape: if isinstance(dim, int): @@ -128,6 +131,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. @@ -147,6 +202,307 @@ def _format_seconds(value: Optional[float]) -> str: return "n/a" if value is None else f"{value:.4f}s" +# --------------------------------------------------------------------------- +# Subprocess worker for per-component (encoder or decoder) discrepancy check. +# ONNX Runtime can segfault on certain model-builder outputs; running each +# component in its own subprocess lets the parent degrade gracefully and +# still measure the other component. +# --------------------------------------------------------------------------- +_SPEECH_COMPONENT_WORKER_SCRIPT = '''\ +"""Single-component discrepancy worker for speech models. + +Compares one ONNX component (encoder or decoder) against the HuggingFace reference. +Runs in its own subprocess so a native ORT crash only affects this component. + +Usage: python worker.py + +request.json fields: + - component: "encoder" or "decoder" + - onnx_path: path to the component .onnx file + - reference_model_path: path to the HuggingFace reference model directory + - encoder_outputs_path: (decoder only) path to .npz with saved encoder outputs +""" +import faulthandler +import json +import re +import sys +import traceback + +faulthandler.enable() +sys.stderr.reconfigure(line_buffering=True) +sys.stdout.reconfigure(line_buffering=True) + +import numpy as np + + +def _infer_shape(dynamic_shape): + default_values = { + "batch_size": 1, + "past_sequence_length": 0, + "sequence_length": 8, + "total_sequence_length": 8, + } + result = [] + for dim in dynamic_shape: + if isinstance(dim, int): + result.append(dim) + elif dim in default_values: + result.append(default_values[dim]) + else: + raise KeyError(f"Unsupported symbolic dimension: {dim}") + return tuple(result) + + +_ONNX_TYPE_TO_NP = { + "tensor(float)": np.float32, + "tensor(float16)": np.float16, + "tensor(double)": np.float64, + "tensor(int32)": np.int32, + "tensor(int64)": np.int64, + "tensor(int8)": np.int8, + "tensor(uint8)": np.uint8, + "tensor(bool)": np.bool_, +} + + +def _ort_type_to_numpy(ort_type_str): + return _ONNX_TYPE_TO_NP.get(ort_type_str, np.float32) + + +def compare_encoder(onnx_path, reference_model_path, save_outputs_path=None): + import onnxruntime as ort + import torch + from transformers import AutoModelForSpeechSeq2Seq + + enc_sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) + enc_inputs = {} + for inp in enc_sess.get_inputs(): + shape = _infer_shape([d if isinstance(d, int) else d for d in inp.shape]) + dtype = _ort_type_to_numpy(inp.type) + enc_inputs[inp.name] = np.random.randn(*shape).astype(dtype) + + onnx_enc_out = enc_sess.run(None, enc_inputs) + enc_output_names = [o.name for o in enc_sess.get_outputs()] + + ref_model = AutoModelForSpeechSeq2Seq.from_pretrained(reference_model_path) + ref_model.eval() + audio_key = list(enc_inputs.keys())[0] + hf_encoder = ref_model.get_encoder() + hf_input = torch.tensor(enc_inputs[audio_key], dtype=torch.float32) + with torch.no_grad(): + hf_enc_out = hf_encoder(hf_input) + + hf_hidden = hf_enc_out.last_hidden_state.to(torch.float64).cpu() + onnx_hidden = torch.as_tensor(onnx_enc_out[0]).to(torch.float64).cpu() + diff = torch.abs(hf_hidden - onnx_hidden) + + # Save encoder outputs + audio input for the decoder subprocess + if save_outputs_path: + save_dict = {"audio_input": enc_inputs[audio_key], "audio_key": audio_key} + for i, name in enumerate(enc_output_names): + save_dict[f"onnx_output_{name}"] = onnx_enc_out[i] + save_dict["output_names"] = np.array(enc_output_names, dtype=object) + np.savez(save_outputs_path, **save_dict) + + return { + "max_abs_error": float(torch.max(diff)), + "elements_above_0_1": int(torch.sum(diff > 0.1)), + "elements_above_0_01": int(torch.sum(diff > 0.01)), + "total_elements": int(diff.numel()), + "output_compared": "hidden_states", + } + + +def compare_decoder(onnx_path, reference_model_path, encoder_outputs_path=None): + import onnxruntime as ort + import torch + from transformers import AutoModelForSpeechSeq2Seq + from transformers.modeling_outputs import BaseModelOutput + + dec_sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) + + # Load the HF reference model first -- we may need its encoder. + ref_model = AutoModelForSpeechSeq2Seq.from_pretrained(reference_model_path) + ref_model.eval() + + # Build default decoder inputs from ONNX graph metadata. + dec_inputs = {} + for inp in dec_sess.get_inputs(): + shape = _infer_shape([d if isinstance(d, int) else d for d in inp.shape]) + dtype = _ort_type_to_numpy(inp.type) + if "past_key" in inp.name or "past_value" in inp.name: + dec_inputs[inp.name] = np.zeros(shape, dtype=dtype) + elif inp.name == "input_ids": + dec_inputs[inp.name] = np.full(shape, 50258, dtype=dtype) + else: + dec_inputs[inp.name] = np.random.randn(*shape).astype(dtype) + + # Load ONNX encoder outputs if available (produced by encoder subprocess). + audio_input = None + encoder_source = "onnx" + hf_encoder_hidden = None + if encoder_outputs_path: + saved = np.load(encoder_outputs_path, allow_pickle=True) + enc_output_names = list(saved["output_names"]) + audio_input = saved["audio_input"] + for name in enc_output_names: + if name in dec_inputs: + output_value = saved[f"onnx_output_{name}"] + dec_inputs[name] = output_value + if hf_encoder_hidden is None and output_value.ndim == 3: + hf_encoder_hidden = output_value + for name, value in dec_inputs.items(): + if "encoder_hidden_states" in name and isinstance(value, np.ndarray) and value.ndim == 3: + hf_encoder_hidden = value + break + else: + # ONNX encoder was unavailable (crashed). Run the HF encoder instead + # so both ONNX decoder and HF decoder receive identical encoder outputs. + # This isolates the decoder discrepancy from the encoder failure. + encoder_source = "hf_fallback" + audio_input = np.random.randn(1, 80, 3000).astype(np.float32) + hf_audio = torch.tensor(audio_input, dtype=torch.float32) + with torch.no_grad(): + hf_enc_out = ref_model.get_encoder()(hf_audio) + hf_hidden = hf_enc_out.last_hidden_state.cpu().numpy() + hf_encoder_hidden = hf_hidden + + # Wire HF encoder hidden_states into the ONNX decoder cross-attention + # inputs. Model-builder names these after the original encoder output + # names (e.g. "encoder_hidden_states") or cross-attention past_key/ + # past_value tensors whose shape matches [batch, heads, seq, head_dim]. + for inp in dec_sess.get_inputs(): + name = inp.name + shape = _infer_shape([d if isinstance(d, int) else d for d in inp.shape]) + if "encoder_hidden_states" in name: + # Encoder hidden states -- broadcast / pad to expected shape. + dec_inputs[name] = hf_hidden.astype(_ort_type_to_numpy(inp.type)) + elif "cross" in name and ("key" in name or "value" in name): + # Cross-attention KV caches -- derive from HF encoder outputs. + # The ONNX decoder expects [batch, heads, encoder_seq, head_dim]. + # We fill with the HF encoder hidden_states projected through the + # *reference* decoder's cross-attention layers. + pass # handled below + + # Project HF encoder hidden states through cross-attention K/V + # projections of the HF decoder so the ONNX decoder receives proper + # cross-attention inputs. + hf_hidden_t = torch.tensor(hf_hidden, dtype=torch.float32) + hf_decoder = ref_model.get_decoder() + + # Build a mapping from ONNX cross-attention input name → layer index. + # Model-builder uses names like "past_key_cross_0" or + # "past_key_values.0.cross.key" — extract the layer index from either. + cross_inputs = {} # layer_idx -> {"key": input_meta, "value": input_meta} + for inp in dec_sess.get_inputs(): + name = inp.name + if "cross" not in name: + continue + is_key = "key" in name + is_value = "value" in name + if not (is_key or is_value): + continue + # Try "past_key_cross_0" style (layer index as trailing digits) + m = re.search(r"_(\\d+)$", name) + if not m: + # Try "past_key_values.0.cross.key" style + m = re.search(r"\\.(\\d+)\\.", name) + if m: + idx = int(m.group(1)) + cross_inputs.setdefault(idx, {}) + cross_inputs[idx]["key" if is_key else "value"] = inp + + for layer_idx, layer in enumerate(hf_decoder.layers): + if layer_idx not in cross_inputs: + continue + cross_attn = layer.encoder_attn + num_heads = cross_attn.num_heads + head_dim = cross_attn.head_dim + with torch.no_grad(): + k_proj = cross_attn.k_proj(hf_hidden_t) + v_proj = cross_attn.v_proj(hf_hidden_t) + batch = k_proj.shape[0] + enc_seq = k_proj.shape[1] + k_proj = k_proj.reshape(batch, enc_seq, num_heads, head_dim).permute(0, 2, 1, 3).cpu().numpy() + v_proj = v_proj.reshape(batch, enc_seq, num_heads, head_dim).permute(0, 2, 1, 3).cpu().numpy() + for kv_type, proj_data in [("key", k_proj), ("value", v_proj)]: + if kv_type in cross_inputs[layer_idx]: + inp_meta = cross_inputs[layer_idx][kv_type] + dec_inputs[inp_meta.name] = proj_data.astype(_ort_type_to_numpy(inp_meta.type)) + + onnx_dec_out = dec_sess.run(None, dec_inputs) + + # HF full model for logits comparison. + # Use the same encoder hidden states as the ONNX decoder so discrepancy + # reflects decoder behavior rather than encoder differences. + input_ids = torch.tensor(dec_inputs["input_ids"], dtype=torch.long) + if hf_encoder_hidden is None: + if audio_input is None: + audio_input = np.random.randn(1, 80, 3000).astype(np.float32) + hf_audio = torch.tensor(audio_input, dtype=torch.float32) + with torch.no_grad(): + hf_out = ref_model(input_features=hf_audio, decoder_input_ids=input_ids) + else: + encoder_outputs = BaseModelOutput(last_hidden_state=torch.tensor(hf_encoder_hidden, dtype=torch.float32)) + with torch.no_grad(): + hf_out = ref_model(encoder_outputs=encoder_outputs, decoder_input_ids=input_ids) + hf_logits = hf_out.logits.to(torch.float64).cpu() + + dec_output_names = [o.name for o in dec_sess.get_outputs()] + logits_idx = dec_output_names.index("logits") if "logits" in dec_output_names else 0 + onnx_logits = torch.as_tensor(onnx_dec_out[logits_idx]).to(torch.float64).cpu() + + diff = torch.abs(hf_logits - onnx_logits) + result = { + "max_abs_error": float(torch.max(diff)), + "elements_above_0_1": int(torch.sum(diff > 0.1)), + "elements_above_0_01": int(torch.sum(diff > 0.01)), + "total_elements": int(diff.numel()), + "output_compared": "logits", + } + if encoder_source == "hf_fallback": + result["encoder_source"] = "hf_fallback" + result["note"] = ( + "ONNX encoder unavailable; used HF encoder outputs as " + "cross-attention input to isolate decoder discrepancy." + ) + return result + + +if __name__ == "__main__": + if len(sys.argv) != 3: + sys.stderr.write("usage: component_worker.py \\n") + sys.exit(2) + + with open(sys.argv[1]) as f: + request = json.load(f) + + try: + component = request["component"] + if component == "encoder": + result = compare_encoder( + request["onnx_path"], + request["reference_model_path"], + save_outputs_path=request.get("save_outputs_path"), + ) + elif component == "decoder": + result = compare_decoder( + request["onnx_path"], + request["reference_model_path"], + encoder_outputs_path=request.get("encoder_outputs_path"), + ) + else: + sys.stderr.write(f"Unknown component: {component}\\n") + sys.exit(2) + except Exception: + traceback.print_exc() + sys.exit(1) + + with open(sys.argv[2], "w") as f: + json.dump(result, f) +''' + + # --------------------------------------------------------------------------- # Helper script executed inside the ``llama_env`` virtual environment. # All llama-cpp-python / gguf imports are intentionally isolated to this @@ -432,7 +788,18 @@ def _run_for_config( ) _, torch_device = self._resolve_devices() ref_model = self._cast_reference_model(ref_model, None, torch_device) + + # Measure per-component (encoder/decoder) discrepancies before generation comparison + component_disc = self._compute_speech_component_discrepancy(model, ref_path) results = self._run_speech_generation_comparison(model, config, ref_model, ref_path) + if component_disc: + results["component_discrepancy"] = component_disc.get("components", {}) + if "max_abs_error" in component_disc: + results["max_abs_error"] = component_disc["max_abs_error"] + results["elements_above_0_1"] = component_disc["elements_above_0_1"] + results["elements_above_0_01"] = component_disc["elements_above_0_01"] + results["total_elements"] = component_disc["total_elements"] + self._save_results(model, results, report_dir) return model @@ -499,7 +866,7 @@ def _load_or_make_image(config): """Return a PIL image for the VLM generation comparison. Loads the image from ``config.generate_image_path`` when the path points to an existing - file; otherwise returns a small synthetic solid-colour image. The synthetic image is + file; otherwise returns a small synthetic solid-color image. The synthetic image is intentionally tiny (32x32) so that it is cheap to process even with a real visual encoder. """ try: @@ -513,6 +880,173 @@ def _load_or_make_image(config): return Image.open(image_path).convert("RGB") return Image.new("RGB", (32, 32), color=(128, 128, 128)) + def _compute_speech_component_discrepancy(self, model, ref_model_path): + """Compare encoder and decoder outputs between HuggingFace and ONNX for speech models. + + Each component runs in its own subprocess so a native ORT crash (segfault) in one + component does not prevent the other from being measured. + + Returns a dict with per-component discrepancy metrics and aggregate max_abs_error. + """ + from olive.model import CompositeModelHandler + + if not isinstance(model, CompositeModelHandler): + logger.warning("Speech component discrepancy requires a CompositeModelHandler; skipping.") + return {} + + components = dict(model.get_model_components()) + # Component names may be bare ("encoder") or include the extension ("encoder.onnx") + encoder_handler = None + decoder_handler = None + for name, handler in components.items(): + if "encoder" in name.lower(): + encoder_handler = handler + elif "decoder" in name.lower(): + decoder_handler = handler + if not encoder_handler or not decoder_handler: + logger.warning( + "Could not find encoder/decoder components (found: %s); skipping component discrepancy.", + list(components.keys()), + ) + return {} + + encoder_path = encoder_handler.model_path + decoder_path = decoder_handler.model_path + if not ref_model_path or not Path(ref_model_path).is_dir(): + logger.warning( + "Reference model path %r is not a local directory; skipping component discrepancy.", ref_model_path + ) + return {} + + component_results = {} + + with tempfile.TemporaryDirectory(prefix="olive_component_disc_") as work_dir: + work_path = Path(work_dir) + script_path = work_path / "component_worker.py" + script_path.write_text(_SPEECH_COMPONENT_WORKER_SCRIPT) + encoder_outputs_path = str(work_path / "encoder_outputs.npz") + + # --- Encoder subprocess --- + component_results["encoder"] = self._run_component_subprocess( + script_path, + work_path, + { + "component": "encoder", + "onnx_path": str(encoder_path), + "reference_model_path": str(ref_model_path), + "save_outputs_path": encoder_outputs_path, + }, + ) + + # --- Decoder subprocess --- + decoder_request = { + "component": "decoder", + "onnx_path": str(decoder_path), + "reference_model_path": str(ref_model_path), + } + if Path(encoder_outputs_path).is_file(): + decoder_request["encoder_outputs_path"] = encoder_outputs_path + + component_results["decoder"] = self._run_component_subprocess( + script_path, + work_path, + decoder_request, + ) + + for comp_name, comp_res in component_results.items(): + if "error" not in comp_res: + logger.info( + "OnnxDiscrepancyCheck %s: max_abs_error=%.6f, elements_above_0.1=%d/%d, " + "elements_above_0.01=%d/%d (%s)", + comp_name, + comp_res["max_abs_error"], + comp_res["elements_above_0_1"], + comp_res["total_elements"], + comp_res["elements_above_0_01"], + comp_res["total_elements"], + comp_res.get("output_compared", "?"), + ) + else: + logger.warning("OnnxDiscrepancyCheck %s comparison failed: %s", comp_name, comp_res["error"]) + + # Aggregate + max_errors = [] + total_above_0_1 = 0 + total_above_0_01 = 0 + total_elements = 0 + for comp_res in component_results.values(): + if "error" not in comp_res: + max_errors.append(comp_res["max_abs_error"]) + total_above_0_1 += comp_res["elements_above_0_1"] + total_above_0_01 += comp_res["elements_above_0_01"] + total_elements += comp_res["total_elements"] + + aggregate = {} + if max_errors: + aggregate = { + "max_abs_error": max(max_errors), + "elements_above_0_1": total_above_0_1, + "elements_above_0_01": total_above_0_01, + "total_elements": total_elements, + } + + return {"components": component_results, **aggregate} + + @staticmethod + def _run_component_subprocess(script_path, work_dir, request): + """Run a single component discrepancy check in a subprocess.""" + request_path = work_dir / f"request_{request['component']}.json" + result_path = work_dir / f"result_{request['component']}.json" + request_path.write_text(json.dumps(request)) + + proc = subprocess.run( + [sys.executable, str(script_path), str(request_path), str(result_path)], + capture_output=True, + text=True, + check=False, + ) + + if proc.returncode != 0 or not result_path.is_file(): + stderr_raw = (proc.stderr or "").strip() + # Extract the faulthandler traceback (between "Fatal Python error" and "Extension modules") + # to avoid displaying the very long extension modules list. + stderr_lines = stderr_raw.split("\n") + traceback_lines = [] + in_traceback = False + for line in stderr_lines: + if "Fatal Python error" in line or "Traceback" in line: + in_traceback = True + if in_traceback and "Extension modules:" in line: + break + if in_traceback: + traceback_lines.append(line) + stderr_tail = "\n".join(traceback_lines).strip() if traceback_lines else stderr_raw[-500:] + + stdout_tail = (proc.stdout or "").strip()[-500:] + output_parts = [] + if stderr_tail: + output_parts.append(f"stderr: {stderr_tail}") + if stdout_tail: + output_parts.append(f"stdout: {stdout_tail}") + output_detail = "; ".join(output_parts) if output_parts else "(no output captured)" + return { + "error": ( + f"{request['component']} discrepancy subprocess failed " + f"(exit code {proc.returncode}). {output_detail}" + ) + } + + result_text = result_path.read_text() + try: + return json.loads(result_text) + except json.JSONDecodeError as e: + return { + "error": ( + f"{request['component']} discrepancy subprocess produced invalid JSON " + f"at {result_path}: {e}. output tail: {result_text[-500:]}" + ) + } + def _compute_final_metrics(self, results: dict) -> None: def _ratio(numer_key: str, denom_key: str, out_key: str) -> None: numer = results.get(numer_key) @@ -541,6 +1075,11 @@ def _prepare_dataloader(self, model: ONNXModelHandler): else: input_shapes = [] known = {} + # onnxruntime-genai exports the KV cache head size as the symbolic `kv_cache_dim`, + # so the concrete value has to be read back from the genai_config.json next to the model. + decoder_config = get_genai_decoder_config(model.model_path) + if decoder_config: + known["kv_cache_dim"] = decoder_config["head_size"] for shape in io_config.get("input_shapes"): new_shape = _infer_shape(shape, known) input_shapes.append(new_shape) @@ -698,7 +1237,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 @@ -971,6 +1510,10 @@ def _run_speech_generation_comparison(self, model, config, ref_model, ref_path): # GenAI-vs-transformers comparison could not be completed. if gen_results.get("genai_error"): results["genai_generation_error"] = gen_results["genai_error"] + # Only transformers metrics are available — mark the comparison as partial so the + # user knows the GenAI side could not be evaluated. + if results.get("status") != "failed": + results["status"] = "partial" logger.warning( "OnnxDiscrepancyCheck speech generation: reported transformers-only metrics; the " "GenAI comparison was skipped (%s). This is typically an onnxruntime-genai / " @@ -1139,12 +1682,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") @@ -1271,11 +1814,18 @@ def _run_genai_speech_subprocess(self, genai_model_path, audio, sample_rate, *, ) if proc.returncode != 0 or not result_path.is_file(): stderr_tail = (proc.stderr or "").strip()[-2000:] + stdout_tail = (proc.stdout or "").strip()[-2000:] + output_parts = [] + if stderr_tail: + output_parts.append(f"stderr: {stderr_tail}") + if stdout_tail: + output_parts.append(f"stdout: {stdout_tail}") + output_detail = "; ".join(output_parts) if output_parts else "(no output captured)" raise RuntimeError( f"onnxruntime-genai speech generation subprocess failed (exit code " f"{proc.returncode}). This typically indicates a native crash in onnxruntime-genai " f"for this Whisper build (e.g. a genai / model-builder version incompatibility). " - f"stderr tail: {stderr_tail}" + f"{output_detail}" ) with result_path.open() as f: return json.load(f) diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index e2d10fc0c4..bad4714356 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -47,9 +47,29 @@ class MobiusBuilder(Pass): whose components are individual :class:`~olive.model.ONNXModelHandler` objects. Single-component models return a plain :class:`~olive.model.ONNXModelHandler`. - Requires ``mobius-ai`` to be installed:: + Use ``components_to_export`` to export only a subset of components. This is + useful when some components (e.g. a text decoder) are already exported and + you only need the remaining ones (e.g. vision encoder and embedding):: + + { + "type": "MobiusBuilder", + "model_path": "mistralai/Ministral-3B-Instruct-2512", + "components_to_export": ["vision_encoder", "embedding"] + } + + Raises :class:`ValueError` if ``components_to_export`` is an empty list or + contains names not present in the built package. + + ORT GenAI config generation (``genai_config.json``, tokenizer files, processor + configs) is skipped when ``components_to_export`` is set, since mobius has no + API to scope that config to a subset of a package — the caller is responsible + for producing a ``genai_config.json`` that covers the full pipeline (e.g. by + combining this pass's output with another tool's output, as in a recipe that + exports the decoder separately). + + Requires ``mobius-onnx`` to be installed:: - pip install mobius-ai + pip install mobius-onnx See https://github.com/onnxruntime/mobius """ @@ -91,6 +111,20 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon "quantization pass (e.g. OnnxMatMulNBits) after this pass." ), ), + "components_to_export": PassConfigParam( + type_=list[str], + required=False, + default_value=None, + description=( + "Optional list of component names to export from a multi-component model " + "(e.g. ['vision', 'embedding'] to skip the decoder). " + "When set, only the named components are written by ``pkg.save()`` and " + "returned by this pass; all others are skipped entirely. " + "When not set (None), all components are exported (default, backward compatible). " + "Raises ValueError if the list is empty or if any specified name is not found in " + "the model's components." + ), + ), } def _run_for_config( @@ -103,7 +137,7 @@ def _run_for_config( from mobius import build except ImportError as exc: raise ImportError( - "mobius-ai is required to run MobiusBuilder. Install with: pip install mobius-ai" + "mobius-onnx is required to run MobiusBuilder. Install with: pip install mobius-onnx" ) from exc if not isinstance(model, HfModelHandler): @@ -136,6 +170,13 @@ def _run_for_config( if trust_remote_code: logger.warning("MobiusBuilder: trust_remote_code=True — only use with trusted model sources.") + # Validate components_to_export early (before the expensive build step). + if config.components_to_export is not None and len(config.components_to_export) == 0: + raise ValueError( + "MobiusBuilder: components_to_export cannot be empty. " + "Pass None to export all components, or specify at least one component name." + ) + output_dir = Path(output_model_path) output_dir.mkdir(parents=True, exist_ok=True) @@ -147,19 +188,60 @@ def _run_for_config( trust_remote_code=trust_remote_code, ) + # Determine which package components to export. + all_keys = list(pkg.keys()) + if config.components_to_export is not None: + requested = set(config.components_to_export) + unknown = requested - set(all_keys) + if unknown: + raise ValueError( + f"MobiusBuilder: components_to_export contains unknown component(s): {sorted(unknown)}. " + f"Available components from this model: {sorted(all_keys)}" + ) + package_keys = [k for k in all_keys if k in requested] + logger.info( + "MobiusBuilder: exporting subset of components %s (skipping %s)", + package_keys, + [k for k in all_keys if k not in requested], + ) + + def components_filter(name: str) -> bool: + return name in requested + else: + package_keys = all_keys + components_filter = None + # ModelPackage.save() handles both single and multi-component layouts: # single component → /model.onnx # multi-component → //model.onnx for each key - pkg.save(str(output_dir)) - - # Generate ORT GenAI config artifacts (genai_config.json, tokenizer - # files, processor configs) alongside the ONNX models. - genai_artifacts = self._write_genai_config(pkg, str(output_dir), model_id, ep_str) + pkg.save(str(output_dir), components=components_filter) + + # ORT GenAI config generation assumes every component in `pkg` was actually saved to + # disk (e.g. it unconditionally writes a "decoder/model.onnx" filename reference for + # multimodal packages). For a partial export (components_to_export set), that produces + # a genai_config.json/tokenizer set that references components we deliberately omitted + # — invalid artifacts, since mobius has no API to generate GenAI config for a subset of + # a package. Skip config generation entirely in that case and let the caller (e.g. a + # recipe combining this partial export with another tool's output) assemble the final + # genai_config.json itself. + if components_filter is not None: + logger.info( + "MobiusBuilder: components_to_export is set; skipping ORT GenAI config generation " + "since it cannot be scoped to a subset of components. The caller is responsible for " + "producing a genai_config.json that covers the full pipeline." + ) + genai_artifacts = {} + else: + # Generate ORT GenAI config artifacts (genai_config.json, tokenizer + # files, processor configs) alongside the ONNX models. + genai_artifacts = self._write_genai_config(pkg, str(output_dir), model_id, ep_str) - package_keys = list(pkg.keys()) logger.info("MobiusBuilder: saved components %s to '%s'", package_keys, output_dir) - if len(package_keys) == 1: + # Use the single-component (root layout) path only when the model is + # architecturally single-component. A multi-component model filtered + # down to one component still uses component sub-directories on disk. + if len(all_keys) == 1: # Single-component model (most LLMs): return a plain ONNXModelHandler. onnx_path = output_dir / "model.onnx" if not onnx_path.exists(): diff --git a/olive/passes/onnx/model_builder.py b/olive/passes/onnx/model_builder.py index d862fe214d..6b458f2df8 100644 --- a/olive/passes/onnx/model_builder.py +++ b/olive/passes/onnx/model_builder.py @@ -19,6 +19,7 @@ from packaging import version from olive.common.hf.utils import has_test_model_weights, is_test_model_dir +from olive.common.quant.patterns import match_override from olive.constants import Precision from olive.hardware.accelerator import AcceleratorSpec, Device from olive.hardware.constants import ExecutionProvider @@ -60,6 +61,14 @@ class AccuracyLevel(IntEnum): ExecutionProvider.NvTensorRTRTXExecutionProvider: "NvTensorRtRtx", } + # Olive exposes these model builder options as lists, but the model builder only understands + # the joined string form its CLI produces. Keys are matched with the deprecated `int4_` prefix + # stripped, since the model builder renames those aliases itself. + LIST_OPTION_SEPARATORS: ClassVar[dict[str, str]] = { + "op_types_to_quantize": "/", + "nodes_to_exclude": ",", + } + @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: return { @@ -214,7 +223,7 @@ def _run_for_config( config: type[BasePassConfig], output_model_path: str, ) -> ONNXModelHandler: - from onnxruntime_genai.models.builder import create_model + from onnxruntime_genai.models import builder self.maybe_patch_quant() @@ -294,7 +303,16 @@ def _run_for_config( try: logger.debug("Building model with the following args: %s", extra_args) - create_model( + self._check_extra_options( + model_path, + input_path, + output_model_filepath.parent, + precision, + target_execution_provider, + extra_args, + ) + + builder.create_model( model_name=model_path, input_path=input_path, output_dir=str(output_model_filepath.parent), @@ -429,6 +447,45 @@ def _run_for_config( return output_model + @staticmethod + def _check_extra_options( + model_name, + input_path, + output_dir, + precision, + execution_provider, + extra_options, + ): + """Run the model builder's pre-checks on ``extra_options`` before calling ``create_model``. + + The model builder validates the options, renames deprecated option aliases and looks up the + Hugging Face config in ``check_extra_options``; ``create_model`` fails outright when the + resulting ``hf_details`` entry is missing. ``check_extra_options`` is written against the + string values that ``--extra_options key=value`` produces, so Olive's typed config values + are serialized the same way first and the model builder stays the single owner of which + option means what. ``extra_options`` is updated in place and can be forwarded to + ``create_model`` afterwards. + """ + from onnxruntime_genai.models.builder import check_extra_options + + for key, value in list(extra_options.items()): + if isinstance(value, bool): + extra_options[key] = str(value).lower() + elif isinstance(value, (list, tuple)): + separator = ModelBuilder.LIST_OPTION_SEPARATORS.get(key.removeprefix("int4_")) + if separator: + extra_options[key] = separator.join(map(str, value)) + + check_extra_options( + model_name, + input_path, + str(output_dir), + precision, + execution_provider, + HF_HUB_CACHE, + extra_options, + ) + @staticmethod def maybe_patch_quant(): from onnxruntime_genai import __version__ as genai_version @@ -458,8 +515,18 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme from onnxruntime_genai.models.quantized_model import QuantizedDecoderLayer, QuantizedTensorModule, TensorModule from safetensors.torch import load_file + from olive.common.quant.state_dict import QWEIGHT_SUFFIX, QZEROS_SUFFIX, SCALES_SUFFIX + config = quant_attrs["config"] + if config.get("moe"): + raise NotImplementedError( + "ModelBuilder does not support loading Olive-quantized MoE checkpoints " + "(``quantization_config.moe == True``). Use the Mobius model builder for " + "MoE models or rerun the RTN pass with ``moe=False`` to leave experts in " + "their original precision." + ) + self.quant_type = quant_type self.embedding = QuantizedTensorModule() if config["embeds"] else TensorModule() self.final_norm = TensorModule() @@ -482,13 +549,20 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme overrides = config["overrides"] or {} - def get_layer_bits(layer_name): + def get_override(layer_name: str) -> dict: + # ``overrides`` keys support ``re:``-prefixed regex patterns (see + # ``olive.common.quant.patterns``); a plain ``dict.get`` would silently ignore + # them and fall back to the global bits/group_size, which then miscomputes + # ``in_features`` and reshapes the packed ``qweight`` incorrectly. name = ".".join(layer_name.split(".")[:-1]) - return overrides.get(name, {}).get("bits", config["bits"]) + matched = match_override(name, list(overrides.keys())) + return overrides[matched] if matched is not None else {} + + def get_layer_bits(layer_name): + return get_override(layer_name).get("bits", config["bits"]) def get_layer_group_size(layer_name): - name = ".".join(layer_name.split(".")[:-1]) - return overrides.get(name, {}).get("group_size", config["group_size"]) + return get_override(layer_name).get("group_size", config["group_size"]) def set_tensor(module, tensor_name, tensor_value, local_bits, local_group_size): submodule = module @@ -502,19 +576,30 @@ def set_tensor(module, tensor_name, tensor_value, local_bits, local_group_size): child = QuantizedTensorModule() setattr(submodule, sub_name, child) submodule = child + + attr_name = tensor_name.split(".")[-1] if isinstance(submodule, QuantizedTensorModule): + # Olive's native quantized checkpoints store buffers as + # ``_qweight`` / ``_scales`` / ``_qzeros`` (typically + # ``weight_*``); ``QuantizedTensorModule`` expects bare + # ``qweight`` / ``scales`` / ``qzeros`` attributes. + for suffix in (QWEIGHT_SUFFIX, SCALES_SUFFIX, QZEROS_SUFFIX): + if attr_name.endswith(suffix): + attr_name = suffix.lstrip("_") + break + for q_attr, q_value in [("bits", local_bits), ("_group_size", local_group_size)]: setattr(submodule, q_attr, q_value) # in_features is always a multiple of group_size, group_size is a power of 2 # assumes no padding - if tensor_name.endswith("qweight"): + if attr_name == "qweight": out_features, in_features_packed = tensor_value.shape in_features = in_features_packed * 8 // local_bits submodule.in_features = in_features submodule.out_features = out_features num_blocks = in_features // local_group_size if local_group_size != -1 else 1 tensor_value = tensor_value.reshape(out_features, num_blocks, -1) - setattr(submodule, tensor_name.split(".")[-1], tensor_value) + setattr(submodule, attr_name, tensor_value) for weight_file in Path(input_path).iterdir(): if weight_file.suffix == ".safetensors": @@ -591,7 +676,7 @@ def __init__(self): self.group_size = q_matmul.group_size matmul = PackedMatMul() - return self.make_matmul_int4(matmul, basename, root_input, **kwargs) + return self.make_matmul_nbits(matmul, basename, root_input, **kwargs) def patched_make_embedding(self, embedding): diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py index 5b76f8967f..64c233f337 100644 --- a/olive/passes/onnx/static_llm.py +++ b/olive/passes/onnx/static_llm.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import logging +from copy import deepcopy from pathlib import Path import onnx @@ -71,12 +72,25 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon default_value=64, description="Input length of the context model.", ), + "prefill_decode_models": PassConfigParam( + type_=bool, + default_value=True, + description=("To generate prefill and decode models. Specifically for QNN GPU"), + ), "group_session_options": PassConfigParam( type_=dict, description=( "Session options for the context and iterator models. Only used for models with genai_config." ), ), + "update_genai_config": PassConfigParam( + type_=bool, + default_value=False, + description=( + "Whether to update the genai_config.json pipeline for the generated static LLM model(s)." + " Only applicable to QNN GPU." + ), + ), } def _run_for_config(self, model, config: type[BasePassConfig], output_model_path: str): @@ -199,57 +213,159 @@ def process_context_iterator(component_models, llm_pipeline, output_dir): ) def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: Path): + """QNN_GPU path: generate one or more static ONNX models for different context lengths. + + - If config.prefill_decode_models is false: generate single model. + - If config.prefill_decode_models is true: generate multiple models (prefill/arN and decode/ar1) and return + CompositeModelHandler. + """ output_model_dir = Path(output_model_path).with_suffix("") model_path = Path(model.model_path) # --- Step 1: Load model (handle both single and external data) --- try: - model_proto = onnx.load(model_path, load_external_data=True) + base_model_proto = onnx.load(model_path, load_external_data=True) except Exception as e: raise RuntimeError(f"Failed to load ONNX model: {e}") from e # --- Step 2: Fix symbolic dimensions --- - batch_size, sequence_length = _proto_io_shape(model_proto, "input_ids") + batch_size, sequence_length = _proto_io_shape(base_model_proto, "input_ids") if not (isinstance(batch_size, str) and isinstance(sequence_length, str)): raise ValueError("Input dimensions must be symbolic before static shape fixing.") - param_mapping = {batch_size: config.batch_size, sequence_length: config.context_length} - self.fix_shape(model_proto, param_mapping) + prefill_decode_models = getattr(config, "prefill_decode_models", True) + + if not prefill_decode_models: + # Single model mode + ctx_lengths_list = [int(config.context_length)] + else: + # Composite model mode → AR1 + AR-N + n = int(config.context_length) + ctx_lengths_list = [n, 1] + + multiple = len(ctx_lengths_list) > 1 + + generated_handlers: dict[int, ONNXModelHandler] = {} + generated_names: dict[int, str] = {} + + for ctx_len in ctx_lengths_list: + # --- Clone base model proto for this variant --- + model_proto = onnx.ModelProto() + model_proto.CopyFrom(base_model_proto) + + # --- Step 3: Fix symbolic dimensions for this context length --- + param_mapping = {batch_size: config.batch_size, sequence_length: ctx_len} + self.fix_shape(model_proto, param_mapping) + + add_version_metadata_to_model_proto(model_proto) + + # --- Step 4: Save as external-data ONNX --- + # single model: "model", composite: "prefill" (AR-N) or "decode" (AR-1) + if not multiple: + logical_name = "model" + elif ctx_len == 1: + logical_name = "decode" + else: + logical_name = "prefill" + onnx_file_name = f"{logical_name}.onnx" + output_model_file = Path(output_model_dir) / onnx_file_name + # share a single external-data file. + external_data_file = Path(output_model_dir) / "model.onnx.data" + + output_model_dir.mkdir(parents=True, exist_ok=True) + external_data_file.unlink(missing_ok=True) + onnx.save( + model_proto, + str(output_model_file), + save_as_external_data=True, + all_tensors_to_one_file=True, + location=external_data_file.name, + convert_attribute=False, + ) + + # Build handler for this static model + new_model_attributes = deepcopy(model.model_attributes) or {} + handler = ONNXModelHandler( + model_path=output_model_dir, + onnx_file_name=output_model_file.name, + model_attributes=new_model_attributes, + ) + + # Store handler + logical component name + generated_handlers[ctx_len] = handler + generated_names[ctx_len] = logical_name + + # --- Step 5: Update genai_config.json --- + # For single model: pipeline with one component. + # For multiple models: pipeline with multiple components (composite). + if not multiple: + # Single context length + ctx_len = ctx_lengths_list[0] + handler = generated_handlers[ctx_len] + + decoder_config_extra = { + "inputs": { + "past_sequence_length": "past_seq_len", + "total_sequence_length": "total_seq_len", + }, + "sliding_window": { + "window_size": ctx_len, + "pad_value": 0, + "alignment": "left", + "slide_key_value_cache": False, + }, + } + + if not config.update_genai_config: + return handler + + return update_llm_pipeline_genai_config_gpu( + model=handler, + output_model_dir=output_model_dir, + decoder_config_extra=decoder_config_extra, + composite_components=None, + ) + + # Multiple context lengths -> wrap in CompositeModelHandler and create composite pipeline + components = [] + component_names = [] - # --- Step 3: Save model as external-data format --- - output_model_file = Path(output_model_dir) / "model.onnx" - external_data_file = Path(output_model_dir) / "model.onnx.data" + for ctx_len, handler in generated_handlers.items(): + components.append(handler) + component_names.append(generated_names[ctx_len]) - onnx.save( - model_proto, - str(output_model_file), - save_as_external_data=True, - all_tensors_to_one_file=True, - location=external_data_file.name, - convert_attribute=False, + new_model_attributes = deepcopy(model.model_attributes) or {} + # prefill and decode components share a single external-data file (model.onnx.data); + # tell OliveCache.save_model to preserve that shared name instead of renaming it after + # whichever component is copied first. + new_model_attributes["keep_shared_external_data_names"] = True + + composite = CompositeModelHandler( + model_components=components, model_component_names=component_names, model_attributes=new_model_attributes ) - decoder_config_extra = { + if not config.update_genai_config: + return composite + + # Build per-component sliding_window config keyed by name + composite_decoder_extra = { "inputs": { "past_sequence_length": "past_seq_len", "total_sequence_length": "total_seq_len", }, "sliding_window": { - "window_size": config.context_length, + "window_size": max(ctx_lengths_list), "pad_value": 0, "alignment": "left", "slide_key_value_cache": False, }, } - input_model_path = model.model_path - model_static = ONNXModelHandler(model_path=output_model_dir, onnx_file_name=output_model_file.name) - return update_llm_pipeline_genai_config_gpu( - model_static, - output_model_dir, - input_model_path, - decoder_config_extra, + model=composite, + output_model_dir=output_model_dir, + decoder_config_extra=composite_decoder_extra, + composite_components=list(zip(component_names, components)), ) @staticmethod diff --git a/olive/passes/pytorch/autoclip.py b/olive/passes/pytorch/autoclip.py index 5b8d7b2960..52e5af3038 100644 --- a/olive/passes/pytorch/autoclip.py +++ b/olive/passes/pytorch/autoclip.py @@ -131,9 +131,9 @@ def _get_oc_batch_size(out_features: int) -> int: @staticmethod def accumulate_inputs(module: torch.nn.Module, inputs: tuple, _: torch.Tensor) -> None: - if module.quant_info.data is None: - module.quant_info.data = {"inputs": []} - module.quant_info.data["inputs"].append(inputs[0].detach().cpu()) + if module.weight.quant_info.data is None: + module.weight.quant_info.data = {"inputs": []} + module.weight.quant_info.data["inputs"].append(inputs[0].detach().cpu()) @classmethod def process_module( @@ -144,11 +144,11 @@ def process_module( max_shrink: float, n_sample_token: int, ) -> None: - if module.quant_info.data is None or not module.quant_info.data.get("inputs"): + if module.weight.quant_info.data is None or not module.weight.quant_info.data.get("inputs"): raise ValueError(f"Module {module} does not have cached inputs initialized!") - input_feat = torch.cat(module.quant_info.data["inputs"], dim=0) - module.quant_info.data = None + input_feat = torch.cat(module.weight.quant_info.data["inputs"], dim=0) + module.weight.quant_info.data = None module.to(device) cls._auto_clip_layer( @@ -173,7 +173,7 @@ def _auto_clip_layer( if weight.dim() != 2: raise ValueError("AutoClip expects a 2D linear weight tensor.") - quantizer = module.quant_info.quantizer + quantizer = module.weight.quant_info.quantizer effective_group_size = weight.shape[1] if quantizer.group_size <= 0 else quantizer.group_size if weight.shape[1] % effective_group_size != 0: raise ValueError("Weight in_features must be divisible by group_size.") diff --git a/olive/passes/pytorch/gptq.py b/olive/passes/pytorch/gptq.py index c24381b3aa..778a521d9a 100644 --- a/olive/passes/pytorch/gptq.py +++ b/olive/passes/pytorch/gptq.py @@ -14,9 +14,16 @@ from olive.data.config import DataConfig from olive.passes import Pass from olive.passes.pass_config import BasePassConfig, PassConfigParam +from olive.passes.pytorch.moe_calib import ( + DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE, + DEFAULT_MOE_FALLBACK_THRESHOLD, + MoeCalibrationSession, +) from olive.passes.pytorch.quant_utils import ( + _module_weight_has_quant_info, finalize, get_quantizer_config, + module_quant_info_param_names, prepare_model, run_layerwise_quantization, ) @@ -38,7 +45,7 @@ class Gptq(Pass): @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: return { - **get_quantizer_config(), + **get_quantizer_config(allow_moe=True), "damp_percent": PassConfigParam( type_=float, default_value=0.01, @@ -62,6 +69,40 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon " Required for PyTorch models." ), ), + "moe_fallback_threshold": PassConfigParam( + type_=float, + default_value=DEFAULT_MOE_FALLBACK_THRESHOLD, + description=( + "Only used when moe=True. Fraction of the calibration tokens reaching an experts" + " module below which an individual expert is quantized with round-to-nearest" + " instead of GPTQ. MoE routing sends only a fraction of the tokens to each" + " expert, so cold experts get a rank-deficient (or missing) Hessian for which" + " GPTQ's correction is driven by the damping prior rather than by data." + " Default is 0.005 (0.5%, matching GPTQModel's default). This threshold measures" + " routing skew (is this expert under-served relative to its peers?), which is" + " scale-invariant: it does not by itself guarantee the expert's Hessian is" + " numerically well-formed -- see moe_fallback_min_k_multiple, which measures" + " statistical sufficiency instead. An expert falls back to RTN if it fails" + " EITHER condition." + ), + ), + "moe_fallback_min_k_multiple": PassConfigParam( + type_=float, + default_value=DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE, + description=( + "Only used when moe=True. Minimum number of calibration tokens an expert must" + " have seen, expressed as a multiple of K (the expert weight's last dimension)," + " below which it is quantized with round-to-nearest instead of GPTQ. Each" + " expert's Hessian is a (K, K) matrix accumulated from its routed tokens, so" + " rank(H) <= num_tokens_seen: below K tokens, the Hessian is necessarily" + " rank-deficient, but N >= K alone does not guarantee GPTQ beats RTN -- damping" + " reweights rank-deficient directions rather than eliminating their influence," + " so GPTQ's correction can still be noisier than RTN somewhat above the bare" + " N=K floor. Default is 2.0, set past the empirically measured crossover" + " (roughly 1x-2x K, model/config dependent) rather than at N=K itself. An expert" + " falls back to RTN if it fails EITHER this condition or moe_fallback_threshold." + ), + ), } @classmethod @@ -81,6 +122,14 @@ def validate_config( logger.info("desc_act can only be True when group_size is -1.") return False + if not 0 <= config.moe_fallback_threshold < 1: + logger.info("moe_fallback_threshold must be in [0, 1).") + return False + + if not math.isfinite(config.moe_fallback_min_k_multiple) or config.moe_fallback_min_k_multiple < 0: + logger.info("moe_fallback_min_k_multiple must be finite and >= 0.") + return False + return True @torch.no_grad() @@ -99,16 +148,30 @@ def _run_for_config( """ wrapper, qcfg, _ = prepare_model(model, config) + moe_session = ( + MoeCalibrationSession.create( + wrapper, + fallback_threshold=config.moe_fallback_threshold, + fallback_min_k_multiple=config.moe_fallback_min_k_multiple, + ) + if getattr(config, "moe", False) + else None + ) device = run_layerwise_quantization( model, wrapper, config.data_config, input_hook=self.accumulate_hessian, process_module=lambda module, _: self.process_module( - module, percdamp=config.damp_percent, actorder=config.desc_act + module, + percdamp=config.damp_percent, + actorder=config.desc_act, + moe_fallback_threshold=config.moe_fallback_threshold, + moe_fallback_min_k_multiple=config.moe_fallback_min_k_multiple, ), update_before_process=False, include_lm_head=config.lm_head, + moe_session=moe_session, ) return finalize(model, output_model_path, wrapper, qcfg, device) @@ -123,8 +186,8 @@ def accumulate_hessian(module: torch.nn.Module, inp: tuple, _: Any) -> None: _: Unused output parameter. """ - if module.quant_info.data is None: - module.quant_info.data = { + if module.weight.quant_info.data is None: + module.weight.quant_info.data = { "H": torch.zeros((module.in_features, module.in_features), device=inp[0].device), "N": 0, } @@ -132,122 +195,330 @@ def accumulate_hessian(module: torch.nn.Module, inp: tuple, _: Any) -> None: batch_size = inp[0].shape[0] inp = inp[0].reshape(-1, module.in_features).t() - module.quant_info.data["H"] *= module.quant_info.data["N"] / (module.quant_info.data["N"] + batch_size) - module.quant_info.data["N"] += batch_size - inp = math.sqrt(2 / module.quant_info.data["N"]) * inp.float() - module.quant_info.data["H"] += inp.matmul(inp.t()) + module.weight.quant_info.data["H"] *= module.weight.quant_info.data["N"] / ( + module.weight.quant_info.data["N"] + batch_size + ) + module.weight.quant_info.data["N"] += batch_size + inp = math.sqrt(2 / module.weight.quant_info.data["N"]) * inp.float() + module.weight.quant_info.data["H"] += inp.matmul(inp.t()) @staticmethod def process_module( - module: torch.nn.Module, blocksize: int = 128, percdamp: float = 0.01, actorder: bool | None = False + module: torch.nn.Module, + blocksize: int = 128, + percdamp: float = 0.01, + actorder: bool | None = False, + moe_fallback_threshold: float = DEFAULT_MOE_FALLBACK_THRESHOLD, + moe_fallback_min_k_multiple: float = DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE, ) -> None: - """Process a module for GPTQ quantization using the accumulated Hessian. + """Process a module for GPTQ quantization using the accumulated calibration data. + + Dispatches on how the module's selected parameters were calibrated: + + * ``nn.Linear`` / ``nn.Embedding`` ``weight`` -- a single ``(K, K)`` Hessian + collected by :meth:`accumulate_hessian` from a forward hook; + * fused-3D MoE experts parameters (``gate_up_proj`` / ``down_proj``) -- one + independent ``(K, K)`` Hessian *per expert*, collected by + :mod:`olive.passes.pytorch.moe_calib`. Experts that saw too few calibration + tokens fall back to RTN. Args: - module: The linear module to quantize. + module: The module to quantize. blocksize: Block size for processing weights. percdamp: Damping factor for numerical stability. actorder: Whether to use act-order quantization scheme. + moe_fallback_threshold: Fraction of the calibration tokens reaching an experts + module below which an expert is quantized with the RTN fallback. + moe_fallback_min_k_multiple: Minimum calibration tokens for an expert, expressed + as a multiple of K, below which an expert is quantized with the RTN fallback. """ - if module.quant_info.data is None: - raise ValueError(f"Module {module} does not have quant_info.data initialized!") - - if actorder is None: - actorder = module.quant_info.quantizer.group_size == -1 - elif actorder is True: - assert module.quant_info.quantizer.group_size == -1, ( - "actorder can only be True when group_size is -1, but got group_size=" - f"{module.quant_info.quantizer.group_size}" - ) - - H = module.quant_info.data["H"] - W = module.weight.data.clone().float().to(H.device) - num_cols = H.shape[0] - - dead = torch.diag(H) == 0 - H[dead, dead] = 1 - W[:, dead] = 0 - - if actorder: - perm = torch.argsort(torch.diag(H), descending=True) - W = W[:, perm] - H = H[perm][:, perm] - invperm = torch.argsort(perm) - - Losses = torch.zeros_like(W) - Q = torch.zeros_like(W) - - damp = percdamp * torch.mean(torch.diag(H)) - diag = torch.arange(num_cols, device=H.device) - H[diag, diag] += damp - Hinv = torch.linalg.cholesky(H) # pylint: disable=not-callable - del H - Hinv = torch.cholesky_inverse(Hinv) - Hinv = torch.linalg.cholesky(Hinv, upper=True) # pylint: disable=not-callable - - all_scales = [] - all_zp = [] - now_idx = 1 - # create a per-channel quantizer - quantizer = WeightQuantizer( - bits=module.quant_info.quantizer.bits, symmetric=module.quant_info.quantizer.symmetric, group_size=-1 - ) - if module.quant_info.quantizer.group_size == -1: - # this can be before or after actorder permutation since there's only one group - active_scale, active_zp = quantizer.find_qparams(W) + if _module_weight_has_quant_info(module): + Gptq._process_dense_module(module, blocksize=blocksize, percdamp=percdamp, actorder=actorder) else: - active_scale, active_zp = None, None - - for i1 in range(0, num_cols, blocksize): - i2 = min(i1 + blocksize, num_cols) - count = i2 - i1 - - W1 = W[:, i1:i2].clone() - Q1 = torch.zeros_like(W1) - Err1 = torch.zeros_like(W1) - Losses1 = torch.zeros_like(W1) - Hinv1 = Hinv[i1:i2, i1:i2] - - for i in range(count): - w = W1[:, i] - d = Hinv1[i, i] - - if module.quant_info.quantizer.group_size != -1: - if (i1 + i) % module.quant_info.quantizer.group_size == 0: - active_scale, active_zp = quantizer.find_qparams( - W[:, (i1 + i) : (i1 + i + module.quant_info.quantizer.group_size)] - ) + for pname in module_quant_info_param_names(module): + Gptq._process_moe_param( + module, + pname, + blocksize=blocksize, + percdamp=percdamp, + actorder=actorder, + fallback_threshold=moe_fallback_threshold, + fallback_min_k_multiple=moe_fallback_min_k_multiple, + ) - if ((i1 + i) // module.quant_info.quantizer.group_size) - now_idx == -1: - all_scales.append(active_scale) - all_zp.append(active_zp) - now_idx += 1 - - q = quantizer.fake_quantize(w.unsqueeze(1), active_scale, active_zp).flatten() - Q1[:, i] = q - Losses1[:, i] = (w - q) ** 2 / d**2 - - err1 = (w - q) / d - W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0)) - Err1[:, i] = err1 + torch.cuda.empty_cache() - Q[:, i1:i2] = Q1 - Losses[:, i1:i2] = Losses1 / 2 + @staticmethod + def _process_dense_module( + module: torch.nn.Module, blocksize: int = 128, percdamp: float = 0.01, actorder: bool | None = False + ) -> None: + """Quantize ``module.weight`` with GPTQ using its accumulated Hessian.""" + info = module.weight.quant_info + if info.data is None: + raise ValueError(f"Module {module} does not have quant_info.data initialized!") - W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:]) + Q, scales, zero_points = gptq_quantize_weight( + module.weight.data.clone().float().to(info.data["H"].device), + info.data["H"], + info.quantizer, + blocksize=blocksize, + percdamp=percdamp, + actorder=actorder, + ) - if actorder: - Q = Q[:, invperm] + module.weight.data = Q.to(module.weight.data.device).to(module.weight.data.dtype) + info.scales = scales.to("cpu") + info.zero_points = zero_points.to("cpu") + info.data = None - if not all_scales: - all_scales.append(active_scale) - all_zp.append(active_zp) + @staticmethod + def _process_moe_param( + module: torch.nn.Module, + pname: str, + blocksize: int = 128, + percdamp: float = 0.01, + actorder: bool | None = False, + fallback_threshold: float = DEFAULT_MOE_FALLBACK_THRESHOLD, + fallback_min_k_multiple: float = DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE, + ) -> None: + """Quantize one fused-3D MoE parameter, expert by expert. + + Each expert is quantized from its *own* Hessian. An expert falls back to RTN -- + instead of GPTQ -- when it fails EITHER of two independent conditions, since they + measure different things: + + * **Routing skew** (``fallback_threshold``, a fraction of the calibration tokens + reaching this experts module): is this expert under-served *relative to its + peers*? Scale-invariant -- 10x more calibration data means this expert's own + token count also scales ~10x, so the ratio is unchanged. + * **Statistical sufficiency** (``fallback_min_k_multiple``, a multiple of K, this + parameter's last dimension): does this expert have *enough absolute samples* to + estimate a well-formed ``(K, K)`` Hessian at all? ``H = sum(x xT)`` accumulated + from ``N`` tokens has ``rank(H) <= N``, so ``N < K`` makes ``H`` provably + singular. This is *necessary* but not *sufficient* for GPTQ to underperform RTN + in general: damping makes ``H`` invertible even when rank-deficient, but it + reweights the null directions by ``1/lambda`` rather than eliminating their + influence, so the sequential column corrections stay coupled and can amplify + noise past ``N = K``. Empirically GPTQ has been measured to still underperform + RTN somewhere in the ``1x-2x K`` range and only reliably beat it above roughly + ``2x K`` -- the default ``fallback_min_k_multiple`` is set past that measured + crossover rather than at the bare ``N = K`` rank floor. This condition is + absolute: more calibration data helps it directly, unlike routing skew. + + Both checks must pass for GPTQ to be used. A cold-but-adequate expert (e.g. 5x K + tokens, but still a small share of a very large calibration set) falls back to RTN + on skew alone. Likewise, a "fair share" expert that is still far below the + sufficiency threshold (e.g. because the whole calibration set is too small) falls + back on insufficiency alone. This dual gate is intended to make GPTQ+fallback no + worse than plain RTN under realistic calibration budgets, but it is a measured + design choice tuned against the crossover above, not a theorem -- see + moe_fallback_min_k_multiple's docstring. + """ + param = module._parameters[pname] # pylint: disable=protected-access + info = param.quant_info + data = info.data + if not data or not data.get("moe"): + raise ValueError( + f"MoE parameter '{pname}' of {type(module).__name__} has no per-expert calibration " + "data. This usually means the experts forward was never intercepted during " + "calibration." + ) - module.weight.data = Q.to(module.weight.data.device).to(module.weight.data.dtype) - module.quant_info.scales = torch.cat(all_scales, dim=1).to("cpu") - module.quant_info.zero_points = torch.cat(all_zp, dim=1).to("cpu") + weight = param.data + num_experts = weight.shape[0] + k = weight.shape[-1] + skew_threshold = fallback_threshold * data["tokens_seen"] + sufficiency_threshold = fallback_min_k_multiple * k + observed = data["token_counts"] + + expert_weights, expert_scales, expert_zero_points = [], [], [] + fallback_experts = [] + fallback_skew_only = [] + fallback_sufficiency_only = [] + for expert_idx in range(num_experts): + entry = data["experts"].get(expert_idx) + W = weight[expert_idx].clone().float() + n = entry["N"] if entry is not None else 0 + skew_fails = n < skew_threshold + sufficiency_fails = n < sufficiency_threshold + if entry is None or skew_fails or sufficiency_fails: + fallback_experts.append(expert_idx) + if skew_fails and not sufficiency_fails: + fallback_skew_only.append(expert_idx) + elif sufficiency_fails and not skew_fails: + fallback_sufficiency_only.append(expert_idx) + # RTN: derive qparams straight from the float weight and fake-quantize with + # them. Fake-quantizing here (rather than deferring all rounding to + # ``finalize``) keeps the true-sequential invariant: the post-quantization + # re-run of the layer sees on-grid weights for *every* expert, GPTQ or + # fallback. After ``Q`` is cast back to the weight's storage dtype below, + # ``finalize`` may round a small number of fp16/bf16 values to an adjacent + # quantized integer, so the packed result is not guaranteed to be + # bit-identical to a standalone Rtn pass. + scales, zero_points = info.quantizer.find_qparams(W) + Q = info.quantizer.fake_quantize(W, scales, zero_points) + else: + Q, scales, zero_points = gptq_quantize_weight( + W.to(entry["H"].device), + entry["H"], + info.quantizer, + blocksize=blocksize, + percdamp=percdamp, + actorder=actorder, + ) + expert_weights.append(Q.to(weight.device).to(weight.dtype)) + expert_scales.append(scales.to("cpu")) + expert_zero_points.append(zero_points.to("cpu")) + + if fallback_experts: + # Report which condition(s) actually drove each fallback -- a min/max token + # count across all fallback experts can otherwise look inconsistent with + # whichever single threshold a reader happens to compare it against (e.g. a + # fallback expert with N=11330 above the skew threshold but below a large + # sufficiency threshold would look like a logging bug if only one number is + # shown). + logger.info( + "GPTQ MoE fallback for '%s': %d/%d experts quantized with RTN (skew threshold " + "%.1f = %.2f%% of %d calibration tokens reaching this module; sufficiency " + "threshold %.1f = %.2fx K=%d; observed min=%d max=%d). %d expert(s) failed " + "skew only: %s. %d expert(s) failed sufficiency only: %s. All fallback " + "experts: %s", + type(module).__name__ + "." + pname, + len(fallback_experts), + num_experts, + skew_threshold, + 100 * fallback_threshold, + data["tokens_seen"], + sufficiency_threshold, + fallback_min_k_multiple, + k, + min(observed, default=0), + max(observed, default=0), + len(fallback_skew_only), + fallback_skew_only, + len(fallback_sufficiency_only), + fallback_sufficiency_only, + fallback_experts, + ) - module.quant_info.data = None - if torch.cuda.is_available(): - torch.cuda.empty_cache() + param.data = torch.stack(expert_weights, dim=0) + info.scales = torch.stack(expert_scales, dim=0) + info.zero_points = torch.stack(expert_zero_points, dim=0) + info.data = None + + +@torch.no_grad() +def gptq_quantize_weight( + W: torch.Tensor, + H: torch.Tensor, + quantizer: WeightQuantizer, + blocksize: int = 128, + percdamp: float = 0.01, + actorder: bool | None = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the GPTQ column sweep on a single 2D weight matrix. + + Args: + W: Float weight of shape ``(out_features, K)``; quantization runs along the last dim. + H: The ``(K, K)`` Hessian accumulated from this weight's calibration inputs. Modified + in place (dead-column patching + damping). + quantizer: The target :class:`WeightQuantizer` (supplies bits / symmetric / group_size). + blocksize: Column block size for the error-compensated sweep. + percdamp: Damping factor for numerical stability. + actorder: Act-order (desc_act) scheme. ``None`` means "True iff per-channel". + + Returns: + ``(Q, scales, zero_points)`` -- the fake-quantized weight and its qparams, with + ``scales``/``zero_points`` of shape ``(out_features, num_groups)``. + + """ + group_size = quantizer.group_size + if actorder is None: + actorder = group_size == -1 + elif actorder is True: + assert group_size == -1, f"actorder can only be True when group_size is -1, but got group_size={group_size}" + + W = W.to(H.device) + num_cols = H.shape[0] + + dead = torch.diag(H) == 0 + H[dead, dead] = 1 + W[:, dead] = 0 + + if actorder: + perm = torch.argsort(torch.diag(H), descending=True) + W = W[:, perm] + H = H[perm][:, perm] + invperm = torch.argsort(perm) + + Losses = torch.zeros_like(W) + Q = torch.zeros_like(W) + + damp = percdamp * torch.mean(torch.diag(H)) + diag = torch.arange(num_cols, device=H.device) + H[diag, diag] += damp + Hinv = torch.linalg.cholesky(H) # pylint: disable=not-callable + del H + Hinv = torch.cholesky_inverse(Hinv) + Hinv = torch.linalg.cholesky(Hinv, upper=True) # pylint: disable=not-callable + + all_scales = [] + all_zp = [] + now_idx = 1 + # create a per-channel quantizer + per_channel_quantizer = WeightQuantizer( + bits=quantizer.bits, + symmetric=quantizer.symmetric, + group_size=-1, + ) + if group_size == -1: + # this can be before or after actorder permutation since there's only one group + active_scale, active_zp = per_channel_quantizer.find_qparams(W) + else: + active_scale, active_zp = None, None + + for i1 in range(0, num_cols, blocksize): + i2 = min(i1 + blocksize, num_cols) + count = i2 - i1 + + W1 = W[:, i1:i2].clone() + Q1 = torch.zeros_like(W1) + Err1 = torch.zeros_like(W1) + Losses1 = torch.zeros_like(W1) + Hinv1 = Hinv[i1:i2, i1:i2] + + for i in range(count): + w = W1[:, i] + d = Hinv1[i, i] + + if group_size != -1: + if (i1 + i) % group_size == 0: + active_scale, active_zp = per_channel_quantizer.find_qparams(W[:, (i1 + i) : (i1 + i + group_size)]) + + if ((i1 + i) // group_size) - now_idx == -1: + all_scales.append(active_scale) + all_zp.append(active_zp) + now_idx += 1 + + q = per_channel_quantizer.fake_quantize(w.unsqueeze(1), active_scale, active_zp).flatten() + Q1[:, i] = q + Losses1[:, i] = (w - q) ** 2 / d**2 + + err1 = (w - q) / d + W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0)) + Err1[:, i] = err1 + + Q[:, i1:i2] = Q1 + Losses[:, i1:i2] = Losses1 / 2 + + W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:]) + + if actorder: + Q = Q[:, invperm] + + if not all_scales: + all_scales.append(active_scale) + all_zp.append(active_zp) + + return Q, torch.cat(all_scales, dim=1), torch.cat(all_zp, dim=1) diff --git a/olive/passes/pytorch/kquant.py b/olive/passes/pytorch/kquant.py index dfb59d31fa..a7dfdeeaef 100644 --- a/olive/passes/pytorch/kquant.py +++ b/olive/passes/pytorch/kquant.py @@ -18,7 +18,13 @@ from olive.passes import Pass from olive.passes.pass_config import PassConfigParam -from olive.passes.pytorch.quant_utils import finalize, get_quantizer_config, prepare_model +from olive.passes.pytorch.moe_support import check_moe_layout_support +from olive.passes.pytorch.quant_utils import ( + _iter_quant_info_params, + finalize, + get_quantizer_config, + prepare_model, +) if TYPE_CHECKING: from olive.hardware.accelerator import AcceleratorSpec @@ -141,11 +147,15 @@ def kquant_find_qparams( minq: int, symmetric: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute k-quant per-group scale and zero point for a 2D weight tensor. + """Compute k-quant per-group scale and zero point for an N-D weight tensor. Args: - weight: 2D tensor of shape ``(out_features, in_features)``. For embeddings - this is ``(num_embeddings, embedding_dim)``. + weight: Tensor with at least two dimensions. Supported shapes include + ``(out_features, in_features)`` for linear weights, + ``(num_embeddings, embedding_dim)`` for embeddings, and + ``(num_experts, out_features, in_features)`` for fused MoE weights. + All leading dimensions are preserved and grouping is along the last + dimension. group_size: Group size along the last dimension. Must be > 0 and evenly divide ``weight.shape[-1]``. maxq: Inclusive maximum integer code (as produced by ``get_maxq_minq``). @@ -156,19 +166,19 @@ def kquant_find_qparams( Returns: Tuple ``(scales, zero_points)`` matching ``WeightQuantizer.find_qparams``: - * ``scales``: shape ``(out_features, num_groups)``, dtype matches the - input weight dtype. - * ``zero_points``: shape ``(out_features, num_groups)``, dtype - ``int32``, values in ``[minq, maxq]``. + * ``scales``: shape ``(*weight.shape[:-1], num_groups)``, dtype + matches the input weight dtype. + * ``zero_points``: shape ``(*weight.shape[:-1], num_groups)``, + dtype ``int32``, values in ``[minq, maxq]``. """ if maxq <= minq: raise ValueError(f"k-quant requires maxq > minq, got maxq={maxq}, minq={minq}.") if group_size <= 0: raise ValueError(f"k-quant requires group_size > 0, got {group_size}.") - if weight.dim() != 2: - raise ValueError(f"Expected a 2D weight tensor, got shape {tuple(weight.shape)}.") - out_features, in_features = weight.shape + if weight.dim() < 2: + raise ValueError(f"Expected a weight tensor with at least 2 dimensions, got shape {tuple(weight.shape)}.") + *batch_shape, in_features = weight.shape if in_features % group_size != 0: raise ValueError(f"in_features ({in_features}) must be divisible by group_size ({group_size}) for k-quant.") @@ -224,8 +234,8 @@ def kquant_find_qparams( zero_point = torch.clamp(torch.round(float(minq) - offset / scale), l_min, l_max).to(torch.int32) num_groups = in_features // group_size - scales = scale.reshape(out_features, num_groups).to(orig_dtype).contiguous() - zero_points = zero_point.reshape(out_features, num_groups).contiguous() + scales = scale.reshape(*batch_shape, num_groups).to(orig_dtype).contiguous() + zero_points = zero_point.reshape(*batch_shape, num_groups).contiguous() return scales, zero_points @@ -235,12 +245,13 @@ class KQuant(Pass): Per-group weight quantization using the iterative weighted-least-squares search from llama.cpp's ggml k-quants. Supports both asymmetric (scale and zero point) and symmetric (scale only) variants for 2-, 4-, and 8-bit - weights of ``nn.Linear`` and ``nn.Embedding`` modules. + weights of ``nn.Linear`` and ``nn.Embedding`` modules, plus K-last fused + MoE expert parameters when ``moe=True``. """ @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: - config = get_quantizer_config(allow_embeds=True) + config = get_quantizer_config(allow_embeds=True, allow_moe=True) config["group_size"] = PassConfigParam( type_=int, default_value=32, @@ -279,18 +290,36 @@ def _run_for_config( """ wrapper, qcfg, retie_word_embeddings = prepare_model(model, config, allow_quantized=True) + # Gate on this invocation's own ``moe`` request, not ``qcfg.moe`` -- ``prepare_model`` + # ORs in any pre-existing checkpoint's ``moe`` flag (see ``quant_utils.prepare_model``), + # so ``qcfg.moe`` can be True even when this run itself passed ``moe=False``. + if getattr(config, "moe", False): + experts_modules = [ + experts + for layer in wrapper.get_layer_wrappers() + if (experts := layer.get_experts(return_name=False)) is not None + ] + if experts_modules: + check_moe_layout_support( + experts_modules, + model_type=wrapper.model_type, + operation="KQuant MoE quantization", + ) device = "cuda" if torch.cuda.is_available() else "cpu" from tqdm.auto import tqdm - modules = [(name, m) for name, m in wrapper.model.named_modules() if hasattr(m, "quant_info")] - pbar = tqdm(modules, desc="Quantizing modules") - for name, module in pbar: - pbar.set_postfix(module=name, refresh=False) - quantizer = module.quant_info.quantizer - - weight = module.weight.data.to(device) - effective_group_size = quantizer.group_size if quantizer.group_size > 0 else weight.shape[1] + module_names = {id(module): name for name, module in wrapper.model.named_modules()} + targets = list(_iter_quant_info_params(wrapper.model)) + pbar = tqdm(targets, desc="Quantizing parameters") + for sub_module, pname, param, quant_info in pbar: + module_name = module_names[id(sub_module)] + parameter_name = f"{module_name}.{pname}" if module_name else pname + pbar.set_postfix(parameter=parameter_name, refresh=False) + quantizer = quant_info.quantizer + + weight = param.data.to(device) + effective_group_size = quantizer.group_size if quantizer.group_size > 0 else weight.shape[-1] scales, zero_points = kquant_find_qparams( weight, group_size=effective_group_size, @@ -298,8 +327,8 @@ def _run_for_config( minq=quantizer.minq, symmetric=quantizer.symmetric, ) - module.quant_info.scales = scales.to("cpu") - module.quant_info.zero_points = zero_points.to("cpu") + quant_info.scales = scales.to("cpu") + quant_info.zero_points = zero_points.to("cpu") if torch.cuda.is_available(): torch.cuda.empty_cache() diff --git a/olive/passes/pytorch/moe_calib.py b/olive/passes/pytorch/moe_calib.py new file mode 100644 index 0000000000..bab8ebfbd5 --- /dev/null +++ b/olive/passes/pytorch/moe_calib.py @@ -0,0 +1,769 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Per-expert calibration for MoE (Mixture-of-Experts) weight quantization. + +Calibrated passes (GPTQ) need the *activations that actually reach each expert*. For fused +MoE architectures routing happens inside a single experts-module forward call, so a plain +``register_forward_hook`` on the experts module sees one undifferentiated activation batch +and cannot attribute rows to experts. + +Interception therefore goes through transformers' own experts-implementation registry +(``ALL_EXPERTS_FUNCTIONS`` / ``@use_experts_implementation``, transformers >= 5.0): we +register one generic recording implementation and point the model at it for the duration of +calibration. Every decorated experts class dispatches to it with no per-model branching, +which is why this file contains no architecture-specific code paths. Layout support is +determined from transformers-owned ``is_transposed`` metadata: fused weights must be stored +``(num_experts, out_features, in_features)`` (K last), matching GPTQ's Hessian layout. + +Transposed-layout architectures such as gpt-oss store ``(num_experts, in, out)``. GPTQ's +``(K, K)`` Hessian math assumes K is the last dim, so they are refused with a clear error +rather than silently mis-quantized. Architectures such as llama4 and aria that do not report +``is_transposed`` are also refused because Olive cannot verify their layout; supporting +either case requires work that is deliberately out of scope here. +""" + +from __future__ import annotations + +import logging +import math +import threading +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import torch + +from olive.passes.pytorch.moe_support import MoeSupportError, check_moe_layout_support + +if TYPE_CHECKING: + from collections.abc import Iterator + + from olive.common.hf.wrapper import ModelWrapper + +logger = logging.getLogger(__name__) + + +#: Key under which the recording forward is registered in ``ALL_EXPERTS_FUNCTIONS``. +OLIVE_MOE_CALIB_IMPLEMENTATION = "olive_moe_calib" + +#: Minimum transformers version exposing ``ALL_EXPERTS_FUNCTIONS`` / ``set_experts_implementation``. +MIN_TRANSFORMERS_VERSION = "5.0.0" + +#: Fraction of the calibration tokens reaching an experts module below which an expert is +#: quantized by the RTN fallback instead of GPTQ. Matches GPTQModel's ``"0.5%"`` default. +#: Measures routing skew (is this expert under-served relative to its peers?) -- scale +#: invariant, so it does not by itself guarantee an expert's Hessian is well-formed. An +#: expert falls back if it fails this OR ``DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE``. +DEFAULT_MOE_FALLBACK_THRESHOLD = 0.005 + +#: Minimum number of calibration tokens an expert must have seen, expressed as a multiple of +#: K (the expert weight's last dimension), below which an expert is quantized by the RTN +#: fallback instead of GPTQ. An expert's Hessian is a (K, K) matrix accumulated from its +#: routed tokens, so rank(H) <= num_tokens_seen: below K tokens the Hessian is *necessarily* +#: rank-deficient (N < K is necessary but not sufficient for a well-conditioned Hessian -- +#: damping does not make GPTQ's correction reduce to exactly RTN, it reweights the +#: rank-deficient directions rather than eliminating their influence, so a naive "N=K is the +#: floor" framing overstates what's guaranteed). Empirically, GPTQ has been measured to +#: underperform plain RTN somewhere in the 1x-2x K range and only reliably beat it above +#: roughly 2x K; the default below is set conservatively past that empirical crossover rather +#: than at the bare N>=K rank floor. Measures statistical sufficiency (absolute: more +#: calibration data genuinely helps), unlike DEFAULT_MOE_FALLBACK_THRESHOLD's routing-skew +#: measure. An expert falls back if it fails EITHER condition. +DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE = 2.0 + +#: Peak per-layer Hessian working set (bytes) above which :func:`check_moe_gptq_support` +#: warns about a likely out-of-memory during calibration. One float32 ``(K, K)`` Hessian is +#: allocated per (expert, parameter), so a layer needs +#: ``num_experts * (hidden_size**2 + intermediate_size**2) * 4`` bytes on the calibration +#: device, on top of the layer's own weights and activations. 4 GiB is chosen as the +#: threshold because it is small enough to fire well before any mainstream accelerator +#: (16-80 GB) actually OOMs -- note this means even Mixtral-8x7B (~6.6 GB/layer at +#: hidden_size=4096, intermediate_size=14336, num_local_experts=8) is expected to trip this +#: warning; it is not reserved for exceptionally large configs like DeepSeek-V3 (~57 GB/layer). +MOE_HESSIAN_MEMORY_WARN_BYTES = 4 * 1024**3 + + +class MoeCalibrationError(ValueError): + """Raised when calibrated MoE quantization cannot be performed correctly.""" + + +# --------------------------------------------------------------------------- +# recording forward implementation +# --------------------------------------------------------------------------- + +# id(experts_module) -> _ExpertsRecorder, populated only while a layer is being calibrated. +_ACTIVE_RECORDERS: dict[int, _ExpertsRecorder] = {} + +# Session ownership and recorder registration are process-local because both mutate model +# objects in place. One lock makes check-and-claim atomic across concurrent calibration +# threads. +_OWNERSHIP_LOCK = threading.Lock() +_ACTIVE_MODEL_SESSIONS: dict[int, object] = {} +_ACTIVE_EXPERT_SESSIONS: dict[int, object] = {} + + +def olive_moe_calib_experts_forward( + self: torch.nn.Module, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Run the reference per-expert experts forward, recording per-expert calibration inputs. + + Registered in transformers' ``ALL_EXPERTS_FUNCTIONS`` registry, so it replaces the + experts forward of *every* decorated experts module while the calibration + implementation is active. It reproduces the canonical eager loop shared by every + supported architecture (``F.linear`` on ``W[e]`` of shape ``(out, in)``, gated + activation, routing-weighted scatter-add), so outputs are correct on both the + Hessian-collection pass and the post-quantization re-run of the true-sequential loop. + + Recording is driven by :data:`_ACTIVE_RECORDERS`: only experts modules registered by + :meth:`MoeCalibrationSession.record` are recorded, and only while that context is + active -- this is what keeps the second (post-quantization) ``run_layer`` from + double-counting into the Hessians. + """ + recorder = _ACTIVE_RECORDERS.get(id(self)) + + if hidden_states.dim() != 2: + raise MoeCalibrationError( + "Olive's MoE calibration expects the experts forward to receive 2D " + f"(num_tokens, hidden_dim) hidden states, got shape {tuple(hidden_states.shape)}." + ) + + num_experts = self.num_experts + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + # pylint: disable=not-callable + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=num_experts).permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + if recorder is not None: + recorder.note_tokens(hidden_states.shape[0]) + + for hit in expert_hit: + expert_idx = hit[0] + if expert_idx == num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + intermediate = self._apply_gate( # pylint: disable=protected-access + torch.nn.functional.linear(current_state, self.gate_up_proj[expert_idx]) # pylint: disable=not-callable + ) + if recorder is not None: + recorder.record(int(expert_idx), {"gate_up_proj": current_state, "down_proj": intermediate}) + current_hidden_states = torch.nn.functional.linear( # pylint: disable=not-callable + intermediate, self.down_proj[expert_idx] + ) + current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) + + return final_hidden_states + + +def _register_calib_implementation() -> None: + """Register :func:`olive_moe_calib_experts_forward` in transformers' experts registry. + + The registry is a process-global singleton, so the key alone is not proof that *our* + function is what will run. Verify identity: if something else already owns the key, + calibration would silently execute a foreign forward (collecting wrong or zero + Hessians), so fail closed instead. + """ + from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS + + if OLIVE_MOE_CALIB_IMPLEMENTATION not in ALL_EXPERTS_FUNCTIONS: + ALL_EXPERTS_FUNCTIONS.register(OLIVE_MOE_CALIB_IMPLEMENTATION, olive_moe_calib_experts_forward) + return + + registered = ALL_EXPERTS_FUNCTIONS[OLIVE_MOE_CALIB_IMPLEMENTATION] + if registered is not olive_moe_calib_experts_forward: + raise MoeCalibrationError( + f"transformers' experts registry already maps '{OLIVE_MOE_CALIB_IMPLEMENTATION}' to " + f"{registered!r}, which is not Olive's recording forward " + f"({olive_moe_calib_experts_forward!r}). Calibration would silently run the wrong " + "experts forward, so it is refused. Remove the conflicting registration (or restart " + "the process), or re-run with moe=False." + ) + + +# --------------------------------------------------------------------------- +# per-expert recording +# --------------------------------------------------------------------------- + + +class _ExpertsRecorder: + """Accumulates one independent Hessian per (parameter, expert) for a single experts module. + + Each expert gets its *own* ``(K, K)`` Hessian -- activations are never pooled across + experts, because different experts see systematically different activation + distributions by routing design. Hessians are allocated lazily on an expert's first + sample, so experts that never get routed simply have no entry (and are handed to the + RTN fallback). + + The accumulated state is written straight onto each parameter's + ``quant_info.data`` so the quantization math (``Gptq.process_module``) reads it the + same way it reads the dense path's ``{"H": ..., "N": ...}``. + """ + + def __init__(self, experts: torch.nn.Module, pnames: list[str]): + self.experts = experts + self.pnames = pnames + self.num_experts = int(experts.num_experts) + self.tokens_seen = 0 + # pname -> {expert_idx: {"H": (K, K) tensor, "N": int}} + self.hessians: dict[str, dict[int, dict]] = {pname: {} for pname in pnames} + + def note_tokens(self, num_tokens: int) -> None: + self.tokens_seen += int(num_tokens) + + @torch.no_grad() + def record(self, expert_idx: int, inputs: dict[str, torch.Tensor]) -> None: + """Accumulate one expert's activation slice into that expert's Hessian.""" + for pname in self.pnames: + inp = inputs.get(pname) + if inp is None: + continue + self._accumulate(pname, expert_idx, inp) + + def token_counts(self) -> list[int]: + """Return per-expert routed-token counts, derived from the recorded sample counts. + + Every recorded parameter sees exactly one activation row per (token, expert) + routing decision, so the Hessian sample count ``N`` *is* the routed-token count. + Deriving the counts here (rather than incrementing a counter keyed on a hardcoded + parameter name) keeps coverage reporting correct when e.g. ``modules_to_not_convert`` + leaves only ``down_proj`` quantized. + """ + counts = [0] * self.num_experts + for expert_hessians in self.hessians.values(): + for expert_idx, entry in expert_hessians.items(): + if expert_idx < self.num_experts: + counts[expert_idx] = max(counts[expert_idx], int(entry["N"])) + return counts + + @torch.no_grad() + def _accumulate(self, pname: str, expert_idx: int, inp: torch.Tensor) -> None: + num_cols = inp.shape[-1] + entry = self.hessians[pname].get(expert_idx) + if entry is None: + entry = {"H": torch.zeros((num_cols, num_cols), device=inp.device, dtype=torch.float32), "N": 0} + self.hessians[pname][expert_idx] = entry + + num_rows = inp.shape[0] + if num_rows == 0: + return + x = inp.reshape(-1, num_cols).t() + entry["H"] *= entry["N"] / (entry["N"] + num_rows) + entry["N"] += num_rows + x = math.sqrt(2 / entry["N"]) * x.float() + entry["H"] += x.matmul(x.t()) + + def publish(self) -> None: + """Write the collected state onto the parameters' ``quant_info.data``.""" + token_counts = self.token_counts() + for pname in self.pnames: + param = self.experts._parameters[pname] # pylint: disable=protected-access + param.quant_info.data = { + "moe": True, + "experts": self.hessians[pname], + "tokens_seen": self.tokens_seen, + "token_counts": list(token_counts), + } + + +# --------------------------------------------------------------------------- +# coverage report +# --------------------------------------------------------------------------- + + +@dataclass +class LayerCoverage: + """Per-layer routing coverage collected during calibration. + + An expert counts as "starved" if its observed token count is below EITHER threshold -- + the routing-skew threshold (a fraction of this layer's calibration tokens) or the + statistical-sufficiency threshold (a multiple of K) -- matching the dual-condition + fallback gate in ``Gptq._process_moe_param``. + """ + + layer_name: str + num_experts: int + tokens_seen: int + token_counts: list[int] + skew_threshold: float + k_threshold: float + + @property + def threshold(self) -> float: + """The effective (combined) starved-cutoff: an expert starves below EITHER threshold.""" + return max(self.skew_threshold, self.k_threshold) + + @property + def unseen(self) -> int: + return sum(1 for c in self.token_counts if c == 0) + + @property + def starved(self) -> int: + return sum(1 for c in self.token_counts if 0 < c < self.threshold) + + @property + def covered(self) -> int: + return self.num_experts - self.unseen - self.starved + + def format(self) -> str: + counts = sorted(self.token_counts) + median = counts[len(counts) // 2] if counts else 0 + return ( + f"MoE coverage [{self.layer_name}]: {self.covered}/{self.num_experts} experts covered, " + f"{self.starved} starved (< {self.threshold:.1f} tokens = " + f"max(skew {self.skew_threshold:.1f}, sufficiency {self.k_threshold:.1f})), " + f"{self.unseen} unseen; tokens/expert min={min(counts, default=0)} median={median} " + f"max={max(counts, default=0)} (calibration tokens reaching the layer: {self.tokens_seen})" + ) + + +@dataclass +class CoverageReport: + """Aggregated routing coverage across every MoE layer of a calibration run.""" + + layers: list[LayerCoverage] = field(default_factory=list) + + def add(self, coverage: LayerCoverage) -> None: + self.layers.append(coverage) + logger.info("%s", coverage.format()) + + def format_summary(self) -> str: + if not self.layers: + return "MoE coverage summary: no MoE layers were calibrated." + total = sum(lc.num_experts for lc in self.layers) + starved = sum(lc.starved for lc in self.layers) + unseen = sum(lc.unseen for lc in self.layers) + return ( + f"MoE coverage summary: {len(self.layers)} MoE layers, {total} experts total, " + f"{starved} starved ({100 * starved / total:.1f}%), {unseen} unseen " + f"({100 * unseen / total:.1f}%). Starved/unseen experts are quantized with the " + "RTN fallback instead of GPTQ." + ) + + def log_summary(self) -> None: + message = self.format_summary() + if any(lc.starved or lc.unseen for lc in self.layers): + logger.warning("%s", message) + else: + logger.info("%s", message) + + +# --------------------------------------------------------------------------- +# support gating +# --------------------------------------------------------------------------- + + +def _transformers_supports_experts_registry() -> bool: + try: + import transformers.integrations.moe as transformers_moe + except ImportError: + return False + return hasattr(transformers_moe, "ALL_EXPERTS_FUNCTIONS") + + +def check_moe_gptq_support(model_type: str, experts_modules: list[torch.nn.Module]) -> None: + """Fail closed unless calibrated MoE quantization can safely record this model. + + GPTQ calibration intercepts the experts forward through transformers' experts- + implementation registry, so that registry must be available and every experts module + must be a registry-compatible fused implementation. GPTQ groups weights and constructs + Hessians along the last dimension, so fused experts must report a K-last layout through + the shared ``is_transposed`` metadata check. Finally, the recording forward implements + only bias-free experts with a gated activation; modules that declare biases or a + non-gated activation are refused rather than calibrated with the wrong computation. + + Layout support is independent of ``model_type``: transformers-owned + ``is_transposed=False`` metadata is sufficient for any fused-experts architecture. + + Also logs a warning when the estimated per-layer Hessian memory is large enough to risk + an out-of-memory during calibration. + + No weight is touched before this returns, mirroring the fail-closed guards in + :mod:`olive.common.quant.selection`. + + Raises: + MoeCalibrationError: If the experts registry is unavailable, an experts module + cannot be intercepted through that registry, the fused-experts layout cannot be + proven K-last, or an experts module declares bias or a non-gated activation. + + """ + if not _transformers_supports_experts_registry(): + raise MoeCalibrationError( + "Calibrated MoE quantization (moe=True) requires transformers >= " + f"{MIN_TRANSFORMERS_VERSION}, which provides the experts-implementation registry " + "(transformers.integrations.moe.ALL_EXPERTS_FUNCTIONS) used to collect per-expert " + "activations. Upgrade transformers, or re-run with moe=False." + ) + + try: + check_moe_layout_support( + experts_modules, + model_type=model_type, + operation="GPTQ MoE calibration", + ) + except MoeSupportError as exc: + raise MoeCalibrationError(str(exc)) from exc + + for experts in experts_modules: + if isinstance(experts, torch.nn.ModuleList) or not hasattr(experts, "config"): + raise MoeCalibrationError( + f"Experts module '{type(experts).__name__}' is not a fused-experts module registered " + "with transformers' experts-implementation registry, so GPTQ cannot intercept its " + "forward to record per-expert activations. Re-run with moe=False, or quantize the " + "experts with the Rtn pass." + ) + if getattr(experts, "has_bias", False): + raise MoeCalibrationError( + f"Experts module '{type(experts).__name__}' declares expert biases, which Olive's " + "calibrated MoE path does not handle. Re-run with moe=False." + ) + if not getattr(experts, "has_gate", True): + raise MoeCalibrationError( + f"Experts module '{type(experts).__name__}' declares non-gated experts, which " + "Olive's calibrated MoE path does not handle. Re-run with moe=False." + ) + + _warn_on_hessian_memory(experts_modules) + + +def _estimate_layer_hessian_bytes(experts: torch.nn.Module) -> int | None: + """Estimate the peak per-layer Hessian working set in bytes, or ``None`` if unknown. + + One float32 ``(K, K)`` Hessian is held per (expert, quantized parameter), where ``K`` is + that parameter's input dim: ``hidden_size`` for ``gate_up_proj`` and + ``moe_intermediate_size`` for ``down_proj``. + """ + total_cols_squared = 0 + for pname in ("gate_up_proj", "down_proj"): + param = getattr(experts, pname, None) + if param is None or not hasattr(param, "shape") or len(param.shape) != 3: + return None + total_cols_squared += int(param.shape[-1]) ** 2 + num_experts = getattr(experts, "num_experts", None) + if num_experts is None: + return None + return int(num_experts) * total_cols_squared * 4 # float32 + + +def _warn_on_hessian_memory(experts_modules: list[torch.nn.Module]) -> None: + """Warn up front when per-layer Hessian memory is likely to exhaust the device. + + Hessians are allocated per expert on the calibration device and only freed once the + layer is quantized, so the peak is a whole layer's worth. For DeepSeek-V3-class configs + (256 experts, hidden 7168, moe_intermediate 2048) this is ~57 GB -- an OOM that would + otherwise surface as a raw CUDA error minutes into calibration with no explanation. + """ + estimates = [(experts, nbytes) for experts in experts_modules if (nbytes := _estimate_layer_hessian_bytes(experts))] + if not estimates: + return + experts, peak = max(estimates, key=lambda item: item[1]) + if peak <= MOE_HESSIAN_MEMORY_WARN_BYTES: + return + logger.warning( + "Calibrated MoE quantization will allocate up to %.1f GiB of float32 Hessians for a " + "single '%s' layer (%d experts x (%d^2 + %d^2) x 4 bytes), held on the calibration " + "device on top of the layer's weights and activations. This exceeds the %.1f GiB " + "warning threshold and may run out of memory. Consider calibrating on CPU " + "(device='cpu'), reducing the number of quantized MoE parameters via " + "modules_to_not_convert, or re-running with moe=False.", + peak / 1024**3, + type(experts).__name__, + int(experts.num_experts), + int(experts.gate_up_proj.shape[-1]), + int(experts.down_proj.shape[-1]), + MOE_HESSIAN_MEMORY_WARN_BYTES / 1024**3, + ) + + +# --------------------------------------------------------------------------- +# session +# --------------------------------------------------------------------------- + + +class MoeCalibrationSession: + """Owns the experts-implementation swap and the per-layer recording lifecycle. + + Usage (see :func:`olive.passes.pytorch.quant_utils.run_layerwise_quantization`):: + + session = MoeCalibrationSession.create(wrapper, fallback_threshold=0.005) + if session: + session.start() # swap in the recording experts implementation + ... + with session.record(experts_modules): # recording ON for one layer + run_layer(...) # Hessian collection pass + run_layer(...) # true-sequential re-run: recording OFF + ... + session.finish() # restore the original implementation + log summary + """ + + def __init__( + self, + model: torch.nn.Module, + fallback_threshold: float = DEFAULT_MOE_FALLBACK_THRESHOLD, + fallback_min_k_multiple: float = DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE, + ): + self.model = model + self.fallback_threshold = fallback_threshold + self.fallback_min_k_multiple = fallback_min_k_multiple + self.report = CoverageReport() + self.experts_modules: list[torch.nn.Module] = [] + self._saved_implementation = None + self._active = False + self._owned_expert_ids: set[int] = set() + + @classmethod + def create( + cls, + wrapper: ModelWrapper, + fallback_threshold: float = DEFAULT_MOE_FALLBACK_THRESHOLD, + fallback_min_k_multiple: float = DEFAULT_MOE_FALLBACK_MIN_K_MULTIPLE, + ) -> MoeCalibrationSession | None: + """Validate support and build a session, or return ``None`` when the model has no experts.""" + experts_modules = [ + experts for lw in wrapper.get_layer_wrappers() if (experts := lw.get_experts(return_name=False)) is not None + ] + if not experts_modules: + return None + + check_moe_gptq_support(wrapper.model_type, experts_modules) + session = cls( + wrapper.model, fallback_threshold=fallback_threshold, fallback_min_k_multiple=fallback_min_k_multiple + ) + session.experts_modules = experts_modules + return session + + def start(self) -> None: + """Swap the model's experts implementation to the recording one. + + Re-entrancy is refused: a second ``start()`` would overwrite + ``_saved_implementation`` with ``"olive_moe_calib"``, so the eventual ``finish()`` + would "restore" the calibration implementation instead of the model's original one. + Any failure after the swap restores the model before propagating, so a caller that + aborts here never leaves the model in calibration mode. + """ + if not hasattr(self.model, "set_experts_implementation"): + raise MoeCalibrationError( + "Calibrated MoE quantization (moe=True) requires transformers >= " + f"{MIN_TRANSFORMERS_VERSION}: the loaded model has no " + "'set_experts_implementation'. Upgrade transformers, or re-run with moe=False." + ) + self._claim_ownership() + try: + self._start_owned() + except BaseException: + self._release_ownership() + raise + + def _claim_ownership(self) -> None: + """Atomically reserve this model and its experts for the session lifetime.""" + with _OWNERSHIP_LOCK: + if self._active: + raise MoeCalibrationError( + "MoE calibration is already active for this model; MoeCalibrationSession.start() " + "cannot be nested or called twice (the original experts implementation would be " + "lost). Call finish() before starting a new session." + ) + model_owner = _ACTIVE_MODEL_SESSIONS.get(id(self.model)) + expert_owner = next( + ( + owner + for experts in self.experts_modules + if (owner := _ACTIVE_EXPERT_SESSIONS.get(id(experts))) is not None + ), + None, + ) + if model_owner is not None or expert_owner is not None: + raise MoeCalibrationError( + "There is another MoE calibration session active for this model or one of its " + "experts modules. Concurrent/nested calibration sessions are not supported." + ) + self._owned_expert_ids = {id(experts) for experts in self.experts_modules} + _ACTIVE_MODEL_SESSIONS[id(self.model)] = self + for experts_id in self._owned_expert_ids: + _ACTIVE_EXPERT_SESSIONS[experts_id] = self + + def _release_ownership(self) -> None: + """Release only registry entries still owned by this session.""" + with _OWNERSHIP_LOCK: + if _ACTIVE_MODEL_SESSIONS.get(id(self.model)) is self: + _ACTIVE_MODEL_SESSIONS.pop(id(self.model)) + for experts_id in self._owned_expert_ids: + if _ACTIVE_EXPERT_SESSIONS.get(experts_id) is self: + _ACTIVE_EXPERT_SESSIONS.pop(experts_id) + self._owned_expert_ids.clear() + + def _start_owned(self) -> None: + """Start calibration after this session has atomically claimed ownership.""" + _register_calib_implementation() + saved_implementation = self.model.get_experts_implementation() + # transformers returns a dict ({"": impl, : impl}); older/simpler models + # may return a plain string. + saved_values = ( + list(saved_implementation.values()) if isinstance(saved_implementation, dict) else [saved_implementation] + ) + if OLIVE_MOE_CALIB_IMPLEMENTATION in saved_values: + raise MoeCalibrationError( + f"The model is already using the '{OLIVE_MOE_CALIB_IMPLEMENTATION}' experts " + "implementation, which means another MoE calibration session is active on it. " + "Concurrent/nested calibration sessions are not supported." + ) + + try: + self.model.set_experts_implementation(OLIVE_MOE_CALIB_IMPLEMENTATION) + except Exception: + # transformers can mutate some submodels' configs before raising while + # configuring a later one (e.g. a heterogeneous/composite model); attempt a + # best-effort restore of whatever was already swapped, but always propagate the + # original error so the caller sees why start() actually failed. + try: + self.model.set_experts_implementation(saved_implementation) + except Exception: # pylint: disable=broad-except + logger.warning( + "Failed to restore the original experts implementation after a MoE " + "calibration start() error; the model's experts implementation may be left " + "in an inconsistent state.", + exc_info=True, + ) + raise + try: + # ``set_experts_implementation`` silently no-ops when transformers' source-inspection + # heuristic decides the class isn't switchable. Verify the swap actually reached every + # experts module rather than silently collecting zero Hessians. + stale = sorted( + { + type(experts).__name__ + for experts in self.experts_modules + if getattr(experts.config, "_experts_implementation", None) != OLIVE_MOE_CALIB_IMPLEMENTATION + } + ) + except Exception: + # Best-effort restore: transformers' own setter can itself mutate some submodels' + # configs before raising on a later one, so even the "verify the swap landed" step + # (not just the setter call above) can observe a partially-swapped model. Attempt to + # restore regardless, but propagate the original error either way -- a failed restore + # attempt here should not mask why ``start()`` failed in the first place. + try: + self.model.set_experts_implementation(saved_implementation) + except Exception: # pylint: disable=broad-except + logger.warning( + "Failed to restore the original experts implementation after a MoE " + "calibration start() error; the model's experts implementation may be left " + "in an inconsistent state.", + exc_info=True, + ) + raise + + if stale: + try: + self.model.set_experts_implementation(saved_implementation) + except Exception as exc: + logger.warning( + "Failed to restore the original experts implementation after the " + "calibration implementation remained stale for %s; the model may be left " + "in an inconsistent state.", + stale, + exc_info=True, + ) + raise MoeCalibrationError( + "Olive could not switch the experts implementation to " + f"'{OLIVE_MOE_CALIB_IMPLEMENTATION}' for {stale}, so per-expert calibration " + "data cannot be collected; restoring the original experts implementation " + f"also failed ({type(exc).__name__}: {exc}). Re-run with moe=False." + ) from exc + raise MoeCalibrationError( + "Olive could not switch the experts implementation to " + f"'{OLIVE_MOE_CALIB_IMPLEMENTATION}' for {stale}; per-expert " + "calibration data cannot be collected. Re-run with moe=False." + ) + + self._saved_implementation = saved_implementation + self._active = True + logger.debug("Switched experts implementation to '%s' for calibration.", OLIVE_MOE_CALIB_IMPLEMENTATION) + + @contextmanager + def record(self, experts_modules: list[torch.nn.Module]) -> Iterator[None]: + """Record per-expert Hessians for ``experts_modules`` for the duration of the block.""" + recorders = [] + for experts in experts_modules: + pnames = [ + pname + for pname, param in experts.named_parameters(recurse=False) + if param is not None and hasattr(param, "quant_info") + ] + recorder = _ExpertsRecorder(experts, pnames) + recorders.append(recorder) + + with _OWNERSHIP_LOCK: + if not self._active: + raise MoeCalibrationError("record() requires an active session; call start() first.") + conflicts = [type(experts).__name__ for experts in experts_modules if id(experts) in _ACTIVE_RECORDERS] + if conflicts: + raise MoeCalibrationError( + f"Per-expert recording is already active for {sorted(set(conflicts))}; " + "nested or concurrent record() contexts are not supported." + ) + for experts, recorder in zip(experts_modules, recorders): + _ACTIVE_RECORDERS[id(experts)] = recorder + try: + yield + finally: + with _OWNERSHIP_LOCK: + for experts, recorder in zip(experts_modules, recorders): + if _ACTIVE_RECORDERS.get(id(experts)) is recorder: + _ACTIVE_RECORDERS.pop(id(experts)) + for recorder in recorders: + recorder.publish() + + def add_coverage(self, layer_name: str, experts: torch.nn.Module) -> None: + """Log (and remember) the routing coverage recorded for one experts module. + + NOTE: the sufficiency threshold (``k_threshold``) is derived from ``pnames[0]``'s + last-dim size only (typically ``gate_up_proj``). Parameters on the same layer with a + different last-dim size (e.g. ``down_proj``, whose K is the intermediate size rather + than the hidden size) are NOT separately represented in this report -- the actual + per-expert fallback gate in ``Gptq._process_moe_param`` does use each parameter's own + K, so this coverage summary may under- or over-count "starved" experts relative to + what actually happened for parameters other than ``pnames[0]``. + """ + pnames = [ + pname + for pname, param in experts.named_parameters(recurse=False) + if param is not None and hasattr(param, "quant_info") + ] + if not pnames: + return + data = experts._parameters[pnames[0]].quant_info.data # pylint: disable=protected-access + if not data: + return + k = experts._parameters[pnames[0]].shape[-1] # pylint: disable=protected-access + self.report.add( + LayerCoverage( + layer_name=layer_name, + num_experts=len(data["token_counts"]), + tokens_seen=data["tokens_seen"], + token_counts=data["token_counts"], + skew_threshold=self.fallback_threshold * data["tokens_seen"], + k_threshold=self.fallback_min_k_multiple * k, + ) + ) + + def finish(self) -> None: + """Restore the original experts implementation and log the coverage summary.""" + try: + if self._active: + self.model.set_experts_implementation(self._saved_implementation) + finally: + # Clear session state and log whatever coverage was recorded regardless of + # whether the restore above succeeded -- a failed restore should not also + # suppress the coverage summary or leave ``_active``/``_saved_implementation`` + # in a way that could make a later ``start()`` behave inconsistently. + self._saved_implementation = None + self._active = False + self._release_ownership() + self.report.log_summary() diff --git a/olive/passes/pytorch/moe_support.py b/olive/passes/pytorch/moe_support.py new file mode 100644 index 0000000000..a71b9d9780 --- /dev/null +++ b/olive/passes/pytorch/moe_support.py @@ -0,0 +1,95 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Shared layout-safety checks for fused MoE expert quantization.""" + +from __future__ import annotations + +import torch + + +class MoeSupportError(ValueError): + """Raised when an MoE quantization layout cannot be proven safe to quantize.""" + + +def check_moe_layout_support( + experts_modules: list[torch.nn.Module], + *, + model_type: str, + operation: str, +) -> None: + """Fail closed unless every fused-experts module reports a K-last layout. + + The quantizers using this check group along each weight tensor's last dimension. + Consequently, only fused weights stored ``(num_experts, out_features, in_features)`` + are safe: the input/contraction dimension K must be last. + + ``is_transposed`` is a boolean attribute exposed by ``transformers``'s + ``use_experts_implementation`` decorator on fused-experts modules (``False`` for K-last + architectures such as Mixtral/Qwen3-MoE/DeepSeek-V3/etc., ``True`` for architectures such + as gpt-oss). It is assigned once, from the decorator's own argument, inside the wrapped + ``__init__`` -- not derived from ``config`` -- so no checkpoint or config field can + influence it for a standard (non-``trust_remote_code``) experts implementation. Its + absence indicates either an undecorated experts implementation (e.g. an older + ``transformers`` release, or an architecture such as llama4/aria that has not adopted the + fused-experts decorator) or an unrecognized implementation, so it is treated as unsafe. + + Classic per-expert ``nn.ModuleList`` experts (e.g. PhiMoE, DeepSeek-V3, or Mixtral on + older ``transformers`` releases without the fused-experts refactor) are exempt only when + the ``ModuleList`` itself owns no direct 3D parameter: Olive's quantizer selection only + ever groups a fused-experts module's *own* 3D parameters, and a plain ``ModuleList`` + container cannot have one (each per-expert child is a plain 2D ``nn.Linear``, whose + weight is unconditionally ``(out, in)`` with K last). A ``ModuleList`` that does carry a + direct 3D parameter falls through to the normal ``is_transposed`` check below, which + rejects it (such a parameter has no ``is_transposed`` metadata to trust). + + Note: this trusts ``is_transposed`` as reported by the resolved experts class. A custom + experts implementation loaded via ``trust_remote_code`` that misreports its own layout + (e.g. sets ``is_transposed=False`` while actually storing transposed weights) is not + caught here; that is treated as user-introduced misuse of an explicitly opted-in trust + boundary, not a layout Olive can independently verify. + + Args: + experts_modules: The fused-experts (or classic ``nn.ModuleList``) modules discovered + for this model, one per MoE decoder layer. + model_type: The resolved HF ``config.model_type`` of the model being quantized. Not + used to decide the layout (that is entirely driven by ``is_transposed``); included + only for diagnostic context in error messages. + operation: Human-readable label for the calling pass/operation, prefixed onto error + messages. Purely cosmetic: it does not affect validation logic. + + Raises: + MoeSupportError: If a fused-experts module's ``is_transposed`` attribute is missing, + not a ``bool``, or ``True``. + + """ + for experts in experts_modules: + if isinstance(experts, torch.nn.ModuleList) and not any( + param.dim() == 3 for param in experts.parameters(recurse=False) + ): + # No direct 3D parameter: Olive's quantizer selection never groups anything on + # this module itself, only on its per-expert nn.Linear children (K is + # unconditionally last for those), so there is no transposed-layout risk here. + continue + + actual_class = type(experts).__name__ + is_transposed = getattr(experts, "is_transposed", None) + if not isinstance(is_transposed, bool): + raise MoeSupportError( + f"{operation} cannot verify the layout of experts module '{actual_class}' " + f"(model_type='{model_type}') because attribute 'is_transposed' is missing " + "or not a boolean. This typically means an experts implementation that has " + "not adopted the fused-experts decorator (e.g. an older transformers release, " + "or an architecture such as llama4/aria), or an unrecognized implementation. " + "Quantization is refused rather than risking grouping along the wrong " + "dimension. Re-run with moe=False so the experts stay in full precision." + ) + if is_transposed: + raise MoeSupportError( + f"{operation} refuses experts module '{actual_class}' (model_type='{model_type}') " + "because it reports a transposed fused-weight layout (E, K, OUT), while " + "last-dimension grouping requires (E, OUT, K) with K last. Architectures such " + "as gpt-oss store this transposed layout. Re-run with moe=False so the " + "experts stay in full precision." + ) diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py index 4d89e2b338..17c0371532 100644 --- a/olive/passes/pytorch/quant_utils.py +++ b/olive/passes/pytorch/quant_utils.py @@ -2,8 +2,11 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +# pylint: disable=protected-access from __future__ import annotations +import contextlib +import inspect import logging from copy import deepcopy from dataclasses import dataclass @@ -16,10 +19,11 @@ OliveHfQuantizationConfig, OliveHfQuantizationMethod, OliveHfQuantizationOverrideConfig, - replace_matching_submodules, tie_quant_word_embeddings, ) -from olive.common.quant.nn import QuantEmbedding, QuantLinear +from olive.common.quant.selection import iter_quant_targets +from olive.common.quant.state_dict import install_quant_tensor_param +from olive.common.quant.tensor import QuantTensor from olive.common.quant.utils import WeightQuantizer from olive.common.utils import tensor_data_to_device from olive.constants import PrecisionBits @@ -31,12 +35,13 @@ if TYPE_CHECKING: from olive.model import HfModelHandler from olive.passes.pass_config import BasePassConfig + from olive.passes.pytorch.moe_calib import MoeCalibrationSession logger = logging.getLogger(__name__) -def get_quantizer_config(allow_embeds: bool = False) -> dict[str, PassConfigParam]: +def get_quantizer_config(allow_embeds: bool = False, allow_moe: bool = False) -> dict[str, PassConfigParam]: return { "bits": PassConfigParam( type_=PrecisionBits, @@ -62,6 +67,17 @@ def get_quantizer_config(allow_embeds: bool = False) -> dict[str, PassConfigPara search_defaults=Boolean(), description="Whether to quantize the language model head. Default value is False.", ), + "quantize_vision": PassConfigParam( + type_=bool, + default_value=False, + description=( + "Whether to quantize a composite vision-language model's vision tower in this pass. When " + "False (default), the vision tower (``visual``/``vision_tower``/``vision_model``/" + "``vision_encoder``) is left in full precision -- the typical Olive pipeline quantizes it " + "separately downstream (e.g. on the ONNX side). Set to True to quantize the vision tower " + "here too, e.g. when this pass is the only quantization step for the model." + ), + ), **( { "embeds": PassConfigParam( @@ -73,13 +89,37 @@ def get_quantizer_config(allow_embeds: bool = False) -> dict[str, PassConfigPara if allow_embeds else {} ), + **( + { + "moe": PassConfigParam( + type_=bool, + default_value=False, + description=( + "Whether to quantize MoE expert modules / parameters. When False (default), every " + "nn.Module under each experts subtree is skipped. Defaults reuse the pass-level " + "bits/group_size/sym settings; use ``overrides`` to tune experts independently." + ), + ) + } + if allow_moe + else {} + ), + "modules_to_not_convert": PassConfigParam( + type_=list, + default_value=None, + description=( + "Optional list of module name patterns to exclude from quantization. Plain strings use" + " substring matching (HF semantics); entries prefixed with 're:' use re.fullmatch." + ), + ), "overrides": PassConfigParam( type_=dict, default_value=None, description=( - "Optional dictionary to specify overrides for specific modules. The keys are module names and the" - " values are dictionaries with any of the following keys: 'bits', 'symmetric', 'group_size'. These" - " overrides take precedence over the overrides provided in the mixed precision info." + "Optional dictionary to specify overrides for specific modules. The keys are module names" + " (literal match) or 're:' patterns (regex fullmatch). Values are dictionaries with" + " any of the following keys: 'bits', 'symmetric', 'group_size'. These overrides take" + " precedence over the overrides provided in the mixed precision info." ), ), } @@ -96,7 +136,10 @@ def get_qkv_quantization_groups(wrapper: ModelWrapper, module_names: set[str] | module_to_name = {id(module): name for name, module in wrapper.model.named_modules()} qkv_groups = [] for layer_wrapper in wrapper.get_layer_wrappers(): - attn_inputs, _ = layer_wrapper.get_attention_inputs() + # partial_ok: MLA-style attentions (deepseek_v3) expose only a subset of + # q/k/v_proj; QKV grouping treats the result as an unordered set, so dropping the + # missing entries is safe here (unlike the positional consumers in rotate.py). + attn_inputs, _ = layer_wrapper.get_attention_inputs(partial_ok=True) group = tuple( name for name in (module_to_name.get(id(module)) for module in attn_inputs) @@ -177,7 +220,10 @@ def normalize_qkv_quant_config( def _collect_excluded_attn_inputs(wrapper: ModelWrapper) -> set[torch.nn.Module]: excluded: set[torch.nn.Module] = set() for layer_wrapper in wrapper.get_layer_wrappers(): - attn_inputs, _ = layer_wrapper.get_attention_inputs() + # partial_ok: MLA-style attentions (deepseek_v3) expose only a subset of + # q/k/v_proj; QKV grouping treats the result as an unordered set, so dropping the + # missing entries is safe here (unlike the positional consumers in rotate.py). + attn_inputs, _ = layer_wrapper.get_attention_inputs(partial_ok=True) if len(attn_inputs) == 1: excluded.add(attn_inputs[0]) else: @@ -185,6 +231,27 @@ def _collect_excluded_attn_inputs(wrapper: ModelWrapper) -> set[torch.nn.Module] return excluded +def _collect_already_quantized_names(model: torch.nn.Module) -> set[str]: + """Return override-key names of parameters already backed by a ``QuantTensor``. + + Names follow the same convention as :func:`iter_quant_targets`: ``module_name`` for + the ``weight`` parameter of ``nn.Linear`` / ``nn.Embedding`` and ``f"{name}.{pname}"`` + otherwise. These names lock the corresponding modules against re-quantization and QKV + renormalization when merging with an existing checkpoint. + """ + names: set[str] = set() + for name, module in model.named_modules(): + for pname, param in module.named_parameters(recurse=False): + if param is None: + continue + if isinstance(param, QuantTensor) or isinstance(getattr(param, "data", None), QuantTensor): + if pname == "weight" and isinstance(module, (torch.nn.Linear, torch.nn.Embedding)): + names.add(name) + else: + names.add(f"{name}.{pname}" if name else pname) + return names + + def prepare_model( model: HfModelHandler, config: type[BasePassConfig], @@ -226,29 +293,31 @@ def prepare_model( lm_head_name = wrapper.get_lm_head()[1] embeds_name = wrapper.get_embeds()[1][0] - def should_quantize(module: torch.nn.Module, name: str) -> bool: - if module in excluded_attn_inputs: - return False - if isinstance(module, torch.nn.Linear): - return name != lm_head_name or fresh_qcfg.lm_head - if fresh_qcfg.embeds and isinstance(module, torch.nn.Embedding): - return name == embeds_name - return False + skip_patterns = list(getattr(fresh_qcfg, "modules_to_not_convert", None) or []) # Pre-existing quantized weights are immutable. If we're merging with an existing # checkpoint, build the final qcfg first (merge fresh into existing, then renormalize - # QKV with already-quantized modules locked) so that the quant_info we attach below - # uses the same settings the on-disk fusion will require. Every module that is already - # a QuantLinear/QuantEmbedding after load is on-disk-immutable, including those that - # used the existing config's defaults (no explicit override entry). + # QKV with already-quantized parameters locked) so that the quant_info we attach below + # uses the same settings the on-disk fusion will require. Every parameter that is already + # a ``QuantTensor`` after load is on-disk-immutable, including those that used the + # existing config's defaults (no explicit override entry). on_disk_overrides: set[str] = set() already_quantized: set[str] = set() if existing_qcfg: on_disk_overrides = set((existing_qcfg.get("overrides") or {}).keys()) - already_quantized = { - name for name, module in wrapper.model.named_modules() if isinstance(module, (QuantLinear, QuantEmbedding)) + already_quantized = _collect_already_quantized_names(wrapper.model) + fresh_names = { + full_name + for _, _, full_name in iter_quant_targets( + wrapper.model, + quantize_lm_head=fresh_qcfg.lm_head, + quantize_embeds=fresh_qcfg.embeds, + quantize_moe=getattr(fresh_qcfg, "moe", False), + quantize_vision=getattr(fresh_qcfg, "quantize_vision", False), + skip_patterns=skip_patterns, + extra_skip_modules=excluded_attn_inputs, + ) } - fresh_names = {name for name, module in wrapper.model.named_modules() if should_quantize(module, name)} merged = existing_qcfg merged["overrides"] = existing_qcfg.get("overrides") or {} for name in fresh_names: @@ -258,6 +327,10 @@ def should_quantize(module: torch.nn.Module, name: str) -> bool: merged["overrides"][name] = override merged["lm_head"] |= fresh_qcfg.lm_head merged["embeds"] |= fresh_qcfg.embeds + merged["moe"] = merged.get("moe", False) or getattr(fresh_qcfg, "moe", False) + merged["quantize_vision"] = merged.get("quantize_vision", False) or getattr( + fresh_qcfg, "quantize_vision", False + ) qcfg = OliveHfQuantizationConfig(**merged) qcfg = normalize_qkv_quant_config(wrapper, qcfg, locked_modules=already_quantized) else: @@ -265,14 +338,18 @@ def should_quantize(module: torch.nn.Module, name: str) -> bool: new_qargs: dict[str, dict[str, int | bool]] = {} - def add_quant_info(module: torch.nn.Module, name: str) -> torch.nn.Module: - # TODO(jambayk): validate that the module and config are compatible - qargs = qcfg.get_qlinear_init_args(name) - module.quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs)) - new_qargs[name] = qargs - return module - - replace_matching_submodules(wrapper.model, should_quantize, add_quant_info, description="Preparing model") + for module, pname, full_name in iter_quant_targets( + wrapper.model, + quantize_lm_head=qcfg.lm_head, + quantize_embeds=qcfg.embeds, + quantize_moe=getattr(qcfg, "moe", False), + quantize_vision=getattr(qcfg, "quantize_vision", False), + skip_patterns=skip_patterns, + extra_skip_modules=excluded_attn_inputs, + ): + qargs = qcfg.get_qlinear_init_args(full_name) + new_qargs[full_name] = qargs + module._parameters[pname].quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs)) # Drop overrides for modules that won't be quantized this pass. Pre-existing (on-disk) # overrides are preserved verbatim since they describe already-quantized weights. @@ -280,6 +357,9 @@ def add_quant_info(module: torch.nn.Module, name: str) -> torch.nn.Module: # follow-up pass runs, the quantized members in the group will be locked and pull the # remaining members back into the shared config via ``normalize_qkv_quant_config``. for name in list(qcfg.overrides or {}): + # ``re:`` keys aren't tied to a specific module, so leave them in place. + if name.startswith("re:"): + continue if name not in new_qargs and name not in on_disk_overrides: qcfg.overrides.pop(name) @@ -310,6 +390,9 @@ def get_quant_config(model: HfModelHandler, config: type[BasePassConfig]) -> Oli "group_size": config.group_size, "lm_head": config.lm_head, "embeds": getattr(config, "embeds", False), + "moe": getattr(config, "moe", False), + "quantize_vision": getattr(config, "quantize_vision", False), + "modules_to_not_convert": getattr(config, "modules_to_not_convert", None) or [], "overrides": config.overrides or {}, } if mp_info := (model.model_attributes or {}).get("mixed_precision_info"): @@ -385,15 +468,19 @@ def store_input_hook(_, args: tuple, kwargs: dict) -> None: first_layer = wrapper.get_layers(return_name=False)[0] hook = first_layer.register_forward_pre_hook(store_input_hook, with_kwargs=True) - for data in get_calibration_dataset(model, data_config): - try: - wrapper.model(**tensor_data_to_device(data, device)) - except ValueError: - pass - - hook.remove() - for module in pre_layer_modules: - module.to("cpu") + try: + for data in get_calibration_dataset(model, data_config): + try: + wrapper.model(**tensor_data_to_device(data, device)) + except ValueError: + # `store_input_hook` raises ValueError once it has captured the first layer's + # inputs, deliberately aborting the forward pass early since the rest of the + # model's computation isn't needed for calibration. + pass + finally: + hook.remove() + for module in pre_layer_modules: + module.to("cpu") return hidden_states, layer_args, layer_kwargs @@ -447,6 +534,7 @@ def run_layerwise_quantization( update_before_process: bool, include_lm_head: bool, device: str | None = None, + moe_session: MoeCalibrationSession | None = None, ) -> str: """Run a layerwise calibration + processing loop with configurable hook order. @@ -459,6 +547,12 @@ def run_layerwise_quantization( update_before_process: Whether to run the layer forward to get next inputs before processing. include_lm_head: Whether to process the lm_head similarly to other layers. device: Device to run calibration on. If None, uses cuda when available. + moe_session: Optional MoE calibration session. Required when any selected parameter + lives on a fused MoE experts module: routing happens *inside* one experts-module + forward call, so per-expert activations cannot be observed with an ordinary + forward hook. The session intercepts the experts forward instead and records one + independent Hessian per expert. Ordinary ``nn.Linear`` / ``nn.Embedding`` + targets keep using ``input_hook``. Returns: Device string used for calibration. @@ -473,65 +567,166 @@ def run_layerwise_quantization( if original_use_cache is not None: wrapper.model.config.use_cache = False - hidden_states, layer_args, layer_kwargs = get_layer_inputs_for_calibration(model, wrapper, data_config, device) - if not hidden_states: - raise ValueError("Calibration data is empty. Provide a valid data_config.") + # Everything below runs inside try/finally: the experts-implementation swap, the + # ``use_cache`` override, the forward hooks and the progress bar are all process-global + # mutations that must be undone even when calibration raises part way through. + pbar = None + handles: list = [] + try: + hidden_states, layer_args, layer_kwargs = get_layer_inputs_for_calibration(model, wrapper, data_config, device) + if not hidden_states: + raise ValueError("Calibration data is empty. Provide a valid data_config.") + + total_steps = wrapper.num_hidden_layers + (1 if include_lm_head else 0) + pbar = tqdm(total=total_steps, desc="Processing layers...") + + if moe_session is not None: + moe_session.start() + + layers_name = wrapper.get_layers(return_name=True)[1] + for layer_idx, layer in enumerate(wrapper.get_layers(return_name=False)): + pbar.set_postfix(module=f"layers.{layer_idx}", refresh=False) + dense_modules, moe_modules = _split_quantizable_modules(layer) + if moe_modules and moe_session is None: + raise ValueError( + "Fused MoE expert parameters were selected for calibrated quantization but no " + "MoE calibration session was provided. This is an internal error." + ) + handles = [module.register_forward_hook(input_hook) for module in dense_modules] - total_steps = wrapper.num_hidden_layers + (1 if include_lm_head else 0) - pbar = tqdm(total=total_steps, desc="Processing layers...") + record_ctx = ( + moe_session.record(moe_modules) if moe_session is not None and moe_modules else contextlib.nullcontext() + ) + with record_ctx: + if update_before_process: + hidden_states = run_layer( + layer, + hidden_states, + layer_args, + layer_kwargs, + return_output=True, + ) + else: + run_layer(layer, hidden_states, layer_args, layer_kwargs) + + for handle in handles: + handle.remove() + handles = [] + + if moe_session is not None: + layer_name = f"{layers_name}.{layer_idx}" + for experts in moe_modules: + moe_session.add_coverage(layer_name, experts) + + for module in [*dense_modules, *moe_modules]: + process_module(module, device) + + if not update_before_process: + # true-sequential: re-run the layer with the quantized weights so the next layer + # sees realistic inputs. Recording is off here (the ``record`` context above has + # exited), so Hessians are not double-counted. + hidden_states = run_layer( + layer, + hidden_states, + layer_args, + layer_kwargs, + return_output=True, + ) + + if torch.cuda.is_available(): + torch.cuda.empty_cache() - for layer_idx, layer in enumerate(wrapper.get_layers(return_name=False)): - pbar.set_postfix(module=f"layers.{layer_idx}", refresh=False) - quantizable_modules = [module for module in layer.modules() if hasattr(module, "quant_info")] - handles = [module.register_forward_hook(input_hook) for module in quantizable_modules] + pbar.update(1) - if update_before_process: + if include_lm_head: hidden_states = run_layer( - layer, - hidden_states, - layer_args, - layer_kwargs, - return_output=True, + wrapper.get_pre_head_layernorm(return_name=False), hidden_states, return_output=True ) - else: - run_layer(layer, hidden_states, layer_args, layer_kwargs) - + lm_head = wrapper.get_lm_head(return_name=False) + pbar.set_postfix(module="lm_head", refresh=False) + handles = [lm_head.register_forward_hook(input_hook)] + run_layer(lm_head, hidden_states, return_output=True) + for handle in handles: + handle.remove() + handles = [] + process_module(lm_head, device) + pbar.update(1) + finally: for handle in handles: handle.remove() + if pbar is not None: + pbar.close() + try: + if moe_session is not None: + moe_session.finish() + finally: + # Always restore ``use_cache`` even if ``moe_session.finish()`` itself raises + # (e.g. the experts-implementation restore call fails) -- these are two + # independent process-global mutations and a failure in one must not leave the + # other un-undone. + if original_use_cache is not None: + wrapper.model.config.use_cache = original_use_cache - for module in quantizable_modules: - process_module(module, device) + return device - if not update_before_process: - hidden_states = run_layer( - layer, - hidden_states, - layer_args, - layer_kwargs, - return_output=True, - ) - if torch.cuda.is_available(): - torch.cuda.empty_cache() +def _module_weight_has_quant_info(module: torch.nn.Module) -> bool: + """Return True if ``module.weight`` carries a ``quant_info`` attribute. - pbar.update(1) + Used by the layerwise discovery in :func:`run_layerwise_quantization` + to locate ``nn.Linear`` modules selected for calibrated quantization + without depending on a module-level attribute. + """ + weight = getattr(module, "weight", None) + return weight is not None and hasattr(weight, "quant_info") - if include_lm_head: - hidden_states = run_layer(wrapper.get_pre_head_layernorm(return_name=False), hidden_states, return_output=True) - lm_head = wrapper.get_lm_head(return_name=False) - pbar.set_postfix(module="lm_head", refresh=False) - handle = lm_head.register_forward_hook(input_hook) - run_layer(lm_head, hidden_states, return_output=True) - handle.remove() - process_module(lm_head, device) - pbar.update(1) - pbar.close() +def module_quant_info_param_names(module: torch.nn.Module) -> list[str]: + """Return the names of ``module``'s direct parameters carrying ``quant_info``. - if original_use_cache is not None: - wrapper.model.config.use_cache = original_use_cache + Unlike :func:`_module_weight_has_quant_info` this does not hardcode the ``weight`` + attribute name, so it also finds fused-3D MoE expert parameters (``gate_up_proj`` / + ``down_proj``), which :func:`prepare_model` selects via + :func:`~olive.common.quant.selection.iter_quant_targets` but which were previously + invisible to the layerwise discovery loop. + """ + return [pname for pname, param in module._parameters.items() if param is not None and hasattr(param, "quant_info")] - return device + +def _split_quantizable_modules( + layer: torch.nn.Module, +) -> tuple[list[torch.nn.Module], list[torch.nn.Module]]: + """Split a layer's quantization targets into ordinary and fused-MoE modules. + + Returns ``(dense_modules, moe_modules)``: + + * ``dense_modules`` own a ``weight`` parameter with ``quant_info`` (``nn.Linear`` / + ``nn.Embedding``) and are calibrated with an ordinary forward hook; + * ``moe_modules`` own other ``quant_info``-carrying parameters -- fused-3D MoE experts + weights (``gate_up_proj`` / ``down_proj``) -- whose per-expert activations are only + observable from inside the experts forward. + """ + dense_modules: list[torch.nn.Module] = [] + moe_modules: list[torch.nn.Module] = [] + for module in layer.modules(): + if _module_weight_has_quant_info(module): + dense_modules.append(module) + elif module_quant_info_param_names(module): + moe_modules.append(module) + return dense_modules, moe_modules + + +def _iter_quant_info_params(model: torch.nn.Module): + """Yield ``(module, pname, param, quant_info)`` for every selected parameter.""" + for sub_module in model.modules(): + for pname in list(sub_module._parameters): + param = sub_module._parameters.get(pname) + if param is None: + continue + info = getattr(param, "quant_info", None) + if info is None: + continue + yield sub_module, pname, param, info def finalize( @@ -542,52 +737,85 @@ def finalize( device: str, retie_word_embeddings: bool = False, ) -> HfModelHandler: - """Finalize quantization by replacing linear and embedding layers with their quantized counterparts. - - Args: - model: The HuggingFace model to finalize. - output_model_path: Path to save the finalized quantized model. - wrapper: ModelWrapper containing the model to finalize. - quant_config: Quantization configuration to use. - device: Device to perform quantization on. - retie_word_embeddings: Whether to retie word embeddings if they were originally tied and have compatible quantization. - - Returns: - HfModelHandler with the finalized quantized model. - + """Finalize quantization by installing ``QuantTensor`` parameters in place. + + Walks every ``nn.Parameter`` whose tensor has a ``quant_info`` + attribute (set by :func:`prepare_model`), builds a ``QuantTensor`` + from the float tensor plus computed qparams, and installs it via + :func:`install_quant_tensor_param` so that: + + * ``module.`` is an ``nn.Parameter(QuantTensor)`` whose + dispatch still drives the original eager forward; + * ``module._qweight`` / ``_scales`` / ``_qzeros`` are plain + buffers (aliasing the QuantTensor's inner tensors), so the model + saves cleanly via ``save_pretrained`` / safetensors with no + tensor subclass on disk. + + The same code path handles 2D linear/embedding weights and any + higher-rank fused parameter (e.g. 3D MoE experts) — quantization + is always along the last dim. """ - - def should_quantize(module: torch.nn.Module, _: str) -> bool: - return hasattr(module, "quant_info") - - def quantize_and_pack(module: torch.nn.Module, _: str) -> QuantLinear | QuantEmbedding: - module.to(device) - quant_cls = QuantEmbedding if isinstance(module, torch.nn.Embedding) else QuantLinear - return quant_cls.from_module( - module.to(device), - bits=module.quant_info.quantizer.bits, - symmetric=module.quant_info.quantizer.symmetric, - group_size=module.quant_info.quantizer.group_size, - scales=module.quant_info.scales, - zero_points=module.quant_info.zero_points, - ).to("cpu") # move the original module to CPU - - replace_matching_submodules( - wrapper.model, - should_quantize, - quantize_and_pack, - description="Quantizing and packing linear layers", - ) + # Group selected params by their owning module so the module is + # moved to ``device`` once even when it owns multiple quantized + # parameters (fused-MoE experts modules carry two 3D tensors). + by_module: dict[int, tuple[torch.nn.Module, list[tuple[str, torch.nn.Parameter, QuantInfo]]]] = {} + for sub_module, pname, param, info in _iter_quant_info_params(wrapper.model): + entry = by_module.setdefault(id(sub_module), (sub_module, [])) + entry[1].append((pname, param, info)) + + for sub_module, params in by_module.values(): + sub_module.to(device) + with torch.no_grad(): + built = [] + for pname, param, info in params: + quantizer = info.quantizer + qt = QuantTensor.from_float( + param.data.detach(), + bits=quantizer.bits, + symmetric=quantizer.symmetric, + group_size=quantizer.group_size, + scales=info.scales, + zero_points=info.zero_points, + ).to("cpu") + built.append((pname, qt)) + sub_module.to("cpu") + for pname, qt in built: + install_quant_tensor_param(sub_module, pname, qt) if retie_word_embeddings: tie_quant_word_embeddings(wrapper.model) quant_config.tie_word_embeddings = True + if getattr(quant_config, "moe", False): + logger.warning( + "MoE weights have been quantized as 3D tensor parameters. The resulting checkpoint is " + "save/load compatible via transformers but is not directly exportable to ONNX with the " + "Olive ONNX conversion pass — consume it via the ORT GenAI model_builder or Mobius." + ) + wrapper.model.quantization_method = quant_config.quant_method wrapper.model.config.quantization_config = quant_config - # save the quantized model - wrapper.model.save_pretrained(output_model_path) + # save the quantized model — state_dict hooks drop QuantTensor entries; + # only plain ``_qweight`` / ``_scales`` / ``_qzeros`` buffers + # are written to safetensors. + # + # Newer ``transformers`` versions (>=5.x) default ``save_pretrained`` to + # ``save_original_format=True``, which for MoE architectures (e.g. Mixtral) + # round-trips the on-disk state dict through a *legacy* per-expert + # ``nn.Linear``-shaped layout (splitting the fused-3D ``experts.gate_up_proj``/ + # ``down_proj`` into ``experts.{i}.w1/w2/w3.weight`` and back). That + # reshape/(un)fuse machinery assumes plain float weight tensors and silently + # corrupts our quantized ``_scales``/``_qzeros`` buffers (their trailing + # group-size dimension gets dropped), which crashes real forward passes after + # a save/reload round-trip. Request the new, non-legacy on-disk format so the + # fused-3D quantized buffers are written byte-for-byte as-is. Older + # ``transformers`` versions don't accept this kwarg (they also don't have the + # legacy-format conversion machinery, so there's nothing to opt out of). + save_kwargs = {} + if "save_original_format" in inspect.signature(wrapper.model.save_pretrained).parameters: + save_kwargs["save_original_format"] = False + wrapper.model.save_pretrained(output_model_path, **save_kwargs) model.save_metadata(output_model_path) return inherit_hf_from_hf(model, output_model_path, adapter_path=model.adapter_path) diff --git a/olive/passes/pytorch/rtn.py b/olive/passes/pytorch/rtn.py index c7dbd462ee..fbbb826555 100644 --- a/olive/passes/pytorch/rtn.py +++ b/olive/passes/pytorch/rtn.py @@ -10,6 +10,7 @@ import torch from olive.passes import Pass +from olive.passes.pytorch.moe_support import check_moe_layout_support from olive.passes.pytorch.quant_utils import finalize, get_quantizer_config, prepare_model if TYPE_CHECKING: @@ -26,7 +27,7 @@ class Rtn(Pass): @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: - return get_quantizer_config(allow_embeds=True) + return get_quantizer_config(allow_embeds=True, allow_moe=True) @torch.no_grad() def _run_for_config( @@ -44,5 +45,20 @@ def _run_for_config( """ wrapper, qcfg, retie_word_embeddings = prepare_model(model, config, allow_quantized=True) + # Gate on this invocation's own ``moe`` request, not ``qcfg.moe`` -- ``prepare_model`` + # ORs in any pre-existing checkpoint's ``moe`` flag (see ``quant_utils.prepare_model``), + # so ``qcfg.moe`` can be True even when this run itself passed ``moe=False``. + if getattr(config, "moe", False): + experts_modules = [ + experts + for layer in wrapper.get_layer_wrappers() + if (experts := layer.get_experts(return_name=False)) is not None + ] + if experts_modules: + check_moe_layout_support( + experts_modules, + model_type=wrapper.model_type, + operation="RTN MoE quantization", + ) device = "cuda" if torch.cuda.is_available() else "cpu" return finalize(model, output_model_path, wrapper, qcfg, device, retie_word_embeddings=retie_word_embeddings) diff --git a/olive/passes/pytorch/save_test_model_config.py b/olive/passes/pytorch/save_test_model_config.py index f22d78388f..0b6edaa49b 100644 --- a/olive/passes/pytorch/save_test_model_config.py +++ b/olive/passes/pytorch/save_test_model_config.py @@ -14,7 +14,7 @@ class SaveTestModelConfig(Pass): - """Saves a random-initialised HuggingFace model to the test_model_path directory. + """Saves a random-initialized HuggingFace model to the test_model_path directory. When ``test_model_path`` and ``test_model_config`` are set on the input ``HfModelHandler``, this pass creates the target directory, writes @@ -69,7 +69,7 @@ def _run_for_config( ) if not _has_weights: logger.info("Saving test random model to %s", test_model_path) - # load_model calls load_model_from_task which creates a random-initialised model + # load_model calls load_model_from_task which creates a random-initialized model # from the reduced config and persists it (weights + config.json + marker) to # test_model_path on the first call. model.load_model(cache_model=False) diff --git a/olive/passes/pytorch/train_utils.py b/olive/passes/pytorch/train_utils.py index 7900b827e4..158c3c5c4f 100644 --- a/olive/passes/pytorch/train_utils.py +++ b/olive/passes/pytorch/train_utils.py @@ -255,7 +255,7 @@ def data_generator(dataset): def get_calibration_dataset( model: HfModelHandler | PyTorchModelHandler, data_config: DataConfig | dict | None = None, - split: str = "train[:1000]", + split: str = "train", batch_size: int = 1, max_seq_len: int = 2048, max_samples: int = 128, @@ -265,7 +265,7 @@ def get_calibration_dataset( Args: model: The HuggingFace or PyTorch model to get dataset for. data_config: Configuration object or dictionary containing data settings. - split: The dataset split to use for default data config. Default is 'train[:1000]'. + split: The dataset split to use for default data config. Default is 'train'. batch_size: The batch size to use for default data config. Default is 1. max_seq_len: Maximum sequence length for default data config. Default is 2048. max_samples: Maximum number of samples for default data config. Default is 128. @@ -312,7 +312,7 @@ def get_calibration_data_config( trust_remote_code: bool = False, data_name: str = "Salesforce/wikitext", subset: str = "wikitext-2-raw-v1", - split: str = "train[:1000]", + split: str = "train", batch_size: int = 1, max_seq_len: int = 2048, max_samples: int = 128, @@ -324,7 +324,7 @@ def get_calibration_data_config( trust_remote_code: Whether to trust remote code when loading data. data_name: The name of the dataset to use from Hugging Face Datasets. Default is "Salesforce/wikitext". subset: The subset of the dataset to use. Default is "wikitext-2-raw-v1". - split: The dataset split to use. Default is 'train[:1000]'. + split: The dataset split to use. Default is 'train'. batch_size: The batch size to use. Default is 1. max_seq_len: Maximum sequence length. Default is 2048. max_samples: Maximum number of samples. Default is 128. diff --git a/olive/systems/docker/docker_system.py b/olive/systems/docker/docker_system.py index 8371cccd44..1fff188c2a 100644 --- a/olive/systems/docker/docker_system.py +++ b/olive/systems/docker/docker_system.py @@ -232,6 +232,8 @@ def _prepare_run_params(self) -> dict: def _prepare_environment(self, base_env) -> dict: """Prepare environment variables for container.""" + from olive.telemetry.telemetry import Telemetry, is_ci_environment + # Convert list to dict if needed if isinstance(base_env, list): environment = {env.split("=")[0]: env.split("=")[1] for env in base_env} @@ -241,6 +243,10 @@ def _prepare_environment(self, base_env) -> dict: # Add default environment variables environment.setdefault("PYTHONPYCACHEPREFIX", "/tmp") environment["OLIVE_LOG_LEVEL"] = logging.getLevelName(logger.getEffectiveLevel()) + if is_ci_environment(): + environment["CI"] = "1" + if Telemetry.is_process_telemetry_disabled(): + environment["OLIVE_DISABLE_TELEMETRY"] = "1" # Add HuggingFace token if needed if self.hf_token: diff --git a/olive/systems/docker/workflow_runner.py b/olive/systems/docker/workflow_runner.py index 5842d0bd49..7035bb3ae3 100644 --- a/olive/systems/docker/workflow_runner.py +++ b/olive/systems/docker/workflow_runner.py @@ -5,6 +5,7 @@ from olive.common.hf.login import huggingface_login from olive.logging import get_olive_logger, set_verbosity_from_env +from olive.telemetry.telemetry import Telemetry from olive.workflows import run as olive_run logger = get_olive_logger() @@ -20,7 +21,16 @@ def runner_entry(config): config = json.load(f) logger.info("Running workflow with config: %s", config) - olive_run(config) + try: + olive_run(config, emit_error_telemetry=False, emit_recipe_telemetry=False) + finally: + telemetry = Telemetry.get_existing_instance() + if telemetry is not None: + telemetry.shutdown( + timeout_millis=15_000, + callback_timeout_millis=15_000, + flush_seconds=15, + ) if __name__ == "__main__": diff --git a/olive/telemetry/__init__.py b/olive/telemetry/__init__.py index 0ecbbc7056..f022f6177d 100644 --- a/olive/telemetry/__init__.py +++ b/olive/telemetry/__init__.py @@ -2,7 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -from olive.telemetry.telemetry import Telemetry +from olive.telemetry.telemetry import Telemetry, disable_telemetry from olive.telemetry.telemetry_extensions import action -__all__ = ["Telemetry", "action"] +__all__ = ["Telemetry", "action", "disable_telemetry"] diff --git a/olive/telemetry/constants.py b/olive/telemetry/constants.py deleted file mode 100644 index ca9e150b1b..0000000000 --- a/olive/telemetry/constants.py +++ /dev/null @@ -1,8 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""OneCollector connection string.""" - -CONNECTION_STRING = "SW5zdHJ1bWVudGF0aW9uS2V5PTlkNWRkYWVjNjFlMjQ1NjdiNzg4YTIwYWVhMzI0NjMxLTcyMzdkN2M2LWVlNjEtNGNmZC1iYjdiLTU5MDNhOTcyYzJlNC03MDQ3" diff --git a/olive/telemetry/deviceid/__init__.py b/olive/telemetry/deviceid/__init__.py index 50698c12d6..b467e8973b 100644 --- a/olive/telemetry/deviceid/__init__.py +++ b/olive/telemetry/deviceid/__init__.py @@ -1,3 +1,3 @@ -from olive.telemetry.deviceid.deviceid import get_encrypted_device_id_and_status +from olive.telemetry.deviceid.deviceid import get_hashed_device_id_and_status -__all__ = ["get_encrypted_device_id_and_status"] +__all__ = ["get_hashed_device_id_and_status"] diff --git a/olive/telemetry/deviceid/_store.py b/olive/telemetry/deviceid/_store.py index 97846ed051..dd62644417 100644 --- a/olive/telemetry/deviceid/_store.py +++ b/olive/telemetry/deviceid/_store.py @@ -1,9 +1,22 @@ +import os +import stat +import tempfile +from contextlib import suppress from pathlib import Path from olive.telemetry.utils import get_telemetry_base_dir REGISTRY_PATH = r"SOFTWARE\Microsoft\DeveloperTools\.onnxruntime" REGISTRY_KEY = "deviceid" +MAX_DEVICE_ID_FILE_SIZE = 256 + + +def _chmod_best_effort(path: Path, mode: int) -> None: + try: + path.chmod(mode) + except OSError: + # Permission tightening is best-effort on filesystems that do not support chmod. + pass class Store: @@ -21,23 +34,64 @@ def retrieve_id(self) -> str: :return: The device id. :rtype: str """ - # check if file doesnt exist and raise an Exception - if not self._file_path.is_file(): - raise FileExistsError(f"File {self._file_path.stem} does not exist") - - return self._file_path.read_text(encoding="utf-8").strip() - - def store_id(self, device_id: str) -> None: + flags = os.O_RDONLY + for optional_flag in ("O_CLOEXEC", "O_NOFOLLOW", "O_NONBLOCK"): + flags |= getattr(os, optional_flag, 0) + try: + fd = os.open(self._file_path, flags) + except FileNotFoundError: + raise FileNotFoundError(f"File {self._file_path.stem} does not exist") from None + try: + if not stat.S_ISREG(os.fstat(fd).st_mode): + raise PermissionError(f"File {self._file_path.stem} is not a regular file") + chunks = [] + remaining = MAX_DEVICE_ID_FILE_SIZE + 1 + while remaining: + chunk = os.read(fd, remaining) + if not chunk: + break + chunks.append(chunk) + remaining -= len(chunk) + content = b"".join(chunks) + if len(content) > MAX_DEVICE_ID_FILE_SIZE: + raise ValueError(f"File {self._file_path.stem} is too large") + return content.decode("utf-8").strip() + except UnicodeDecodeError: + raise ValueError(f"File {self._file_path.stem} is not valid UTF-8") from None + finally: + os.close(fd) + + def store_id(self, device_id: str, replace_existing: bool = False) -> bool: """Store the device id in the store location. :param str device_id: The device id to store. :type device_id: str """ - # create the folder location if it does not exist - self._file_path.parent.mkdir(parents=True, exist_ok=True) - - self._file_path.touch() - self._file_path.write_text(device_id, encoding="utf-8") + # create the folder location if it does not exist, owner-only (0700) so other users on the + # machine cannot traverse into it to reach the device id. + self._file_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + _chmod_best_effort(self._file_path.parent, 0o700) + + fd, temp_path = tempfile.mkstemp(prefix="deviceid.tmp.", dir=self._file_path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as temp_file: + temp_file.write(device_id) + temp_file.flush() + os.fsync(temp_file.fileno()) + _chmod_best_effort(Path(temp_path), 0o600) + if replace_existing: + Path(temp_path).replace(self._file_path) + temp_path = "" + return True + try: + os.link(temp_path, self._file_path) + except FileExistsError: + return False + return True + finally: + if temp_path: + with suppress(OSError): + Path(temp_path).unlink() class WindowsStore: @@ -46,15 +100,15 @@ def retrieve_id(self) -> str: """Retrieve the device id from the Windows registry.""" import winreg - device_id: str - with winreg.OpenKeyEx( winreg.HKEY_CURRENT_USER, REGISTRY_PATH, reserved=0, access=winreg.KEY_READ | winreg.KEY_WOW64_64KEY ) as key_handle: - device_id = winreg.QueryValueEx(key_handle, REGISTRY_KEY) - return device_id[0].strip() + device_id, value_type = winreg.QueryValueEx(key_handle, REGISTRY_KEY) + if value_type != winreg.REG_SZ or not isinstance(device_id, str): + raise ValueError(f"Registry value {REGISTRY_KEY} is not a string") + return device_id.strip() - def store_id(self, device_id: str) -> None: + def store_id(self, device_id: str, replace_existing: bool = False) -> bool: """Store the device id in the windows registry. :param str device_id: The device id to store. @@ -65,6 +119,13 @@ def store_id(self, device_id: str) -> None: winreg.HKEY_CURRENT_USER, REGISTRY_PATH, reserved=0, - access=winreg.KEY_ALL_ACCESS | winreg.KEY_WOW64_64KEY, + access=winreg.KEY_QUERY_VALUE | winreg.KEY_SET_VALUE | winreg.KEY_CREATE_SUB_KEY | winreg.KEY_WOW64_64KEY, ) as key_handle: + if not replace_existing: + try: + winreg.QueryValueEx(key_handle, REGISTRY_KEY) + return False + except FileNotFoundError: + pass winreg.SetValueEx(key_handle, REGISTRY_KEY, 0, winreg.REG_SZ, device_id) + return True diff --git a/olive/telemetry/deviceid/deviceid.py b/olive/telemetry/deviceid/deviceid.py index 09087f33c3..d1430f5e99 100644 --- a/olive/telemetry/deviceid/deviceid.py +++ b/olive/telemetry/deviceid/deviceid.py @@ -1,23 +1,138 @@ import hashlib import platform +import threading import uuid +from contextlib import suppress from enum import Enum -from typing import Union +from typing import ClassVar from olive.telemetry.deviceid._store import Store, WindowsStore +from olive.telemetry.process_lock import ProcessDrainLock +from olive.telemetry.utils import get_telemetry_base_dir class DeviceIdStatus(Enum): - NEW = "new" - EXISTING = "existing" - CORRUPTED = "corrupted" - FAILED = "failed" + NEW = "New" + EXISTING = "Existing" + CORRUPTED = "Corrupted" + FAILED = "Failed" _device_id_state = {"device_id": None, "status": DeviceIdStatus.NEW} +_device_id_lock = threading.RLock() -def get_device_id() -> str: +def _fnv1a_hex_bytes(value: bytes) -> str: + """Hash the Windows SID only for the native-compatible mutex name.""" + hash_value = 14695981039346656037 + for byte in value: + hash_value ^= byte + hash_value = (hash_value * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return f"{hash_value:016x}" + + +class _WindowsDeviceIdMutex: + """Named mutex compatible with the native device-id publication protocol.""" + + def __init__(self) -> None: + self._handle = None + self._acquired = False + self._kernel32 = None + + def acquire(self) -> bool: + try: + import ctypes + from ctypes import wintypes + + class SidAndAttributes(ctypes.Structure): + _fields_: ClassVar = [("sid", ctypes.c_void_p), ("attributes", wintypes.DWORD)] + + class TokenUser(ctypes.Structure): + _fields_: ClassVar = [("user", SidAndAttributes)] + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + advapi32 = ctypes.WinDLL("advapi32", use_last_error=True) + kernel32.GetCurrentProcess.restype = wintypes.HANDLE + kernel32.CreateMutexW.argtypes = [ctypes.c_void_p, wintypes.BOOL, wintypes.LPCWSTR] + kernel32.CreateMutexW.restype = wintypes.HANDLE + kernel32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD] + kernel32.WaitForSingleObject.restype = wintypes.DWORD + kernel32.ReleaseMutex.argtypes = [wintypes.HANDLE] + kernel32.CloseHandle.argtypes = [wintypes.HANDLE] + advapi32.OpenProcessToken.argtypes = [ + wintypes.HANDLE, + wintypes.DWORD, + ctypes.POINTER(wintypes.HANDLE), + ] + advapi32.GetTokenInformation.argtypes = [ + wintypes.HANDLE, + ctypes.c_int, + ctypes.c_void_p, + wintypes.DWORD, + ctypes.POINTER(wintypes.DWORD), + ] + advapi32.IsValidSid.argtypes = [ctypes.c_void_p] + advapi32.GetLengthSid.argtypes = [ctypes.c_void_p] + advapi32.GetLengthSid.restype = wintypes.DWORD + + token = wintypes.HANDLE() + if not advapi32.OpenProcessToken(kernel32.GetCurrentProcess(), 0x0008, ctypes.byref(token)): + return False + try: + size = wintypes.DWORD() + advapi32.GetTokenInformation(token, 1, None, 0, ctypes.byref(size)) + if not size.value: + return False + token_info = ctypes.create_string_buffer(size.value) + if not advapi32.GetTokenInformation(token, 1, token_info, size.value, ctypes.byref(size)): + return False + sid = ctypes.cast(token_info, ctypes.POINTER(TokenUser)).contents.user.sid + if not sid or not advapi32.IsValidSid(sid): + return False + sid_size = advapi32.GetLengthSid(sid) + sid_hash = _fnv1a_hex_bytes(ctypes.string_at(sid, sid_size)) + finally: + kernel32.CloseHandle(token) + + handle = kernel32.CreateMutexW( + None, + False, + f"Global\\Microsoft.DeveloperTools.OnnxRuntime.DeviceId.{sid_hash}", + ) + if not handle: + return False + self._handle = handle + self._kernel32 = kernel32 + wait_result = kernel32.WaitForSingleObject(handle, 1000) + self._acquired = wait_result in (0x00000000, 0x00000080) + return self._acquired + except Exception: + self.release() + return False + + def release(self) -> None: + if self._handle is None or self._kernel32 is None: + return + if self._acquired: + with suppress(Exception): + self._kernel32.ReleaseMutex(self._handle) + with suppress(Exception): + self._kernel32.CloseHandle(self._handle) + self._handle = None + self._acquired = False + self._kernel32 = None + + +def _is_valid_device_id(value: str) -> bool: + if not isinstance(value, str) or len(value) != 36: + return False + hyphens = {8, 13, 18, 23} + return all( + char == "-" if index in hyphens else char.lower() in "0123456789abcdef" for index, char in enumerate(value) + ) + + +def _initialize_device_id() -> str: r"""Get the device id from the store or create one if it does not exist. An empty string is returned if an error occurs during saving or retrieval of the device id. @@ -25,77 +140,105 @@ def get_device_id() -> str: Linux id location: $XDG_CACHE_HOME/Microsoft/DeveloperTools/.onnxruntime/deviceid if defined else $HOME/.cache/Microsoft/DeveloperTools/.onnxruntime/deviceid MacOS id location: $HOME/Library/Application Support/Microsoft/DeveloperTools/.onnxruntime/deviceid - Windows id location: HKEY_CURRENT_USER\SOFTWARE\Microsoft\.onnxruntime\deviceid + Windows id location: HKEY_CURRENT_USER\SOFTWARE\Microsoft\DeveloperTools\.onnxruntime\deviceid :return: The device id. :rtype: str """ - device_id: str = "" - store: Union[Store, WindowsStore] - create_new_id = False - - try: - if platform.system() == "Windows": - store = WindowsStore() - elif platform.system() in ("Linux", "Darwin"): + system = platform.system() + if system == "Windows": + store = WindowsStore() + elif system in ("Linux", "Darwin"): + try: store = Store() - else: - _device_id_state["status"] = DeviceIdStatus.FAILED - _device_id_state["device_id"] = device_id - return device_id - - device_id = store.retrieve_id - if len(device_id) > 256: - _device_id_state["status"] = DeviceIdStatus.CORRUPTED - _device_id_state["device_id"] = "" - create_new_id = True - else: - try: - uuid.UUID(device_id) - except ValueError: - _device_id_state["status"] = DeviceIdStatus.CORRUPTED - _device_id_state["device_id"] = "" - create_new_id = True - else: - _device_id_state["status"] = DeviceIdStatus.EXISTING - _device_id_state["device_id"] = device_id - return device_id - except (FileExistsError, FileNotFoundError): - _device_id_state["status"] = DeviceIdStatus.NEW - _device_id_state["device_id"] = "" - create_new_id = True - except (PermissionError, ValueError, NotImplementedError): - _device_id_state["status"] = DeviceIdStatus.FAILED - _device_id_state["device_id"] = device_id - return device_id - except Exception: - _device_id_state["status"] = DeviceIdStatus.FAILED - _device_id_state["device_id"] = device_id - return device_id - - if create_new_id: - device_id = str(uuid.uuid4()).lower() - + except Exception: + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + else: + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + + def read_existing() -> tuple[str, str]: try: - store.store_id(device_id) + existing = store.retrieve_id + except (FileExistsError, FileNotFoundError): + return ("missing", "") + except ValueError: + return ("invalid", "") except Exception: - _device_id_state["status"] = DeviceIdStatus.FAILED - device_id = "" - _device_id_state["device_id"] = device_id - - return device_id - - -def get_encrypted_device_id_and_status() -> tuple[str, DeviceIdStatus]: - """Generate a FIPS-compliant encrypted device ID using SHA256 and returns the deviceIdStatus. + return ("failed", "") + return ("valid", existing) if _is_valid_device_id(existing) else ("invalid", "") + + initial_state, existing = read_existing() + if initial_state == "valid": + _device_id_state.update({"status": DeviceIdStatus.EXISTING, "device_id": existing}) + return existing + if initial_state == "failed": + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + + lock = None + acquired = True + if system == "Windows": + lock = _WindowsDeviceIdMutex() + acquired = lock.acquire() + elif initial_state == "invalid": + lock = ProcessDrainLock(str(get_telemetry_base_dir() / "deviceid.lock")) + acquired = lock.acquire(1.0) - This method uses SHA256 which is FIPS 140-2 approved for cryptographic operations. - The device ID is hashed to ensure deterministic but secure device identification. + try: + if not acquired: + winner_state, winner = read_existing() + generated = winner if winner_state == "valid" else str(uuid.uuid4()).lower() + status = DeviceIdStatus.EXISTING if winner_state == "valid" else DeviceIdStatus.FAILED + _device_id_state.update({"status": status, "device_id": generated}) + return generated + + current_state, current = read_existing() + if current_state == "valid": + _device_id_state.update({"status": DeviceIdStatus.EXISTING, "device_id": current}) + return current + if current_state == "failed": + generated = str(uuid.uuid4()).lower() + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + + corrupted = initial_state == "invalid" or current_state == "invalid" + generated = str(uuid.uuid4()).lower() + try: + stored = store.store_id(generated, replace_existing=corrupted) + except Exception: + stored = False + if stored: + status = DeviceIdStatus.CORRUPTED if corrupted else DeviceIdStatus.NEW + _device_id_state.update({"status": status, "device_id": generated}) + return generated + + winner_state, winner = read_existing() + if winner_state == "valid": + _device_id_state.update({"status": DeviceIdStatus.EXISTING, "device_id": winner}) + return winner + _device_id_state.update({"status": DeviceIdStatus.FAILED, "device_id": generated}) + return generated + finally: + if lock is not None: + lock.release() - Returns: - str: FIPS-compliant encrypted device ID (base64-encoded) - """ - device_id = _device_id_state["device_id"] if _device_id_state["device_id"] is not None else get_device_id() - encrypted_device_id = hashlib.sha256(device_id.encode("utf-8")).digest().hex().upper() if device_id else "" - return encrypted_device_id, _device_id_state["status"] +def get_device_id() -> str: + """Get the process-cached persistent device ID, initializing it once.""" + with _device_id_lock: + if _device_id_state["device_id"] is None: + return _initialize_device_id() + return _device_id_state["device_id"] + + +def get_hashed_device_id_and_status() -> tuple[str, DeviceIdStatus]: + """Get the canonical shared SHA-256 device ID and its status.""" + with _device_id_lock: + device_id = get_device_id() + hashed = hashlib.sha256(device_id.encode("utf-8")).hexdigest() if device_id else "" + return f"c:{hashed}" if hashed else "", _device_id_state["status"] diff --git a/olive/telemetry/library/__init__.py b/olive/telemetry/library/__init__.py index 39831da66e..dcf07e101d 100644 --- a/olive/telemetry/library/__init__.py +++ b/olive/telemetry/library/__init__.py @@ -3,44 +3,14 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""OneCollector Exporter for OpenTelemetry Python. +"""OneCollector building blocks (standard library only). -This package provides an OpenTelemetry exporter that sends telemetry data -to Microsoft OneCollector using the Common Schema JSON format. - -Example usage: - - from onecollector_exporter import ( - OneCollectorLogExporter, - OneCollectorExporterOptions, - get_telemetry_logger, - ) - - # Option 1: Use with OpenTelemetry SDK directly - options = OneCollectorExporterOptions( - connection_string="InstrumentationKey=your-key-here" - ) - exporter = OneCollectorLogExporter(options=options) - - # Add to logger provider - from opentelemetry.sdk._logs import LoggerProvider - from opentelemetry.sdk._logs.export import BatchLogRecordProcessor - - provider = LoggerProvider() - provider.add_log_record_processor(BatchLogRecordProcessor(exporter)) - - # Option 2: Use the simplified telemetry logger - logger = get_telemetry_logger( - connection_string="InstrumentationKey=your-key-here" - ) - logger.log("MyEvent", {"key": "value"}) - logger.shutdown() +Helpers for serializing telemetry to Common Schema JSON and posting it to the +Microsoft OneCollector endpoint. These modules have no third-party dependency +and are driven directly by the SQLite-backed uploader. """ -from olive.telemetry.library.callback_manager import CallbackManager, PayloadTransmittedCallbackArgs from olive.telemetry.library.connection_string_parser import ConnectionStringParser -from olive.telemetry.library.event_source import OneCollectorEventId, OneCollectorEventSource, event_source -from olive.telemetry.library.exporter import OneCollectorLogExporter from olive.telemetry.library.options import ( CompressionType, OneCollectorExporterOptions, @@ -48,37 +18,16 @@ OneCollectorTransportOptions, ) from olive.telemetry.library.payload_builder import PayloadBuilder -from olive.telemetry.library.retry import RetryHandler from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper -from olive.telemetry.library.telemetry_logger import ( - TelemetryLogger, - get_telemetry_logger, - log_event, - shutdown_telemetry, -) -from olive.telemetry.library.transport import HttpJsonPostTransport, ITransport - -__version__ = "0.0.1" +from olive.telemetry.library.transport import HttpJsonPostTransport __all__ = [ - "CallbackManager", "CommonSchemaJsonSerializationHelper", "CompressionType", "ConnectionStringParser", "HttpJsonPostTransport", - "ITransport", - "OneCollectorEventId", - "OneCollectorEventSource", "OneCollectorExporterOptions", "OneCollectorExporterValidationError", - "OneCollectorLogExporter", "OneCollectorTransportOptions", "PayloadBuilder", - "PayloadTransmittedCallbackArgs", - "RetryHandler", - "TelemetryLogger", - "event_source", - "get_telemetry_logger", - "log_event", - "shutdown_telemetry", ] diff --git a/olive/telemetry/library/callback_manager.py b/olive/telemetry/library/callback_manager.py deleted file mode 100644 index ee62553163..0000000000 --- a/olive/telemetry/library/callback_manager.py +++ /dev/null @@ -1,110 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""Callback manager for payload transmission events.""" - -import threading -from dataclasses import dataclass -from typing import Callable, Optional - -from olive.telemetry.library.event_source import event_source - - -@dataclass -class PayloadTransmittedCallbackArgs: - """Arguments passed to payload transmitted callbacks.""" - - succeeded: bool - """Whether the transmission succeeded.""" - - status_code: Optional[int] - """HTTP status code, if available.""" - - payload_size_bytes: int - """Size of the transmitted payload in bytes.""" - - item_count: int - """Number of items in the payload.""" - - payload_bytes: Optional[bytes] = None - """Raw payload bytes (uncompressed), if available.""" - - -class CallbackManager: - """Manages callbacks for payload transmission events. - - Allows registration of callbacks that are invoked when payloads - are successfully transmitted or fail. - """ - - def __init__(self): - """Initialize the callback manager.""" - self._callbacks: list[tuple[Callable[[PayloadTransmittedCallbackArgs], None], bool]] = [] - self._lock = threading.Lock() - self._closed = False - - def register( - self, callback: Callable[[PayloadTransmittedCallbackArgs], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback to be invoked on payload transmission. - - Args: - callback: Function to call when payload is transmitted - include_failures: Whether to invoke callback on transmission failures - - Returns: - Function to call to unregister the callback - - """ - with self._lock: - if self._closed: - return lambda: None # No-op unregister if disposed - entry = (callback, include_failures) - self._callbacks.append(entry) - - def unregister(): - """Unregister this callback.""" - with self._lock: - try: - self._callbacks.remove(entry) - except ValueError: - # The callback was already removed. - pass - - return unregister - - def notify(self, args: PayloadTransmittedCallbackArgs) -> None: - """Notify all registered callbacks. - - Args: - args: Callback arguments - - """ - # Get snapshot of callbacks to avoid holding lock during invocation - with self._lock: - if self._closed: - return - callbacks_snapshot = self._callbacks.copy() - - # Invoke callbacks - for callback, include_failures in callbacks_snapshot: - # Check if we should invoke this callback - if not args.succeeded and not include_failures: - continue - - try: - callback(args) - except Exception as ex: - # Log but don't propagate exceptions from user code - event_source.exception_thrown_from_user_code("PayloadTransmittedCallback", ex) - - def close(self) -> None: - """Close the callback manager and prevent further registrations. - - This method is idempotent and can be called multiple times. - """ - with self._lock: - self._callbacks.clear() - self._closed = True diff --git a/olive/telemetry/library/event_source.py b/olive/telemetry/library/event_source.py deleted file mode 100644 index e65d9d546a..0000000000 --- a/olive/telemetry/library/event_source.py +++ /dev/null @@ -1,257 +0,0 @@ -"""EventSource-style logging for OneCollector exporter. - -Provides structured logging similar to .NET EventSource for diagnostics and monitoring. -""" - -import logging -from enum import IntEnum - - -class OneCollectorEventId(IntEnum): - """Event IDs matching .NET EventSource implementation.""" - - EXPORT_EXCEPTION = 1 - TRANSPORT_DATA_SENT = 2 - SINK_DATA_WRITTEN = 3 - DATA_DROPPED = 4 - TRANSPORT_EXCEPTION = 5 - HTTP_ERROR_RESPONSE = 6 - EVENT_FULL_NAME_DISCARDED = 7 - EVENT_NAMESPACE_INVALID = 8 - EVENT_NAME_INVALID = 9 - USER_CODE_EXCEPTION = 10 - ATTRIBUTE_DROPPED = 11 - - -class OneCollectorEventSource: - """EventSource for OneCollector exporter diagnostics. - - Provides structured logging matching the .NET EventSource implementation. - """ - - def __init__(self): - self.logger = logging.getLogger("OpenTelemetry.Exporter.OneCollector") - # Set default level to INFO to match .NET behavior - if not self.logger.handlers: - handler = logging.StreamHandler() - formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") - handler.setFormatter(formatter) - self.logger.addHandler(handler) - self.logger.setLevel(logging.INFO) - - @property - def is_informational_logging_enabled(self) -> bool: - """Check if informational level logging is enabled.""" - return self.logger.isEnabledFor(logging.INFO) - - @property - def is_warning_logging_enabled(self) -> bool: - """Check if warning level logging is enabled.""" - return self.logger.isEnabledFor(logging.WARNING) - - @property - def is_error_logging_enabled(self) -> bool: - """Check if error level logging is enabled.""" - return self.logger.isEnabledFor(logging.ERROR) - - def export_exception_thrown(self, item_type: str, exception: Exception) -> None: - """Log an exception thrown during export. - - Args: - item_type: Type of item being exported (e.g., 'LogData') - exception: The exception that was thrown - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Exception thrown exporting '%s' batch: %s", - item_type, - exception, - exc_info=exception, - extra={"event_id": OneCollectorEventId.EXPORT_EXCEPTION}, - ) - - def transport_data_sent(self, item_type: str, num_records: int, transport_description: str) -> None: - """Log successful data transmission. - - Args: - item_type: Type of items sent - num_records: Number of records sent - transport_description: Description of transport used - - """ - if self.is_informational_logging_enabled: - self.logger.info( - "Sent '%s' batch of %s item(s) to '%s' transport", - item_type, - num_records, - transport_description, - extra={"event_id": OneCollectorEventId.TRANSPORT_DATA_SENT}, - ) - - def sink_data_written(self, item_type: str, num_records: int, sink_description: str) -> None: - """Log data written to sink. - - Args: - item_type: Type of items written - num_records: Number of records written - sink_description: Description of sink used - - """ - if self.is_informational_logging_enabled: - self.logger.info( - "Wrote '%s' batch of %s item(s) to '%s' sink", - item_type, - num_records, - sink_description, - extra={"event_id": OneCollectorEventId.SINK_DATA_WRITTEN}, - ) - - def data_dropped( - self, item_type: str, num_records: int, during_serialization: int, during_transmission: int - ) -> None: - """Log dropped data. - - Args: - item_type: Type of items dropped - num_records: Total number of records dropped - during_serialization: Number dropped during serialization - during_transmission: Number dropped during transmission - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Dropped %s '%s' item(s). %s item(s) dropped during serialization. %s item(s) dropped due to " - "transmission failure", - num_records, - item_type, - during_serialization, - during_transmission, - extra={"event_id": OneCollectorEventId.DATA_DROPPED}, - ) - - def transport_exception_thrown(self, transport_type: str, exception: Exception) -> None: - """Log transport exception. - - Args: - transport_type: Type of transport - exception: The exception that was thrown - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Exception thrown by '%s' transport: %s", - transport_type, - exception, - exc_info=exception, - extra={"event_id": OneCollectorEventId.TRANSPORT_EXCEPTION}, - ) - - def http_transport_error_response( - self, transport_type: str, status_code: int, error_message: str, error_details: str - ) -> None: - """Log HTTP error response. - - Args: - transport_type: Type of transport - status_code: HTTP status code - error_message: Error message from response - error_details: Additional error details - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Error response received by '%s' transport. StatusCode: %s, ErrorMessage: '%s', ErrorDetails: '%s'", - transport_type, - status_code, - error_message, - error_details, - extra={"event_id": OneCollectorEventId.HTTP_ERROR_RESPONSE}, - ) - - def event_full_name_discarded(self, event_namespace: str, event_name: str) -> None: - """Log event full name discarded. - - Args: - event_namespace: Event namespace - event_name: Event name - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Event full name discarded. EventNamespace: '%s', EventName: '%s'", - event_namespace, - event_name, - extra={"event_id": OneCollectorEventId.EVENT_FULL_NAME_DISCARDED}, - ) - - def event_namespace_invalid(self, event_namespace: str) -> None: - """Log invalid event namespace. - - Args: - event_namespace: The invalid namespace - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Event namespace invalid. EventNamespace: '%s'", - event_namespace, - extra={"event_id": OneCollectorEventId.EVENT_NAMESPACE_INVALID}, - ) - - def event_name_invalid(self, event_name: str) -> None: - """Log invalid event name. - - Args: - event_name: The invalid event name - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Event name invalid. EventName: '%s'", - event_name, - extra={"event_id": OneCollectorEventId.EVENT_NAME_INVALID}, - ) - - def exception_thrown_from_user_code(self, user_code_type: str, exception: Exception) -> None: - """Log exception from user code (e.g., callbacks). - - Args: - user_code_type: Type of user code that threw exception - exception: The exception that was thrown - - """ - if self.is_error_logging_enabled: - self.logger.error( - "Exception thrown by '%s' user code: %s", - user_code_type, - exception, - exc_info=exception, - extra={"event_id": OneCollectorEventId.USER_CODE_EXCEPTION}, - ) - - def attribute_dropped(self, item_type: str, attribute_name: str, reason: str) -> None: - """Log dropped attribute. - - Args: - item_type: Type of item - attribute_name: Name of dropped attribute - reason: Reason for dropping - - """ - if self.is_warning_logging_enabled: - self.logger.warning( - "Dropped %s attribute '%s': %s", - item_type, - attribute_name, - reason, - extra={"event_id": OneCollectorEventId.ATTRIBUTE_DROPPED}, - ) - - def disable(self) -> None: - """Disable telemetry logging.""" - self.logger.disabled = True - - -# Global event source instance -event_source = OneCollectorEventSource() diff --git a/olive/telemetry/library/exporter.py b/olive/telemetry/library/exporter.py deleted file mode 100644 index 045647fe9d..0000000000 --- a/olive/telemetry/library/exporter.py +++ /dev/null @@ -1,335 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""Main OneCollector log exporter implementation.""" - -import threading -from collections.abc import Sequence -from datetime import datetime, timezone -from time import time -from typing import TYPE_CHECKING, Any, Callable, Optional - -import requests -from opentelemetry.sdk._logs import ReadableLogRecord -from opentelemetry.sdk._logs.export import LogExportResult, LogRecordExporter -from opentelemetry.sdk.resources import Resource - -from olive.telemetry.library.callback_manager import CallbackManager -from olive.telemetry.library.event_source import event_source -from olive.telemetry.library.options import OneCollectorExporterOptions -from olive.telemetry.library.payload_builder import PayloadBuilder -from olive.telemetry.library.retry import RetryHandler -from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper -from olive.telemetry.library.transport import HttpJsonPostTransport - -if TYPE_CHECKING: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - - -class OneCollectorLogExporter(LogRecordExporter): - """OpenTelemetry log exporter for Microsoft OneCollector. - - Implements the OpenTelemetry LogRecordExporter interface and sends logs - to OneCollector using the Common Schema JSON format. - """ - - def __init__( - self, - options: Optional[OneCollectorExporterOptions] = None, - excluded_attributes: Optional[set[str]] = None, - ): - """Initialize the OneCollector log exporter. - - Args: - options: Exporter configuration options - excluded_attributes: Attribute keys to exclude from log attributes - - """ - # Validate options - if options is None: - raise ValueError("OneCollectorExporterOptions is required") - options.validate() - - self._options = options - self._shutdown_lock = threading.Lock() - self._shutdown = False - self._shutdown_event = threading.Event() - if excluded_attributes is None: - self._excluded_attributes = { - "code.filepath", - "code.function", - "code.lineno", - "code.file.path", - "code.function.name", - "code.line.number", - } - else: - self._excluded_attributes = set(excluded_attributes) - - # Initialize transport - transport_opts = options.transport_options - - # Create or get HTTP session - if transport_opts.http_client_factory: - self._session = transport_opts.http_client_factory() - self._owns_session = False - else: - self._session = requests.Session() - self._owns_session = True - - try: - # Build iKey with tenant prefix - self._ikey = f"{CommonSchemaJsonSerializationHelper.ONE_COLLECTOR_TENANCY_SYMBOL}:{options.tenant_token}" - - # Initialize callback manager - self._callback_manager = CallbackManager() - - # Initialize transport with callback manager - self._transport = HttpJsonPostTransport( - endpoint=transport_opts.endpoint, - ikey=options.instrumentation_key, - compression=transport_opts.compression, - session=self._session, - callback_manager=self._callback_manager, - ) - - # Initialize payload builder - self._payload_builder = PayloadBuilder( - max_size_bytes=transport_opts.max_payload_size_bytes, max_items=transport_opts.max_items_per_payload - ) - - # Initialize retry handler - self._retry_handler = RetryHandler(max_retries=6) - - # Initialize metadata - self._metadata: dict[str, Any] = {} - - # Cache for resource (populated on first export) - self._resource: Optional[Resource] = None - except Exception: - if self._owns_session: - self._session.close() - raise - - def add_metadata(self, metadata: dict[str, Any]) -> None: - """Add custom metadata fields to all exported logs. - - Args: - metadata: Dictionary of metadata fields to add - - """ - self._metadata.update(metadata) - - def register_payload_transmitted_callback( - self, callback: Callable[["PayloadTransmittedCallbackArgs"], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback that will be invoked on payload transmission. - - Callbacks are invoked after each HTTP request completes. If retries are - enabled, callbacks will be invoked for each retry attempt. - - Args: - callback: Function to call when payload is transmitted. - Receives PayloadTransmittedCallbackArgs with transmission details. - include_failures: If True, callback is invoked on both success and failure. - If False, callback is only invoked on success. - - Returns: - Function to call to unregister the callback. - - Example: - >>> def on_transmitted(args): - ... if args.succeeded: - ... print(f"✅ Sent {args.item_count} items ({args.payload_size_bytes} bytes)") - ... else: - ... print(f"❌ Failed: status={args.status_code}") - >>> - >>> unregister = exporter.register_payload_transmitted_callback( - ... on_transmitted, - ... include_failures=True - ... ) - >>> # Later: unregister() - - """ - return self._transport.register_payload_transmitted_callback(callback, include_failures) - - def export(self, batch: Sequence[ReadableLogRecord]) -> LogExportResult: - """Export a batch of log records. - - Args: - batch: Sequence of log data records to export - - Returns: - LogExportResult indicating success or failure - - """ - if self._shutdown: - return LogExportResult.FAILURE - - try: - # Get resource (cache for subsequent calls) - if self._resource is None: - first_item = batch[0] if batch else None - resource = getattr(first_item, "resource", None) - if resource is None and first_item is not None: - resource = getattr(first_item.log_record, "resource", None) - self._resource = resource or Resource.create() - - # Serialize log records to JSON - serialized_items = [] - for log_data in batch: - try: - item_bytes = self._serialize_log_data(log_data) - serialized_items.append(item_bytes) - except Exception as ex: - event_source.export_exception_thrown("ReadableLogRecord", ex) - # Continue with other items - - if not serialized_items: - return LogExportResult.FAILURE - - # Build payloads respecting size/count limits - payloads = self._build_payloads(serialized_items) - - # Send each payload with retry logic - deadline_sec = time() + self._options.transport_options.timeout_seconds - - for payload in payloads: - # Count items in this payload (approximation based on newlines) - item_count = payload.count(b"\n") + 1 if payload else 0 - success = self._retry_handler.execute_with_retry( - operation=lambda payload=payload, item_count=item_count: self._transport.send( - payload, max(0.1, deadline_sec - time()), item_count=item_count - ), - deadline_sec=deadline_sec, - shutdown_event=self._shutdown_event, - ) - - if not success: - return LogExportResult.FAILURE - - # Check if shutdown occurred - if self._shutdown: - return LogExportResult.FAILURE - - # Log success - event_source.sink_data_written("ReadableLogRecord", len(batch), "OneCollector") - - return LogExportResult.SUCCESS - - except Exception as ex: - event_source.export_exception_thrown("ReadableLogRecord", ex) - return LogExportResult.FAILURE - - def _serialize_log_data(self, log_data: ReadableLogRecord) -> bytes: - """Serialize a single log record to JSON bytes. - - Args: - log_data: Log data to serialize - - Returns: - UTF-8 encoded JSON bytes - - """ - log_record = log_data.log_record - - # Build data dictionary - data = {} - - # Add resource attributes (if available) - if self._resource and self._resource.attributes: - for key, value in self._resource.attributes.items(): - # Map common resource attributes - if key == "service.name" and "app_name" not in data: - data["app_name"] = value - elif key == "service.version" and "app_version" not in data: - data["app_version"] = value - elif key == "service.instance.id" and "app_instance_id" not in data: - data["app_instance_id"] = value - else: - data[key] = value - - # Add log record attributes (override resource attributes) - if log_record.attributes: - data.update( - {key: value for key, value in log_record.attributes.items() if key not in self._excluded_attributes} - ) - - # Add custom metadata - data.update(self._metadata) - - # Format timestamp - if log_record.timestamp: - timestamp = datetime.fromtimestamp(log_record.timestamp / 1e9, tz=timezone.utc) - else: - timestamp = datetime.now(timezone.utc) - - # Create event envelope - event_name = str(log_record.body) if log_record.body else "UnnamedEvent" - - envelope = CommonSchemaJsonSerializationHelper.create_event_envelope( - event_name=event_name, timestamp=timestamp, ikey=self._ikey, data=data - ) - - # Serialize to JSON bytes - return CommonSchemaJsonSerializationHelper.serialize_to_json_bytes(envelope) - - def _build_payloads(self, serialized_items: list[bytes]) -> list[bytes]: - """Build payloads from serialized items respecting size and count limits. - - Args: - serialized_items: List of serialized item bytes - - Returns: - List of payload bytes - - """ - payloads = [] - self._payload_builder.reset() - - for item_bytes in serialized_items: - if not self._payload_builder.can_add(item_bytes) and not self._payload_builder.is_empty: - # Current payload is full, build it and start a new one - payloads.append(self._payload_builder.build()) - self._payload_builder.reset() - - self._payload_builder.add(item_bytes) - - # Build final payload - if not self._payload_builder.is_empty: - payloads.append(self._payload_builder.build()) - - return payloads - - def force_flush(self, timeout_millis: float = 10_000) -> bool: - """Force flush any buffered data. - - Note: This exporter doesn't buffer data internally, so this is a no-op. - - Args: - timeout_millis: Timeout in milliseconds - - Returns: - True (always succeeds) - - """ - return True - - def shutdown(self) -> None: - """Shutdown the exporter and release resources.""" - with self._shutdown_lock: - if self._shutdown: - return - - self._shutdown = True - self._shutdown_event.set() - - # Close HTTP session (only if we own it) - if hasattr(self, "_session") and getattr(self, "_owns_session", True): - self._session.close() - - # Close callback manager - if hasattr(self, "_callback_manager"): - self._callback_manager.close() diff --git a/olive/telemetry/library/options.py b/olive/telemetry/library/options.py index dd934cad2d..7367c0a062 100644 --- a/olive/telemetry/library/options.py +++ b/olive/telemetry/library/options.py @@ -7,9 +7,7 @@ from dataclasses import dataclass, field from enum import Enum -from typing import Callable, Optional - -import requests +from typing import Optional from olive.telemetry.library.connection_string_parser import ConnectionStringParser @@ -35,7 +33,6 @@ class OneCollectorTransportOptions: max_items_per_payload: int = DEFAULT_MAX_ITEMS_PER_PAYLOAD compression: CompressionType = CompressionType.DEFLATE timeout_seconds: float = 10.0 - http_client_factory: Optional[Callable[[], requests.Session]] = None def validate(self) -> None: """Validate the transport options. @@ -62,6 +59,7 @@ class OneCollectorExporterOptions: """Configuration options for OneCollector exporter.""" connection_string: Optional[str] = None + service_name: Optional[str] = None transport_options: OneCollectorTransportOptions = field(default_factory=OneCollectorTransportOptions) # Internal fields populated during validation diff --git a/olive/telemetry/library/retry.py b/olive/telemetry/library/retry.py deleted file mode 100644 index 9f0cc7cfd8..0000000000 --- a/olive/telemetry/library/retry.py +++ /dev/null @@ -1,98 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""Retry logic with exponential backoff for OneCollector exporter.""" - -import random -import threading -from time import time -from typing import Callable, Optional - -from olive.telemetry.library.event_source import event_source -from olive.telemetry.library.transport import HttpJsonPostTransport - - -class RetryHandler: - """Handles retry logic with exponential backoff and jitter. - - Implements retry strategy matching the .NET implementation. - """ - - def __init__(self, max_retries: int = 6, base_delay: float = 1.0, max_delay: float = 60.0): - """Initialize retry handler. - - Args: - max_retries: Maximum number of retry attempts - base_delay: Base delay for exponential backoff (seconds) - max_delay: Maximum delay between retries (seconds) - - """ - self.max_retries = max_retries - self.base_delay = base_delay - self.max_delay = max_delay - - def execute_with_retry( - self, - operation: Callable[[], tuple[bool, Optional[int]]], - deadline_sec: float, - shutdown_event: threading.Event, - ) -> bool: - """Execute an operation with retry logic. - - Args: - operation: Function that returns (success, status_code) - deadline_sec: Absolute deadline timestamp - shutdown_event: Event to signal shutdown - - Returns: - True if operation succeeded, False otherwise - - """ - for retry_num in range(self.max_retries): - # Check if we've exceeded the deadline - remaining_time = deadline_sec - time() - if remaining_time <= 0: - return False - - try: - # Execute the operation - success, status_code = operation() - - if success: - return True - - # Check if response is retryable - if not HttpJsonPostTransport.is_retryable(status_code): - return False - - except Exception as ex: - event_source.export_exception_thrown("RetryHandler", ex) - - # Last retry - don't wait - if retry_num + 1 == self.max_retries: - return False - - # Last retry - failed - if retry_num + 1 == self.max_retries: - return False - - # Calculate backoff with exponential increase and jitter - backoff = min(self.base_delay * (2**retry_num), self.max_delay) - # Add +/-20% jitter - backoff *= random.uniform(0.8, 1.2) - - # Don't wait longer than remaining time - remaining_time = deadline_sec - time() - wait_time = min(backoff, remaining_time) - - if wait_time <= 0: - return False - - # Wait with ability to interrupt on shutdown - if shutdown_event.wait(wait_time): - # Shutdown occurred - return False - - return False diff --git a/olive/telemetry/library/serialization.py b/olive/telemetry/library/serialization.py index 069f85d7e1..8bc0a1a6bd 100644 --- a/olive/telemetry/library/serialization.py +++ b/olive/telemetry/library/serialization.py @@ -7,10 +7,13 @@ import base64 import json +import math from datetime import date, datetime, time, timedelta, timezone from typing import Any from uuid import UUID +from olive.telemetry.telemetry_redaction import scrub_string_for_telemetry + class CommonSchemaJsonSerializationHelper: """Helper class for serializing values to Common Schema JSON format. @@ -41,7 +44,11 @@ def serialize_value(value: Any) -> Any: return value # Numeric types - if isinstance(value, (int, float)): + if isinstance(value, int): + return value + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError("Telemetry numeric values must be finite") return value # String @@ -80,22 +87,30 @@ def serialize_value(value: Any) -> Any: return base64.b64encode(bytes(value)).decode("ascii") # Arrays/Lists - if isinstance(value, (list, tuple)): - return [CommonSchemaJsonSerializationHelper.serialize_value(item) for item in value] + if isinstance(value, (list, tuple, set, frozenset)): + items = [CommonSchemaJsonSerializationHelper.serialize_value(item) for item in value] + if isinstance(value, (set, frozenset)): + items.sort(key=lambda item: (type(item).__name__, repr(item))) + return items # Dictionary/Map if isinstance(value, dict): result = {} - for k, v in value.items(): - if k: # Skip empty keys - result[str(k)] = CommonSchemaJsonSerializationHelper.serialize_value(v) - return result + collisions = set() + for k, v in sorted(value.items(), key=lambda item: str(item[0])): + key = scrub_string_for_telemetry(str(CommonSchemaJsonSerializationHelper.serialize_value(k))) + if key: + if key in result: + collisions.add(key) + else: + result[key] = CommonSchemaJsonSerializationHelper.serialize_value(v) + return {key: value for key, value in result.items() if key not in collisions} # Default: convert to string try: - return str(value) + return scrub_string_for_telemetry(str(value)) except Exception: - return f"ERROR: type {type(value).__name__} is not supported" + return f"[unsupported:{type(value).__name__}]" @staticmethod def create_event_envelope( @@ -138,4 +153,6 @@ def serialize_to_json_bytes(envelope: dict[str, Any]) -> bytes: UTF-8 encoded JSON bytes """ - return json.dumps(envelope, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + return json.dumps(envelope, ensure_ascii=False, separators=(",", ":"), allow_nan=False, sort_keys=True).encode( + "utf-8" + ) diff --git a/olive/telemetry/library/telemetry_logger.py b/olive/telemetry/library/telemetry_logger.py deleted file mode 100644 index 7eb236e759..0000000000 --- a/olive/telemetry/library/telemetry_logger.py +++ /dev/null @@ -1,197 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- - -"""High-level telemetry logger facade for easy usage.""" - -import logging -import uuid -from typing import Any, Callable, Optional - -from opentelemetry._logs import set_logger_provider -from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler -from opentelemetry.sdk._logs.export import BatchLogRecordProcessor -from opentelemetry.sdk.resources import Resource - -from olive.telemetry.library.exporter import OneCollectorLogExporter -from olive.telemetry.library.options import OneCollectorExporterOptions -from olive.version import __version__ as VERSION - - -class TelemetryLogger: - """Singleton telemetry logger for simplified OneCollector integration. - - Provides a simple interface for logging telemetry events without - needing to configure OpenTelemetry directly. - """ - - _instance: Optional["TelemetryLogger"] = None - _default_logger: Optional["TelemetryLogger"] = None - _logger: Optional[logging.Logger] = None - _logger_exporter: Optional[OneCollectorLogExporter] = None - _logger_provider: Optional[LoggerProvider] = None - - def __new__(cls, options: Optional[OneCollectorExporterOptions] = None): - """Create or return the singleton instance. - - Args: - options: Exporter options (only used on first instantiation) - - """ - if cls._instance is None: - cls._instance = super().__new__(cls) - cls._instance._initialize(options) - - return cls._instance - - def _initialize(self, options: Optional[OneCollectorExporterOptions]) -> None: - """Initialize the logger (called only once). - - Args: - options: Exporter configuration options - - """ - try: - # Create exporter - self._logger_exporter = OneCollectorLogExporter(options=options) - - # Create logger provider - self._logger_provider = LoggerProvider( - resource=Resource.create( - { - "service.name": __name__.split(".", maxsplit=1)[0], - "service.version": VERSION, - "service.instance.id": str(uuid.uuid4()), # Unique instance ID; can double as session ID - } - ) - ) - - # Set as global logger provider - set_logger_provider(self._logger_provider) - - # Add batch processor - self._logger_provider.add_log_record_processor( - BatchLogRecordProcessor( - self._logger_exporter, - schedule_delay_millis=1000, - ) - ) - - # Create logging handler - handler = LoggingHandler(level=logging.INFO, logger_provider=self._logger_provider) - - # Set up Python logger - logger = logging.getLogger(__name__) - logger.propagate = False - logger.setLevel(logging.INFO) - logger.addHandler(handler) - - self._logger = logger - - except Exception: - # Silently fail initialization - logger will be None - self._logger = None - self._logger_provider = None - self._logger_exporter = None - - def add_global_metadata(self, metadata: dict[str, Any]) -> None: - """Add metadata fields to all telemetry events. - - Args: - metadata: Dictionary of metadata to add - - """ - if self._logger_exporter: - self._logger_exporter.add_metadata(metadata) - - def register_payload_transmitted_callback( - self, callback, include_failures: bool = False - ) -> Optional[Callable[[], None]]: - """Register a callback for payload transmission events.""" - if self._logger_exporter: - return self._logger_exporter.register_payload_transmitted_callback(callback, include_failures) - return None - - def log(self, event_name: str, attributes: Optional[dict[str, Any]] = None) -> None: - """Log a telemetry event. - - Args: - event_name: Name of the event - attributes: Optional event attributes - - """ - if self._logger: - extra = attributes if attributes else {} - self._logger.info(event_name, extra=extra) - - def disable_telemetry(self) -> None: - """Disable telemetry logging.""" - if self._logger: - self._logger.disabled = True - - def enable_telemetry(self) -> None: - """Enable telemetry logging.""" - if self._logger: - self._logger.disabled = False - - def shutdown(self) -> None: - """Shutdown the telemetry logger and flush pending data.""" - if self._logger_provider: - self._logger_provider.shutdown() - - @classmethod - def get_default_logger(cls, connection_string: Optional[str] = None) -> "TelemetryLogger": - """Get or create the default telemetry logger. - - Args: - connection_string: OneCollector connection string (only used on first call) - - Returns: - TelemetryLogger instance - - """ - if cls._default_logger is None: - options = None - if connection_string: - options = OneCollectorExporterOptions(connection_string=connection_string) - cls._default_logger = cls(options=options) - - return cls._default_logger - - @classmethod - def shutdown_default_logger(cls) -> None: - """Shutdown the default telemetry logger.""" - if cls._default_logger: - cls._default_logger.shutdown() - cls._default_logger = None - - -def get_telemetry_logger(connection_string: Optional[str] = None) -> TelemetryLogger: - """Get or create the default telemetry logger. - - Args: - connection_string: OneCollector connection string (only used on first call) - - Returns: - TelemetryLogger instance - - """ - return TelemetryLogger.get_default_logger(connection_string=connection_string) - - -def log_event(event_name: str, attributes: Optional[dict[str, Any]] = None) -> None: - """Log a telemetry event using the default logger. - - Args: - event_name: Name of the event - attributes: Optional event attributes - - """ - logger = get_telemetry_logger() - logger.log(event_name, attributes) - - -def shutdown_telemetry() -> None: - """Shutdown the default telemetry logger.""" - TelemetryLogger.shutdown_default_logger() diff --git a/olive/telemetry/library/transport.py b/olive/telemetry/library/transport.py index 93db8ab27f..f75d8d3c28 100644 --- a/olive/telemetry/library/transport.py +++ b/olive/telemetry/library/transport.py @@ -3,259 +3,171 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""HTTP transport implementation for OneCollector exporter.""" +"""HTTP transport for the OneCollector exporter (standard library only). + +Posts Common Schema JSON to the OneCollector endpoint using ``urllib`` so the +telemetry pipeline has no third-party dependency. +""" import gzip +import queue +import threading +import time +import urllib.error +import urllib.request import zlib -from abc import ABC, abstractmethod from io import BytesIO -from typing import TYPE_CHECKING, Callable, Optional - -import requests +from typing import Callable, Optional -from olive.telemetry.library.event_source import event_source from olive.telemetry.library.options import CompressionType -if TYPE_CHECKING: - from olive.telemetry.library.callback_manager import CallbackManager, PayloadTransmittedCallbackArgs - - -class ITransport(ABC): - """Abstract base class for transports.""" - - @abstractmethod - def send(self, payload: bytes, timeout_sec: float, item_count: int = 1) -> tuple[bool, Optional[int]]: - """Send a payload. - - Args: - payload: The data to send - timeout_sec: Timeout in seconds - item_count: Number of items in the payload (for callbacks) - - Returns: - Tuple of (success, status_code) - - """ - - @abstractmethod - def register_payload_transmitted_callback( - self, callback: Callable[["PayloadTransmittedCallbackArgs"], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback for payload transmission events. - - Args: - callback: Function to call when payload is transmitted - include_failures: Whether to invoke callback on failures - - Returns: - Function to call to unregister the callback - - """ - -class HttpJsonPostTransport(ITransport): - """HTTP JSON POST transport implementation. - - Sends telemetry data to OneCollector via HTTP POST with JSON payload. - """ +class HttpJsonPostTransport: + """HTTP JSON POST transport using ``urllib`` (no third-party dependency).""" def __init__( self, endpoint: str, ikey: str, compression: CompressionType, - session: requests.Session, - callback_manager: Optional["CallbackManager"] = None, - sdk_version: str = "OTel-python-1.0.0", + sdk_version: str = "py-olive-1.0.0", ): - """Initialize the HTTP transport. - - Args: - endpoint: OneCollector endpoint URL - ikey: Instrumentation key - compression: Compression type to use - session: Requests session for connection pooling - callback_manager: Optional callback manager for payload events - sdk_version: SDK version string - - """ self.endpoint = endpoint self.ikey = ikey self.compression = compression - self.session = session self.sdk_version = sdk_version - self.callback_manager = callback_manager - # Build base headers self.headers = { "x-apikey": ikey, - "User-Agent": "Python/3 HttpClient", - "Host": "mobile.events.data.microsoft.com", + "User-Agent": "Python/3 urllib", "Content-Type": "application/x-json-stream; charset=utf-8", "sdk-version": sdk_version, "NoResponseBody": "true", } - if compression != CompressionType.NO_COMPRESSION: self.headers["Content-Encoding"] = compression.value + self._worker_lock = threading.Lock() + self._inflight_worker: Optional[threading.Thread] = None + self._inflight_results = None + self._inflight_request_key = None - def register_payload_transmitted_callback( - self, callback: Callable[["PayloadTransmittedCallbackArgs"], None], include_failures: bool = False - ) -> Callable[[], None]: - """Register a callback for payload transmission events. - - Args: - callback: Function to call when payload is transmitted - include_failures: Whether to invoke callback on failures - - Returns: - Function to call to unregister the callback - - """ - if self.callback_manager is None: - # Import here to avoid circular dependency - from olive.telemetry.library.callback_manager import CallbackManager - - self.callback_manager = CallbackManager() - - return self.callback_manager.register(callback, include_failures) - - def send(self, payload: bytes, timeout_sec: float, item_count: int = 1) -> tuple[bool, Optional[int]]: - """Send payload via HTTP POST. - - Args: - payload: Uncompressed payload bytes - timeout_sec: Request timeout in seconds - item_count: Number of items in the payload (for callbacks) - - Returns: - Tuple of (success, status_code) - - """ - payload_size_bytes = len(payload) - + def send( + self, + payload: bytes, + timeout_sec: float, + item_count: int = 1, + on_send_admitted: Optional[Callable[[], None]] = None, + ) -> tuple[bool, Optional[int]]: + """Send payload via HTTP POST. Returns (success, status_code).""" try: - # Compress payload compressed_payload = self._compress(payload) - - # Update headers with content length headers = {**self.headers, "Content-Length": str(len(compressed_payload))} + request = urllib.request.Request(url=self.endpoint, data=compressed_payload, headers=headers, method="POST") - # Send request - try: - response = self.session.post( - url=self.endpoint, data=compressed_payload, headers=headers, timeout=timeout_sec - ) - except (requests.exceptions.ConnectionError, requests.exceptions.Timeout): - # Retry once on transient transport errors - response = self.session.post( - url=self.endpoint, data=compressed_payload, headers=headers, timeout=timeout_sec - ) - - # Check response - success = response.ok - status_code = response.status_code - - # Invoke callbacks - if self.callback_manager: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - - self.callback_manager.notify( - PayloadTransmittedCallbackArgs( - succeeded=success, - status_code=status_code, - payload_size_bytes=payload_size_bytes, - item_count=item_count, - payload_bytes=payload, - ) - ) - - if success: - return True, status_code - else: - # Log error response - if event_source.is_error_logging_enabled: - collector_error = response.headers.get("Collector-Error", "") - error_details = response.text[:100] if response.text else "" - event_source.http_transport_error_response( - "HttpJsonPost", status_code, collector_error, error_details - ) - return False, status_code - - except requests.exceptions.Timeout: - # Invoke failure callbacks - if self.callback_manager: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - - self.callback_manager.notify( - PayloadTransmittedCallbackArgs( - succeeded=False, - status_code=None, - payload_size_bytes=payload_size_bytes, - item_count=item_count, - payload_bytes=payload, - ) - ) - - event_source.transport_exception_thrown("HttpJsonPost", Exception("Request timeout")) + if on_send_admitted is not None: + on_send_admitted() + success, status_code = self._do_request(request, timeout_sec) + return success, status_code + except Exception: return False, None - except Exception as ex: - # Invoke failure callbacks - if self.callback_manager: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs - self.callback_manager.notify( - PayloadTransmittedCallbackArgs( - succeeded=False, - status_code=None, - payload_size_bytes=payload_size_bytes, - item_count=item_count, - payload_bytes=payload, - ) - ) + def _do_request(self, request: "urllib.request.Request", timeout_sec: float) -> tuple[bool, Optional[int]]: + """Run the request behind a wall-clock deadline, including DNS resolution.""" + request_key = (request.full_url, bytes(request.data or b"")) + with self._worker_lock: + worker = self._inflight_worker + if worker is not None: + if worker.is_alive(): + return (False, None) + result = self._consume_inflight_result() + if self._inflight_request_key == request_key: + self._clear_inflight() + return result + self._clear_inflight() + + results = queue.Queue(maxsize=1) + worker = threading.Thread( + target=self._run_request_worker, + args=(request, timeout_sec, results), + name="olive-telemetry-http", + daemon=True, + ) + self._inflight_worker = worker + self._inflight_results = results + self._inflight_request_key = request_key + worker.start() + + worker.join(max(0.0, timeout_sec)) + with self._worker_lock: + if worker.is_alive(): + return (False, None) + result = self._consume_inflight_result() + self._clear_inflight() + return result - event_source.transport_exception_thrown("HttpJsonPost", ex) - return False, None + @staticmethod + def _run_request_worker(request, timeout_sec: float, results) -> None: + try: + result = HttpJsonPostTransport._do_request_blocking(request, timeout_sec) + except Exception: + result = (False, None) + try: + results.put_nowait(result) + except queue.Full: + pass - def _compress(self, data: bytes) -> bytes: - """Compress data according to configured compression type. + def _consume_inflight_result(self) -> tuple[bool, Optional[int]]: + try: + return self._inflight_results.get_nowait() + except (AttributeError, queue.Empty): + return (False, None) - Args: - data: Uncompressed data + def _clear_inflight(self) -> None: + self._inflight_worker = None + self._inflight_results = None + self._inflight_request_key = None - Returns: - Compressed data + @staticmethod + def _do_request_blocking(request: "urllib.request.Request", timeout_sec: float) -> tuple[bool, Optional[int]]: + """Perform the request, retrying once on a transient connection error.""" + deadline = time.monotonic() + max(0.0, timeout_sec) + for attempt in range(2): + remaining = deadline - time.monotonic() + if remaining <= 0.0: + return (False, None) + try: + with urllib.request.urlopen(request, timeout=remaining) as response: + status = getattr(response, "status", response.getcode()) + return (200 <= status < 300, status) + except urllib.error.HTTPError as http_err: + # Server responded with a non-2xx status (4xx/5xx): not retried here. + try: + http_err.close() + except Exception: + # The HTTP status remains authoritative if response cleanup fails. + pass + return (False, http_err.code) + except (urllib.error.URLError, TimeoutError, OSError): + # Connection-level failure: retry once, then give up. + if attempt == 0: + continue + return (False, None) + return (False, None) - """ + def _compress(self, data: bytes) -> bytes: if self.compression == CompressionType.DEFLATE: - # Raw deflate (no zlib header) compressor = zlib.compressobj(wbits=-zlib.MAX_WBITS) - compressed = compressor.compress(data) - compressed += compressor.flush() - return compressed - + return compressor.compress(data) + compressor.flush() elif self.compression == CompressionType.GZIP: gzip_buffer = BytesIO() with gzip.GzipFile(fileobj=gzip_buffer, mode="w") as gzip_file: gzip_file.write(data) return gzip_buffer.getvalue() - - else: # NO_COMPRESSION - return data + return data @staticmethod def is_retryable(status_code: Optional[int]) -> bool: - """Check if a response status code indicates the request should be retried. - - Args: - status_code: HTTP status code, or None if request failed - - Returns: - True if request should be retried - - """ + """Whether a response status indicates the request should be retried.""" if status_code is None: return True # Network errors are retryable - - # Retryable status codes - return status_code in {408, 429, 500, 502, 503, 504} + return status_code in {408, 429} or 500 <= status_code <= 599 diff --git a/olive/telemetry/offline_store.py b/olive/telemetry/offline_store.py new file mode 100644 index 0000000000..16603d4602 --- /dev/null +++ b/olive/telemetry/offline_store.py @@ -0,0 +1,274 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""SQLite-backed durable queue for telemetry events. + +A deliberately small subset of the Microsoft 1DS C++ SDK offline store +(cpp_client_telemetry/lib/offline/OfflineStorage_SQLite.cpp): a single FIFO +table of serialized event payloads. An uploader drains it, deleting rows on +success, dropping them on a permanent (non-retryable) send result, and leaving +them for the next attempt on a transient failure. Because every event is +written to disk before any network call, the process can exit at any time +without losing data and without an exit-time flush. + +Uses only the Python standard library (``sqlite3``), so it adds no dependency. + +Intentionally omitted from the full 1DS store (not needed for low-volume CLI +telemetry): per-event priority (``latency``), persistence classes, leases, +per-row retry counters, tenant multiplexing, and the ``settings`` table. The +schema version is tracked with SQLite's built-in ``PRAGMA user_version``. +""" + +import os +import sqlite3 +import threading +import time +from contextlib import contextmanager, suppress +from pathlib import Path +from typing import Optional + +SCHEMA_VERSION = 3 + + +def _chmod_best_effort(path: str, mode: int) -> None: + if os.name == "nt" or not path: + return + try: + Path(path).chmod(mode) + except OSError: + # Permission tightening is best-effort on filesystems that do not support chmod. + pass + + +class OfflineEventStore: + """Durable FIFO queue of serialized telemetry event payloads. + + All methods are best-effort and swallow storage errors: telemetry must + never crash the host application. Thread-safe via a per-instance lock; + tolerant of concurrent processes via WAL mode + ``busy_timeout``. + """ + + def __init__(self, db_path: str, max_records: int = 2048, busy_timeout_ms: int = 3000): + self._db_path = db_path + self._max_records = max_records + # When full, trim back to this watermark so we don't trim on every insert. + self._trim_target = max(1, (max_records * 3) // 4) + self._busy_timeout_ms = busy_timeout_ms + self._lock = threading.Lock() + self._conn: Optional[sqlite3.Connection] = None + self._initialize() + + def _initialize(self) -> None: + parent = os.path.dirname(self._db_path) + try: + os.makedirs(parent, mode=0o700, exist_ok=True) + _chmod_best_effort(parent, 0o700) + except Exception: + # sqlite3.connect below reports whether storage can actually be opened. + pass + conn = None + try: + conn = sqlite3.connect( + self._db_path, + timeout=self._busy_timeout_ms / 1000.0, + check_same_thread=False, + ) + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute(f"PRAGMA busy_timeout={self._busy_timeout_ms}") + conn.execute("BEGIN IMMEDIATE") + conn.execute( + "CREATE TABLE IF NOT EXISTS events " + "(id INTEGER PRIMARY KEY AUTOINCREMENT, payload BLOB NOT NULL, " + "available_at REAL NOT NULL DEFAULT 0, acknowledged INTEGER NOT NULL DEFAULT 0)" + ) + columns = {row[1] for row in conn.execute("PRAGMA table_info(events)")} + if "available_at" not in columns: + conn.execute("ALTER TABLE events ADD COLUMN available_at REAL NOT NULL DEFAULT 0") + if "acknowledged" not in columns: + conn.execute("ALTER TABLE events ADD COLUMN acknowledged INTEGER NOT NULL DEFAULT 0") + conn.execute(f"PRAGMA user_version={SCHEMA_VERSION}") + conn.commit() + conn.execute(f"PRAGMA busy_timeout={self._busy_timeout_ms}") + self._conn = conn + self._harden_permissions() + except Exception: + if conn is not None: + with suppress(Exception): + conn.close() + self._conn = None + + def _ensure_open(self) -> bool: + return self._conn is not None + + def _harden_permissions(self) -> None: + _chmod_best_effort(os.path.dirname(self._db_path), 0o700) + for path in (self._db_path, self._db_path + "-wal", self._db_path + "-shm"): + if os.path.exists(path): + _chmod_best_effort(path, 0o600) + + @contextmanager + def _bounded_busy_timeout(self, deadline: Optional[float]): + if deadline is None or self._conn is None: + yield + return + remaining_ms = max(0, int((deadline - time.monotonic()) * 1000)) + self._conn.execute(f"PRAGMA busy_timeout={min(self._busy_timeout_ms, remaining_ms)}") + try: + yield + finally: + with suppress(Exception): + self._conn.execute(f"PRAGMA busy_timeout={self._busy_timeout_ms}") + + @property + def is_open(self) -> bool: + with self._lock: + return self._ensure_open() + + @property + def db_path(self) -> str: + return self._db_path + + def store(self, payload: bytes) -> bool: + """Append one serialized event; trims the oldest rows if over capacity.""" + return self._store(payload, 0.0) is not None + + def reserve(self, payload: bytes, available_after_seconds: float) -> Optional[int]: + """Persist an event but defer draining until it is released or the delay expires.""" + return self._store(payload, time.time() + max(0.0, available_after_seconds)) + + def _store(self, payload: bytes, available_at: float) -> Optional[int]: + if not payload: + return None + with self._lock: + if not self._ensure_open(): + return None + try: + cursor = self._conn.execute( + "INSERT INTO events (payload, available_at) VALUES (?, ?)", + (sqlite3.Binary(payload), available_at), + ) + count = self._conn.execute("SELECT COUNT(*) FROM events").fetchone()[0] + if count > self._max_records: + self._conn.execute( + "DELETE FROM events WHERE id IN " + "(SELECT id FROM events WHERE available_at <= ? ORDER BY id ASC LIMIT ?)", + (time.time(), count - self._trim_target), + ) + self._conn.commit() + self._harden_permissions() + return int(cursor.lastrowid) + except Exception: + with suppress(Exception): + self._conn.rollback() + return None + + def get_batch(self, max_count: int) -> list[tuple[int, bytes]]: + """Return up to ``max_count`` oldest events as (id, payload) pairs.""" + return self.get_batch_for_upload(max_count) or [] + + def get_batch_for_upload( + self, max_count: int, deadline: Optional[float] = None + ) -> Optional[list[tuple[int, bytes]]]: + """Return an uploadable batch, or None when local storage failed.""" + with self._lock: + if not self._ensure_open(): + return None + try: + with self._bounded_busy_timeout(deadline): + rows = self._conn.execute( + "SELECT id, payload FROM events " + "WHERE available_at <= ? AND acknowledged=0 ORDER BY id ASC LIMIT ?", + (time.time(), max_count if max_count > 0 else -1), + ).fetchall() + return [(r[0], bytes(r[1])) for r in rows] + except Exception: + return None + + def get_acknowledged_ids(self, max_count: int, deadline: Optional[float] = None) -> Optional[list[int]]: + """Return terminally handled rows that only need local deletion.""" + with self._lock: + if not self._ensure_open(): + return None + try: + with self._bounded_busy_timeout(deadline): + rows = self._conn.execute( + "SELECT id FROM events WHERE acknowledged=1 ORDER BY id ASC LIMIT ?", + (max_count if max_count > 0 else -1,), + ).fetchall() + return [int(row[0]) for row in rows] + except Exception: + return None + + def release(self, row_id: int, payload: Optional[bytes] = None) -> bool: + """Make a reserved event drainable, optionally replacing its payload.""" + with self._lock: + if not self._ensure_open(): + return False + try: + if payload is None: + cursor = self._conn.execute("UPDATE events SET available_at=0 WHERE id=?", (row_id,)) + else: + cursor = self._conn.execute( + "UPDATE events SET payload=?, available_at=0 WHERE id=?", + (sqlite3.Binary(payload), row_id), + ) + self._conn.commit() + return cursor.rowcount == 1 + except Exception: + with suppress(Exception): + self._conn.rollback() + return False + + def acknowledge(self, ids: list[int], deadline: Optional[float] = None) -> bool: + """Persist that rows were delivered or permanently rejected.""" + if not ids: + return True + with self._lock: + if not self._ensure_open(): + return False + try: + with self._bounded_busy_timeout(deadline): + self._conn.executemany("UPDATE events SET acknowledged=1 WHERE id=?", [(i,) for i in ids]) + self._conn.commit() + return True + except Exception: + with suppress(Exception): + self._conn.rollback() + return False + + def delete(self, ids: list[int], deadline: Optional[float] = None) -> bool: + """Remove rows by id (after a successful upload or a permanent drop).""" + if not ids: + return True + with self._lock: + if not self._ensure_open(): + return False + try: + with self._bounded_busy_timeout(deadline): + self._conn.executemany("DELETE FROM events WHERE id=?", [(i,) for i in ids]) + self._conn.commit() + return True + except Exception: + # Failed deletes leave rows durable for a later drain attempt. + with suppress(Exception): + self._conn.rollback() + return False + + def count(self) -> int: + with self._lock: + if not self._ensure_open(): + return 0 + try: + return int(self._conn.execute("SELECT COUNT(*) FROM events").fetchone()[0]) + except Exception: + return 0 + + def close(self) -> None: + with self._lock: + if self._conn is not None: + with suppress(Exception): + self._conn.close() + self._conn = None diff --git a/olive/telemetry/process_lock.py b/olive/telemetry/process_lock.py new file mode 100644 index 0000000000..14ce23fcb4 --- /dev/null +++ b/olive/telemetry/process_lock.py @@ -0,0 +1,117 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Cross-platform single-holder advisory lock (standard library only). + +Used so that, when several processes on one device share a telemetry database, +only one of them runs the uploader's drain loop at a time. Other processes keep +writing events durably to the store; the lock holder drains everyone's rows. +This avoids the same event being uploaded twice by concurrent drainers without +needing per-row reservation bookkeeping. + +The lock is an OS advisory lock on a sidecar file (``msvcrt`` on Windows, +``fcntl`` on POSIX). It is released explicitly and also by the OS when the +process exits, so a crashed holder never blocks other processes permanently. +""" + +import errno +import os +import time +from contextlib import suppress + + +class ProcessDrainLock: + """Non-blocking exclusive advisory lock backed by a sidecar file.""" + + def __init__(self, lock_path: str): + self._lock_path = lock_path + self._fh = None + self._posix_lock_api = None + + @property + def held(self) -> bool: + return self._fh is not None + + def acquire(self, timeout_seconds: float = 0.0) -> bool: + """Try to acquire the lock without blocking. Returns True if held.""" + if self._fh is not None: + return True + deadline = time.monotonic() + max(0.0, timeout_seconds) + while True: + fh = None + try: + with suppress(Exception): + os.makedirs(os.path.dirname(self._lock_path), exist_ok=True) + # The handle must remain open while the advisory lock is held. + fh = open(self._lock_path, "a+b") # noqa: SIM115 # pylint: disable=consider-using-with + if os.name == "nt": + import msvcrt + + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_NBLCK, 1) + else: + import fcntl + + try: + fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + self._posix_lock_api = "flock" + except AttributeError: + fcntl.lockf(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + self._posix_lock_api = "lockf" + except OSError as exc: + unsupported_errors = {errno.ENOSYS} + if hasattr(errno, "ENOTSUP"): + unsupported_errors.add(errno.ENOTSUP) + if hasattr(errno, "EOPNOTSUPP"): + unsupported_errors.add(errno.EOPNOTSUPP) + if exc.errno not in unsupported_errors: + raise + fcntl.lockf(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + self._posix_lock_api = "lockf" + self._fh = fh + return True + except Exception: + self._posix_lock_api = None + if fh is not None: + with suppress(Exception): + fh.close() + if time.monotonic() >= deadline: + return False + time.sleep(min(0.01, max(0.0, deadline - time.monotonic()))) + + def release(self) -> None: + if self._fh is None: + return + fh = self._fh + posix_lock_api = self._posix_lock_api + self._fh = None + self._posix_lock_api = None + try: + if os.name == "nt": + import msvcrt + + try: + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1) + except Exception: + # The OS releases the lock when the handle closes below. + pass + else: + import fcntl + + try: + if posix_lock_api == "lockf": + fcntl.lockf(fh.fileno(), fcntl.LOCK_UN) + else: + fcntl.flock(fh.fileno(), fcntl.LOCK_UN) + except Exception: + # The OS releases the lock when the handle closes below. + pass + finally: + try: + fh.close() + except Exception: + # Lock cleanup must never fail telemetry callers. + pass diff --git a/olive/telemetry/recipe_telemetry.py b/olive/telemetry/recipe_telemetry.py new file mode 100644 index 0000000000..18f42bf048 --- /dev/null +++ b/olive/telemetry/recipe_telemetry.py @@ -0,0 +1,457 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import functools +import json +import re +from copy import deepcopy +from os import PathLike +from pathlib import Path, PurePosixPath, PureWindowsPath +from typing import TYPE_CHECKING, Any, Optional, Union + +from olive.common.config_utils import load_config_file +from olive.common.utils import hash_dict +from olive.package_config import OlivePackageConfig +from olive.systems.common import SystemType +from olive.telemetry.telemetry import is_ci_environment +from olive.telemetry.telemetry_redaction import ( + is_environment_config_key_for_telemetry, + is_path_like_config_key_for_telemetry, + is_sensitive_config_key_for_telemetry, + normalize_config_key_for_telemetry, + scrub_string_for_telemetry, +) + +if TYPE_CHECKING: + from olive.workflows.run.config import RunConfig + +RECIPE_HASH_REDACTED_VALUE = "" +CONFIG_REFERENCE_REDACTED_VALUE = "" +CONFIG_CALLABLE_REDACTED_VALUE = "" +CONFIG_UNKNOWN_REDACTED_VALUE = "" +CONFIG_SECRET_REDACTED_VALUE = "" +RECIPE_HASH_REDACTED_KEYS = { + "output_dir", + "cache_dir", + "tempdir", + "additional_files", + "dockerfile", + "build_context_path", + "python_environment_path", + "prepend_to_path", + "script_dir", + "model_script", + # package_config is tracked separately via package_config_provided and + # package_config_overrides, but excluded from recipe_hash because it is an + # environment/infrastructure path. + "package_config", + "work_dir", +} +CONFIG_SNAPSHOT_REDACTED_KEYS = RECIPE_HASH_REDACTED_KEYS | { + "model_path", + "_name_or_path", + "adapter_path", + "user_script", +} +HF_MODEL_IDENTIFIER_KEYS = {"model_path", "_name_or_path"} +CONFIG_REFERENCE_KEYS = {"host", "target", "evaluator"} +LOCAL_MODEL_FILE_SUFFIXES = {".bin", ".model", ".onnx", ".pb", ".pt", ".pth", ".safetensors", ".tflite"} +HF_CACHE_MODEL_PATTERN = re.compile(r"(?:^|[\\/])models--([^\\/]+)--([^\\/]+)(?:[\\/]|$)") +HF_REPO_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*(/[A-Za-z0-9][A-Za-z0-9._-]*)?$") +_NO_OVERRIDE = object() + + +def _build_recipe_result_metadata( + run_config_input: Union[str, Path, dict], + run_config_telemetry_input: Optional[Any], + run_config: Optional["RunConfig"], + recipe_telemetry_metadata: Optional[dict[str, Any]], + *, + list_required_packages: bool, + package_config_input: Optional[Union[str, Path, dict]], + package_config_provided: bool, +) -> dict[str, Any]: + metadata = dict(recipe_telemetry_metadata or {}) + default_source, default_format = _classify_run_config_source(run_config_input) + metadata.setdefault("recipe_source", default_source) + metadata.setdefault("recipe_format", default_format) + metadata.setdefault("execution_mode", "list_required_packages" if list_required_packages else "run") + metadata.setdefault("package_config_provided", package_config_provided) + config_overrides = metadata.pop("config_overrides", _NO_OVERRIDE) + if config_overrides is _NO_OVERRIDE: + config_overrides = _build_config_overrides(run_config_telemetry_input) + else: + config_overrides = _sanitize_provided_config_snapshot(config_overrides) + if config_overrides is not None: + metadata["config_overrides"] = config_overrides + package_config_overrides = metadata.pop("package_config_overrides", _NO_OVERRIDE) + if package_config_overrides is not _NO_OVERRIDE: + package_config_overrides = _sanitize_provided_config_snapshot(package_config_overrides) + elif package_config_provided: + package_config_overrides = _build_package_config_overrides(package_config_input) + else: + package_config_overrides = None + if package_config_overrides is not None: + metadata["package_config_overrides"] = package_config_overrides + metadata["is_ci"] = is_ci_environment() + + if run_config is None: + metadata.setdefault("recipe_name", metadata.get("recipe_command") or "WorkflowRun") + return metadata + + run_config_json = run_config.to_json(make_absolute=False) + model_metadata = _extract_input_model_metadata(run_config_json["input_model"]) + target_metadata = _extract_target_metadata(run_config) + host_metadata = _extract_host_metadata(run_config) + pass_types = _get_used_pass_types(run_config) + + metadata.setdefault("recipe_name", metadata.get("recipe_command") or run_config.workflow_id) + metadata.setdefault("workflow_id", run_config.workflow_id) + metadata.setdefault("recipe_hash", _build_recipe_hash(run_config_json)) + metadata.setdefault("input_model_type", run_config.input_model.type) + metadata.setdefault("input_model_source", model_metadata["input_model_source"]) + metadata.setdefault("model_task", model_metadata["model_task"]) + _set_metadata_if_present(metadata, target_metadata) + _set_metadata_if_present(metadata, host_metadata) + metadata.setdefault("pass_types", ";".join(pass_types)) + metadata.setdefault("pass_count", len(pass_types)) + metadata.setdefault("data_config_count", len(run_config.data_configs)) + metadata.setdefault("search_enabled", bool(run_config.engine.search_strategy)) + return metadata + + +def _classify_run_config_source(run_config_input: Any) -> tuple[str, str]: + if isinstance(run_config_input, dict): + return "config_dict", "dict" + + if isinstance(run_config_input, (str, PathLike)): + suffix = Path(run_config_input).suffix.lstrip(".").lower() + return "config_file", suffix or "unknown" + + return "config_object", "object" + + +def _build_config_overrides(config_input: Any) -> Optional[str]: + try: + config_data = _load_config_input_for_telemetry(config_input) + if config_data is None: + return None + + snapshot = _sanitize_config_snapshot(config_data) + if snapshot in (None, {}, []): + return None + + return json.dumps(snapshot, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + except Exception: + return None + + +def _sanitize_provided_config_snapshot(value: Any) -> Optional[str]: + if isinstance(value, str): + try: + value = json.loads(value) + except (TypeError, ValueError): + return None + return _build_config_overrides(value) + + +def _build_package_config_overrides(config_input: Any) -> Optional[str]: + try: + config_data = _load_config_input_for_telemetry(config_input) + if not isinstance(config_data, dict): + return None + + default_config = _load_default_package_config_for_telemetry() + baseline = ( + _normalize_package_config_snapshot(default_config) if isinstance(default_config, dict) else _NO_OVERRIDE + ) + overrides = _extract_config_overrides(_normalize_package_config_snapshot(config_data), baseline) + if overrides is _NO_OVERRIDE: + return None + + snapshot = _sanitize_config_snapshot(overrides) + if not isinstance(snapshot, dict): + return None + + return json.dumps(snapshot, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + except Exception: + return None + + +@functools.lru_cache +def _load_default_package_config_for_telemetry() -> Optional[dict[str, Any]]: + try: + default_config = load_config_file(OlivePackageConfig.get_default_config_path()) + except Exception: + return None + + return default_config if isinstance(default_config, dict) else None + + +def _normalize_package_config_snapshot(config_data: Any) -> Any: + if not isinstance(config_data, dict): + return config_data + + normalized = deepcopy(config_data) + passes = normalized.get("passes") + if isinstance(passes, dict): + normalized["passes"] = {str(pass_name).lower(): pass_config for pass_name, pass_config in passes.items()} + return normalized + + +def _extract_config_overrides(value: Any, baseline: Any = _NO_OVERRIDE) -> Any: + if baseline is _NO_OVERRIDE: + return deepcopy(value) + + if isinstance(value, dict) and isinstance(baseline, dict): + overrides = {} + for key, child_value in value.items(): + child_override = _extract_config_overrides(child_value, baseline.get(key, _NO_OVERRIDE)) + if child_override is not _NO_OVERRIDE: + overrides[key] = child_override + if overrides: + return overrides + return _NO_OVERRIDE + + if isinstance(value, list): + if isinstance(baseline, list) and value == baseline: + return _NO_OVERRIDE + return deepcopy(value) + + if isinstance(value, tuple): + value_list = list(value) + baseline_list = list(baseline) if isinstance(baseline, tuple) else baseline + if isinstance(baseline_list, list) and value_list == baseline_list: + return _NO_OVERRIDE + return value_list + + return deepcopy(value) if value != baseline else _NO_OVERRIDE + + +def _load_config_input_for_telemetry(config_input: Any) -> Optional[Any]: + if config_input is None: + return None + if isinstance(config_input, dict): + return deepcopy(config_input) + if isinstance(config_input, (str, PathLike)): + return load_config_file(config_input) + + model_dump = getattr(config_input, "model_dump", None) + if callable(model_dump): + return model_dump(exclude_defaults=True, exclude_none=True, by_alias=True) + return None + + +def _sanitize_config_snapshot(value: Any, key: Optional[str] = None, model_type: Optional[str] = None) -> Any: + normalized_key = normalize_config_key_for_telemetry(key) + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return CONFIG_SECRET_REDACTED_VALUE + if normalized_key in HF_MODEL_IDENTIFIER_KEYS: + if str(model_type).lower() == "hfmodel": + hf_model_id = _extract_huggingface_model_id(value) + if hf_model_id: + return hf_model_id + return RECIPE_HASH_REDACTED_VALUE + if normalized_key in CONFIG_SNAPSHOT_REDACTED_KEYS or is_path_like_config_key_for_telemetry(normalized_key): + return RECIPE_HASH_REDACTED_VALUE + if normalized_key in CONFIG_REFERENCE_KEYS and isinstance(value, str): + return CONFIG_REFERENCE_REDACTED_VALUE + + if isinstance(value, dict): + child_model_type = _get_model_type(value) or model_type + if normalized_key == "systems": + return [_sanitize_config_snapshot(system, "system", child_model_type) for system in value.values()] + if normalized_key == "passes": + passes = [] + for pass_configs in value.values(): + if isinstance(pass_configs, list): + passes.extend(pass_configs) + else: + passes.append(pass_configs) + return [_sanitize_config_snapshot(pass_config, "pass", child_model_type) for pass_config in passes] + if normalized_key == "evaluators": + return [ + _sanitize_config_snapshot(evaluator, "evaluator_config", child_model_type) + for evaluator in value.values() + ] + return { + child_key: _sanitize_config_snapshot(child_value, child_key, child_model_type) + for child_key, child_value in value.items() + if child_value is not None + } + if isinstance(value, list): + return [_sanitize_config_snapshot(item, key, model_type) for item in value] + if isinstance(value, tuple): + return [_sanitize_config_snapshot(item, key, model_type) for item in value] + if isinstance(value, Path): + return RECIPE_HASH_REDACTED_VALUE + if callable(value): + return CONFIG_CALLABLE_REDACTED_VALUE + if isinstance(value, str): + return scrub_string_for_telemetry(value) + if isinstance(value, (int, float, bool)) or value is None: + return value + if hasattr(value, "value") and isinstance(value.value, (str, int, float, bool)): + return scrub_string_for_telemetry(value.value) if isinstance(value.value, str) else value.value + return CONFIG_UNKNOWN_REDACTED_VALUE + + +def _get_model_type(config: dict[str, Any]) -> Optional[str]: + model_type = config.get("type") + return str(model_type).lower() if model_type is not None else None + + +def _extract_huggingface_model_id(model_identifier: Any) -> Optional[str]: + if not isinstance(model_identifier, str): + return None + + identifier = model_identifier.strip() + if not identifier: + return None + + if identifier.startswith("https://huggingface.co/"): + parts = identifier.removeprefix("https://huggingface.co/").strip("/").split("/") + if len(parts) >= 2: + return f"{parts[0]}/{parts[1]}" + if parts and parts[0]: + return parts[0] + + if match := HF_CACHE_MODEL_PATTERN.search(identifier): + return f"{match.group(1)}/{match.group(2)}" + + if HF_REPO_ID_PATTERN.match(identifier) and not _has_local_model_file_suffix(identifier): + return identifier + + return None + + +def _extract_input_model_metadata(input_model_config: dict[str, Any]) -> dict[str, Optional[str]]: + model_config = input_model_config.get("config", {}) + model_attributes = model_config.get("model_attributes", {}) + model_task = model_attributes.get("hf_task") or model_config.get("task") + raw_identifier = model_attributes.get("_name_or_path") or model_config.get("model_path") + return { + "input_model_source": _classify_input_model_source(raw_identifier), + "model_task": str(model_task) if model_task is not None else None, + } + + +def _classify_input_model_source(model_identifier: Any) -> str: + if model_identifier is None: + return "unknown" + if isinstance(model_identifier, dict): + resource_type = model_identifier.get("type") + if resource_type == "azureml_registry_model": + return "azureml" + return "structured_resource" + + identifier = str(model_identifier) + if identifier.startswith("azureml://"): + return "azureml" + if identifier.startswith("https://huggingface.co/"): + return "huggingface_url" + if identifier.startswith(("http://", "https://")): + return "url" + + if _is_explicit_local_model_path(identifier): + suffix = PureWindowsPath(identifier).suffix or PurePosixPath(identifier).suffix + return "local_file" if suffix else "local_folder" + return "string_name" + + +def _is_explicit_local_model_path(identifier: str) -> bool: + if _has_local_model_file_suffix(identifier): + return True + return ( + identifier.startswith(("./", "../", ".\\", "..\\", "~/", "~\\", "/", "\\\\")) + or PureWindowsPath(identifier).is_absolute() + or PurePosixPath(identifier).is_absolute() + ) + + +def _has_local_model_file_suffix(identifier: str) -> bool: + suffix = PureWindowsPath(identifier).suffix or PurePosixPath(identifier).suffix + return suffix.lower() in LOCAL_MODEL_FILE_SUFFIXES + + +def _extract_target_metadata(run_config: "RunConfig") -> dict[str, Optional[str]]: + target_system = run_config.engine.target + return _extract_system_metadata(target_system, "target") + + +def _extract_host_metadata(run_config: "RunConfig") -> dict[str, Optional[str]]: + host_system = run_config.engine.host + if host_system is None: + return { + "host_system_type": SystemType.Local.value, + } + return _extract_system_metadata(host_system, "host") + + +def _extract_system_metadata(system_config: Optional[Any], field_prefix: str) -> dict[str, Optional[str]]: + system_type = system_config.type.value if system_config is not None else None + device = None + execution_provider = None + execution_providers = None + + accelerators = system_config.config.accelerators if system_config and system_config.config else None + if accelerators: + accelerator = accelerators[0] + device = str(accelerator.device) if accelerator.device is not None else None + ep_values = accelerator.get_ep_strs() or [] + if ep_values: + execution_provider = ep_values[0] + execution_providers = ";".join(ep_values) + + return { + f"{field_prefix}_system_type": system_type, + f"{field_prefix}_device": device, + f"{field_prefix}_execution_provider": execution_provider, + f"{field_prefix}_execution_providers": execution_providers, + } + + +def _set_metadata_if_present(metadata: dict[str, Any], values: dict[str, Optional[str]]) -> None: + for key, value in values.items(): + if value is not None: + metadata.setdefault(key, value) + + +def _get_used_pass_types(run_config: "RunConfig") -> list[str]: + return ( + [pass_config.type for _, pass_configs in run_config.passes.items() for pass_config in pass_configs] + if run_config.passes + else [] + ) + + +def _build_recipe_hash(run_config_json: dict[str, Any]) -> str: + sanitized = deepcopy(run_config_json) + _redact_recipe_hash_keys(sanitized) + return hash_dict(sanitized)[:16] + + +def _redact_recipe_hash_keys(value: Any, key: Optional[str] = None) -> Any: + normalized_key = normalize_config_key_for_telemetry(key) + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return CONFIG_SECRET_REDACTED_VALUE + if normalized_key in RECIPE_HASH_REDACTED_KEYS or is_path_like_config_key_for_telemetry(normalized_key): + return RECIPE_HASH_REDACTED_VALUE + if isinstance(value, dict): + for child_key in list(value): + value[child_key] = _redact_recipe_hash_keys(value[child_key], child_key) + elif isinstance(value, list): + for index, item in enumerate(value): + value[index] = _redact_recipe_hash_keys(item, key) + elif isinstance(value, tuple): + return [_redact_recipe_hash_keys(item, key) for item in value] + elif isinstance(value, Path): + return RECIPE_HASH_REDACTED_VALUE + elif isinstance(value, str): + return scrub_string_for_telemetry(value) + elif callable(value): + return CONFIG_CALLABLE_REDACTED_VALUE + elif hasattr(value, "value") and isinstance(value.value, (str, int, float, bool)): + return value.value + return value diff --git a/olive/telemetry/telemetry.py b/olive/telemetry/telemetry.py index 0ddb690e2a..ceb1b13e85 100644 --- a/olive/telemetry/telemetry.py +++ b/olive/telemetry/telemetry.py @@ -2,58 +2,73 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Thin wrapper around the OneCollector telemetry logger with event helpers.""" +"""Telemetry singleton backed by a durable SQLite event queue. + +Detailed events are serialized to Common Schema JSON and written to a per-app +SQLite store; a background uploader drains the store to Microsoft OneCollector. +The one-per-process Heartbeat uses the same durable queue. The pipeline uses only +the Python standard library (no OpenTelemetry, no requests). +""" import base64 -import errno import json import os import platform import threading import time -from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional - -from olive.telemetry.constants import CONNECTION_STRING -from olive.telemetry.deviceid import get_encrypted_device_id_and_status -from olive.telemetry.library.event_source import event_source -from olive.telemetry.library.telemetry_logger import TelemetryLogger, get_telemetry_logger -from olive.telemetry.utils import ( - _decode_cache_line, - _encode_cache_line, - _exclusive_file_lock, - get_telemetry_base_dir, +import uuid +from datetime import datetime, timezone +from typing import Any, Optional + +from olive.telemetry.deviceid import get_hashed_device_id_and_status +from olive.telemetry.library.options import OneCollectorExporterOptions +from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper +from olive.telemetry.offline_store import OfflineEventStore +from olive.telemetry.telemetry_redaction import ( + MAX_TELEMETRY_STRING_LENGTH, + scrub_config_snapshot_for_telemetry, + scrub_error_message_for_telemetry, + scrub_value_for_telemetry, ) +from olive.telemetry.uploader import EventUploader +from olive.telemetry.utils import get_telemetry_base_dir -if TYPE_CHECKING: - from olive.telemetry.library.callback_manager import PayloadTransmittedCallbackArgs +try: + from olive.version import __version__ as VERSION +except Exception: + VERSION = "unknown" # Default event names used by the high-level telemetry helpers. HEARTBEAT_EVENT_NAME = "OliveHeartbeat" +RECIPE_EVENT_NAME = "OliveRecipe" +ACTION_EVENT_NAME = "OliveAction" +ERROR_EVENT_NAME = "OliveError" # CI/CD environment variables whose presence indicates an automated pipeline. _CI_ENV_VARS = ( "CI", # GitHub Actions, GitLab CI, Travis CI, CircleCI, generic "TF_BUILD", # Azure Pipelines "GITHUB_ACTIONS", # GitHub Actions + "GITLAB_CI", # GitLab CI + "CIRCLECI", # CircleCI + "TRAVIS", # Travis CI "JENKINS_URL", # Jenkins "CODEBUILD_BUILD_ID", # AWS CodeBuild "BUILDKITE", # Buildkite + "TEAMCITY_VERSION", # TeamCity + "APPVEYOR", # AppVeyor + "BITBUCKET_BUILD_NUMBER", # Bitbucket Pipelines "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", # Azure DevOps ) -ACTION_EVENT_NAME = "OliveAction" -ERROR_EVENT_NAME = "OliveError" ALLOWED_KEYS = { HEARTBEAT_EVENT_NAME: { "device_id", - "id_status", - "os.name", - "os.version", - "os.release", - "os.arch", - "app_version", - "app_instance_id", + "device_id_status", + "os", + "os_version", + "os_release", + "os_arch", "initTs", }, ACTION_EVENT_NAME: { @@ -61,656 +76,479 @@ "action_name", "duration_ms", "success", - "app_version", - "app_instance_id", "initTs", }, ERROR_EVENT_NAME: { "exception_type", "exception_message", - "app_version", - "app_instance_id", + "initTs", + }, + RECIPE_EVENT_NAME: { + "recipe_name", + "recipe_hash", + "recipe_source", + "recipe_format", + "recipe_command", + "execution_mode", + "workflow_id", + "config_overrides", + "success", + "input_model_type", + "input_model_source", + "model_task", + "target_system_type", + "target_device", + "target_execution_provider", + "target_execution_providers", + "host_system_type", + "host_device", + "host_execution_provider", + "host_execution_providers", + "pass_types", + "pass_count", + "data_config_count", + "search_enabled", + "package_config_provided", + "package_config_overrides", + "is_ci", "initTs", }, } -CRITICAL_EVENTS = {HEARTBEAT_EVENT_NAME} -MAX_CACHE_SIZE_BYTES = 5 * 1024 * 1024 -HARD_MAX_CACHE_SIZE_BYTES = 10 * 1024 * 1024 -CACHE_FILE_NAME = "olive.json" +FIELD_NAMES = { + "device_id": "deviceId", + "device_id_status": "deviceIdStatus", + "os_version": "osVersion", + "os_release": "osRelease", + "os_arch": "osArchitecture", + "invoked_from": "invokedFrom", + "action_name": "actionName", + "duration_ms": "durationMs", + "exception_type": "exceptionType", + "exception_message": "exceptionMessage", + "recipe_name": "recipeName", + "recipe_hash": "recipeHash", + "recipe_source": "recipeSource", + "recipe_format": "recipeFormat", + "recipe_command": "recipeCommand", + "execution_mode": "executionMode", + "workflow_id": "workflowId", + "config_overrides": "configOverrides", + "input_model_type": "inputModelType", + "input_model_source": "inputModelSource", + "model_task": "modelTask", + "target_system_type": "targetSystemType", + "target_device": "targetDevice", + "target_execution_provider": "targetExecutionProvider", + "target_execution_providers": "targetExecutionProviders", + "host_system_type": "hostSystemType", + "host_device": "hostDevice", + "host_execution_provider": "hostExecutionProvider", + "host_execution_providers": "hostExecutionProviders", + "pass_types": "passTypes", + "pass_count": "passCount", + "data_config_count": "dataConfigCount", + "search_enabled": "searchEnabled", + "package_config_provided": "packageConfigProvided", + "package_config_overrides": "packageConfigOverrides", + "is_ci": "isCI", + "app_version": "LibraryVersion", + "app_instance_id": "AppSessionGuid", +} +# Per-app database file. Olive and other apps use separate files so a process +# never drains another app's events (which carry a different tenant key). +DB_FILE_NAME = "olive_telemetry.db" +CI_DB_FILE_NAME = "olive_recipe_telemetry.db" +_HEARTBEAT_RELEASE_SECONDS = 60.0 -class TelemetryCacheHandler: - """Handles caching of failed telemetry events for offline resilience. - Design decisions: - - Single shared cache file (olive.json) for simplicity - - Cache writes are synchronous (fast JSON operations don't need async) - - Cache flush runs in a separate thread (slow network I/O) - - Flush triggered on success when cached events exist - - All critical sections protected by lock to prevent race conditions - - Newline-delimited JSON format for human readability and partial corruption recovery +def _is_environment_signal_truthy(value: str) -> bool: + return value.strip().lower() not in {"", "0", "false", "no", "off"} - Assumptions: - - File I/O (JSON lines) is fast enough for synchronous execution (~microseconds) - - Network I/O is slow and should not block the callback thread - - Successful send indicates network is available to retry cached events - - Cache persists across sessions for offline resilience - """ - def __init__(self, telemetry: "Telemetry") -> None: - self._telemetry = telemetry - # Single shared cache file for all processes - self._cache_file_name = CACHE_FILE_NAME - self._shutdown = False - # Protects all shared state to prevent race conditions - self._lock = threading.Lock() - self._callback_condition = threading.Condition() - self._callbacks_item_count = 0 - self._events_logged = 0 - # Prevents concurrent flush operations - self._is_flushing = False - - def shutdown(self) -> None: - """Signal shutdown to prevent new operations. - - Note: Does NOT flush the cache. Cache persists across sessions for - offline resilience. If network is working, success callbacks already - flushed. If network is down, flushing would fail anyway. - """ - with self._lock: - self._shutdown = True +def is_ci_environment() -> bool: + """Detect CI/CD environments by checking well-known environment variables.""" + return any(_is_environment_signal_truthy(os.environ.get(var, "")) for var in _CI_ENV_VARS) - def __del__(self): - """Cleanup cache handler resources on garbage collection. - - Safety net to ensure shutdown is called even if not done explicitly. - """ - try: - self.shutdown() - except Exception: - # Silently ignore errors during cleanup - pass - def on_payload_transmitted(self, args: "PayloadTransmittedCallbackArgs") -> None: - """Telemetry payload transmission callback. +def is_telemetry_disabled_by_environment() -> bool: + """Return whether any supported environment variable requests full suppression.""" + return any( + os.environ.get(variable, "").strip().lower() in {"1", "true", "yes", "on", "y"} + for variable in ("ORT_DISABLE_TELEMETRY", "OLIVE_DISABLE_TELEMETRY") + ) - Design decisions: - - Ignore callbacks during flush (unlikely to fail during successful flush) - - On success: flush cache if any cached events exist - - On failure: write to cache immediately (synchronous for simplicity) - Assumptions: - - Successful transmission indicates network is available to retry cached events - - If flush is in progress, we already successfully sent an event, so unlikely an event would suddenly fail - - Multiple concurrent successes don't need multiple flush operations - - Failed payloads should be cached immediately to avoid loss - """ - try: - payload = None - should_flush = False - - with self._lock: - if self._shutdown: - return - - # Skip callbacks from replayed events during flush - # If a flush is in progress it means we successfully sent an event, - # so it's unlikely that an event would suddenly fail and need to be cached - # and we don't need to flush again. - if self._is_flushing: - with self._callback_condition: - self._callbacks_item_count += args.item_count - self._callback_condition.notify_all() - return - - if args.succeeded: - # Only flush if cache exists and no flush is in progress - cache_path = self.cache_path - if cache_path and cache_path.exists(): - should_flush = True - else: - payload = args.payload_bytes - - if should_flush: - # Release lock before scheduling (flush runs in separate thread) - self._schedule_flush() - elif payload: - # Write synchronously - JSON operations are fast enough - self._write_payload_to_cache(payload) - except Exception: - # Fail silently - telemetry should never crash the application - pass - finally: - with self._callback_condition: - self._callbacks_item_count += args.item_count - self._callback_condition.notify_all() - - def wait_for_callbacks(self, timeout_sec: float, during_flush: bool = False) -> bool: - deadline = time.time() + timeout_sec - while True: - with self._callback_condition: - callbacks_item_count = self._callbacks_item_count - expected_items = self._events_logged - if (during_flush or not self.is_flushing) and callbacks_item_count >= expected_items: - return True - remaining = deadline - time.time() - if remaining <= 0: - return False - with self._callback_condition: - self._callback_condition.wait(timeout=remaining) - - def record_event_logged(self, count: int = 1) -> None: - with self._callback_condition: - self._events_logged += count - - def _schedule_flush(self) -> None: - """Schedule cache flush in a separate thread to avoid blocking the callback. - - Design decisions: - - Check _is_flushing before spawning thread to avoid unnecessary threads - - Run flush in daemon thread (don't block process exit) - - Acquire lock at start to set _is_flushing flag atomically - - Always clear _is_flushing flag even if flush fails - - Assumptions: - - Flush operations are slow (network I/O) and should not block callbacks - - Daemon thread is acceptable (flush is best-effort) - """ - # Check before spawning thread to avoid unnecessary thread creation - with self._lock: - if self._shutdown or self._is_flushing: - return - self._is_flushing = True +class Telemetry: + """Per-process singleton that persists events to SQLite and uploads them. - def flush_task(): - try: - self._flush_cache() - except Exception: - # Fail silently - pass - finally: - # Always clear flag, even on exception - with self._lock: - self._is_flushing = False + Separate processes get separate in-memory singletons and coordinate only + through the shared SQLite store and its single-drainer file lock. + Use Telemetry() to get the singleton instance. + """ - thread = threading.Thread(target=flush_task, daemon=True) - thread.start() + _instance: Optional["Telemetry"] = None + _lock = threading.RLock() + _process_disabled = False + _heartbeat_enqueued = False - @property - def cache_path(self) -> Optional[Path]: - """Get the path to the telemetry cache file. + @classmethod + def get_existing_instance(cls) -> Optional["Telemetry"]: + """Return the current singleton without creating telemetry.""" + return cls._instance - Returns: - Optional[Path]: Path to cache file, or None if base directory unavailable. + @classmethod + def get_or_create_if_enabled(cls) -> Optional["Telemetry"]: + """Return the singleton only when process telemetry is not fully disabled.""" + with cls._lock: + if cls._process_disabled or is_telemetry_disabled_by_environment(): + cls._process_disabled = True + if cls._instance is not None: + cls._instance.disable_telemetry() + return None + return cls() - """ - telemetry_cache_dir = None - if "OLIVE_TELEMETRY_CACHE_DIR" in os.environ: - telemetry_cache_dir = os.environ["OLIVE_TELEMETRY_CACHE_DIR"] - if not telemetry_cache_dir: - telemetry_cache_dir = get_telemetry_base_dir() / "cache" - return telemetry_cache_dir / self._cache_file_name - - def _write_payload_to_cache(self, payload: bytes) -> None: - """Write failed telemetry payload to cache for later retry. - - Design decisions: - - Parse payload to extract individual events (allows filtering) - - Filter to only critical events near size limit (preserves important data) - - Use file locking for multi-process safety (prevents corruption) - - Use exponential backoff for file contention (avoids spinning) - - Fail silently on errors (telemetry should never crash app) - - Assumptions: - - JSON operations are fast enough for synchronous execution - - File contention is rare and transient (retry a few times) - - Cache size limits prevent unbounded growth - - Critical events (heartbeat) are more important than others - """ - try: - cache_path = self.cache_path - if cache_path is None: - return + def __new__(cls): + """Create or return the singleton instance.""" + with cls._lock: + if cls._process_disabled or is_telemetry_disabled_by_environment(): + cls._process_disabled = True + if cls._instance is not None: + cls._instance.disable_telemetry() + return cls._instance + return cls._new_unpublished_instance(disabled=True, initialized=True) + if cls._instance is None: + cls._instance = cls._new_unpublished_instance(disabled=False, initialized=False) + return cls._instance + + @classmethod + def _new_unpublished_instance(cls, *, disabled: bool, initialized: bool) -> "Telemetry": + instance = super().__new__(cls) + instance._initialized = initialized + instance._disabled = disabled + instance._store = None + instance._uploader = None + instance._enabled = not disabled + instance._recipe_only_ci_telemetry = False + instance._global_metadata = {} + instance._instrumentation_key = "" + instance._envelope_ikey = "" + instance._app_session_guid = "" + return instance - # Parse payload into individual events for filtering - entries = _parse_payload(payload) - if not entries: + def __init__(self): + """Initialize the telemetry store and uploader (runs once).""" + with self._lock: + if self._initialized: return + # Mark initialized under the lock before doing any work, so two + # threads whose first Telemetry() calls interleave cannot both run + # the body (which would create two uploaders and two heartbeats). + self._initialized = True - cache_path.parent.mkdir(parents=True, exist_ok=True) - - max_retries = 3 - for attempt in range(max_retries + 1): - try: - cache_size = cache_path.stat().st_size if cache_path.exists() else 0 - - # Hard limit: stop caching entirely to prevent unbounded growth - if cache_size >= HARD_MAX_CACHE_SIZE_BYTES: - return - - # Soft limit: keep only critical events to preserve space - if cache_size >= MAX_CACHE_SIZE_BYTES: - entries = [entry for entry in entries if entry["event_name"] in CRITICAL_EVENTS] - if not entries: - return - - # Append base64-encoded newline-delimited entries - # Use exclusive file lock for multi-process safety - with _exclusive_file_lock(cache_path, mode="a") as cache_file: - for entry in entries: - plain = json.dumps(entry, ensure_ascii=False, separators=(",", ":")) - cache_file.write(_encode_cache_line(plain) + "\n") - return - except OSError as exc: - # Retry only on transient access errors (file locked by another process) - if exc.errno not in {errno.EACCES, errno.EAGAIN, errno.EWOULDBLOCK, errno.EBUSY}: - return - if attempt >= max_retries: - return - # Exponential backoff: 50ms, 100ms, 200ms (aligned with C# implementation) - time.sleep(0.05 * (2**attempt)) - except Exception: - # Fail silently - telemetry errors should not crash the application - return - - def _flush_cache(self) -> None: - """Flush this process's cached events back to telemetry service.""" - cache_path = self.cache_path - if cache_path is None or not cache_path.exists(): - return - - self._flush_cache_file(cache_path) - - def _flush_cache_file(self, cache_path: Path) -> None: - """Flush cached events back to telemetry service. - - Approach: - 1. Atomically rename cache → .flush (claims ownership, prevents concurrent flushes) - 2. Read all events from .flush file - 3. Queue all events for sending via telemetry logger - 4. Force flush with 2-second timeout - 5. On success: delete .flush file - 6. On failure: restore .flush → cache for retry - - Multi-process coordination: - - `replace()` is atomic; only one process can successfully rename the cache file - - If another process already renamed it, we get FileNotFoundError and abort - - Stale .flush files from crashes are overwritten by the atomic rename - - Shutdown handling: - - If shutdown flag set during flush, restore cache before returning - - This preserves events even if callbacks don't fire during shutdown - - Callback behavior: - - Queued events trigger callbacks with success/failure - - Failed events are automatically re-cached via callbacks (unless shutting down) - - The _is_flushing flag prevents re-caching of replayed events during flush - """ - flush_path = None - try: - # Check shutdown before starting (under lock to prevent race) - with self._lock: - if self._shutdown: - return - - if not cache_path.exists(): - return + self._enabled = True + self._recipe_only_ci_telemetry = False - # Atomically rename to .flush file to claim ownership - # Overwrite any stale .flush file from crashed process (C# pattern) - flush_path = cache_path.with_name(f"{cache_path.name}.flush") - try: - # On Windows/POSIX, replace() overwrites existing files atomically - cache_path.replace(flush_path) - except FileNotFoundError: - # Cache already claimed by another flush or doesn't exist + if self._disabled or is_telemetry_disabled_by_environment(): + type(self)._process_disabled = True + self._disabled = True + self._enabled = False return - # Read all cached entries (base64-decoded) - entries = _read_cache_entries(flush_path) + self._recipe_only_ci_telemetry = is_ci_environment() + self._app_session_guid = str(uuid.uuid4()) - if not entries: - # Empty cache, just delete the flush file - flush_path.unlink(missing_ok=True) - return + try: + options = OneCollectorExporterOptions( + connection_string=base64.b64decode( + "SW5zdHJ1bWVudGF0aW9uS2V5PTYyMTUwOTExZGMwMDRmYzliYjY3YmE5NjA2NDI3ZTU2LWVjNjFmOWFmLTVkN2EtNGQxOS1hZjMxLWI5Y2Q2OWU5ODdmMS02OTE1" + ).decode() + ) + options.validate() + self._instrumentation_key = options.instrumentation_key + self._envelope_ikey = ( + f"{CommonSchemaJsonSerializationHelper.ONE_COLLECTOR_TENANCY_SYMBOL}:{options.tenant_token}" + ) - # Replay all events through telemetry logger - # Note: _is_flushing flag (set by caller) prevents these callbacks from re-caching or triggering nested flushes - # (unlikely since we just successfully sent an event, indicating network is available) - for entry in entries: - try: - event_name = entry["event_name"] - event_data = entry["event_data"] - if not event_name or not event_data: - continue - attributes = json.loads(event_data) - if not isinstance(attributes, dict): - continue - # Preserve original timestamp - attributes["initTs"] = entry.get("initTs", entry["ts"]) - self._telemetry.log(event_name, attributes, None) - except Exception: - # Skip malformed entries - continue - - # Check if shutdown happened during flush - with self._lock: - if self._shutdown: - # Restore cache to avoid data loss during shutdown - if flush_path and flush_path.exists(): - try: - cache_path.parent.mkdir(parents=True, exist_ok=True) - flush_path.replace(cache_path) - except Exception: - # Silently ignore errors during cleanup - pass + # Durable on-disk queue + background uploader. The uploader + # retries detailed events until delivery. CI has a separate + # recipe-only queue so it cannot drain local action/error rows. + db_file_name = CI_DB_FILE_NAME if self._recipe_only_ci_telemetry else DB_FILE_NAME + db_path = os.path.join(get_telemetry_base_dir(), db_file_name) + self._store = OfflineEventStore(db_path) + if not self._store.is_open: + self._store = None + self._enabled = False + self._initialized = False return - - # Wait for in-flight callbacks to complete before deciding success/failure - flush_success = self.wait_for_callbacks(timeout_sec=5.0, during_flush=True) - if flush_success: - # Success: delete the flush file (events were sent) - if flush_path: - flush_path.unlink(missing_ok=True) - elif flush_path and flush_path.exists(): - # Failure: restore cache for retry later - cache_path.parent.mkdir(parents=True, exist_ok=True) - flush_path.replace(cache_path) - except Exception: - # Best-effort restore on any exception to prevent data loss - if flush_path and flush_path.exists(): - try: - cache_path.parent.mkdir(parents=True, exist_ok=True) - flush_path.replace(cache_path) - except Exception: - # If restore fails, we lose the data (acceptable for telemetry) - pass - return - - @property - def is_flushing(self) -> bool: - with self._lock: - return self._is_flushing - - -class Telemetry: - """Wrapper that wires environment configuration into the library logger. - - This is a singleton class - all instances share the same state. - Use Telemetry() to get the singleton instance. - """ - - _instance: Optional["Telemetry"] = None - _lock = threading.Lock() - - def __new__(cls): - """Create or return the singleton instance. - - Thread-safe singleton implementation using double-checked locking. - """ - if cls._instance is None: - with cls._lock: - # Double-check pattern to prevent race conditions - if cls._instance is None: - instance = super().__new__(cls) - instance._initialized = False - cls._instance = instance - return cls._instance - - def __init__(self): - """Initialize the telemetry logger (only runs once for singleton).""" - # Prevent re-initialization - if self._initialized: + self._uploader = EventUploader(self._store, instrumentation_key=self._instrumentation_key) + if not self._recipe_only_ci_telemetry: + self._enqueue_heartbeat_once() + self._uploader.start() + except Exception: + # Fail silently — telemetry must never crash the host application + if self._store is not None: + self._store.close() + self._store = None + self._uploader = None + self._enabled = False + self._initialized = False + + def _enqueue_heartbeat_once(self) -> None: + """Reserve, enrich, and release this process's durable Heartbeat.""" + if type(self)._heartbeat_enqueued or self._store is None: return - - self._logger = None - self._cache_handler = None - try: - self._logger = self._create_logger() - event_source.disable() - - self._cache_handler = TelemetryCacheHandler(self) - self._setup_payload_callbacks() - if self._is_ci_environment(): - self.disable_telemetry() - self._initialized = True + device_id, device_id_status = get_hashed_device_id_and_status() + minimal_payload = self._build_payload( + HEARTBEAT_EVENT_NAME, + { + "device_id": device_id, + "device_id_status": device_id_status.value, + }, + ) + if minimal_payload is None: return - self._log_heartbeat() - if os.environ.get("OLIVE_DISABLE_TELEMETRY") == "1": - self.disable_telemetry() - self._initialized = True - except Exception: - # Fail silently — telemetry must never crash the host application - self._initialized = True - - @staticmethod - def _is_ci_environment() -> bool: - """Detect CI/CD environments by checking well-known environment variables.""" - return any(os.environ.get(var) for var in _CI_ENV_VARS) - - def _create_logger(self) -> Optional[TelemetryLogger]: - try: - return get_telemetry_logger(base64.b64decode(CONNECTION_STRING).decode()) + row_id = self._store.reserve(minimal_payload, _HEARTBEAT_RELEASE_SECONDS) + if row_id is None: + return + type(self)._heartbeat_enqueued = True + try: + full_payload = self._build_payload( + HEARTBEAT_EVENT_NAME, + { + "device_id": device_id, + "device_id_status": device_id_status.value, + "os": platform.system(), + "os_version": platform.version(), + "os_release": platform.release(), + "os_arch": platform.machine(), + }, + ) + if full_payload is not None and self._store.release(row_id, full_payload): + return + except Exception: + pass + # If enrichment fails, release the already-durable minimal event. + self._store.release(row_id) except Exception: - return None - - def _setup_payload_callbacks(self) -> None: - # Register callback for payload transmission events - # No need to store unregister function - logger shutdown will clean up callbacks - if self._logger is None: - return - self._logger.register_payload_transmitted_callback( - self._cache_handler.on_payload_transmitted, - include_failures=True, - ) + pass def add_global_metadata(self, metadata: dict[str, Any]) -> None: - """Add metadata to all telemetry events. - - Args: - metadata: Dictionary of metadata key-value pairs to add to all events. - These will be included in every telemetry event sent. - - Example: - >>> telemetry = Telemetry() - >>> telemetry.add_global_metadata({"user_id": "12345", "environment": "production"}) - - """ + """Merge metadata into every subsequent telemetry event.""" try: - if self._logger is None: - return - self._logger.add_global_metadata(metadata) + if metadata: + self._global_metadata = {**self._global_metadata, **metadata} except Exception: - # Fail silently — telemetry must never crash the host application pass + @property + def accepts_detailed_events(self) -> bool: + """Whether action and error events can currently be persisted.""" + return bool( + self._enabled and not self._recipe_only_ci_telemetry and self._store is not None and self._store.is_open + ) + def log( self, event_name: str, attributes: Optional[dict[str, Any]] = None, metadata: Optional[dict[str, Any]] = None, ) -> None: - """Log a telemetry event. - - Args: - event_name: Name of the event to log (e.g., "UserLogin", "ModelTrained"). - attributes: Optional dictionary of event-specific attributes. - metadata: Optional dictionary of additional metadata to merge with attributes. - - Example: - >>> telemetry = Telemetry() - >>> telemetry.log("ModelOptimized", {"model_type": "bert", "duration_ms": 1500}) - - """ + """Log a telemetry event (persisted durably, uploaded in the background).""" try: - attrs = _merge_metadata(attributes, metadata) - if self._logger is None: - return - self._logger.log(event_name, attrs) - if self._cache_handler: - self._cache_handler.record_event_logged() + with self._lock: + if not self._enabled or self._store is None: + return + if self._recipe_only_ci_telemetry and event_name != RECIPE_EVENT_NAME: + return + payload = self._build_payload(event_name, attributes, metadata) + if payload is None: + return + self._store.store(payload) + if self._uploader is not None: + self._uploader.request_drain() except Exception: # Fail silently — telemetry must never crash the host application pass - def _log_heartbeat( + def _build_payload( self, + event_name: str, + attributes: Optional[dict[str, Any]], metadata: Optional[dict[str, Any]] = None, - ) -> None: - """Log a heartbeat event with system information. - - Args: - metadata: Optional additional metadata to include. + ) -> Optional[bytes]: + """Merge metadata, filter to whitelisted keys, and serialize one event. + Returns the Common Schema JSON bytes, or None if the event is not + whitelisted or filters to nothing. """ - try: - encrypted_device_id, device_id_status = get_encrypted_device_id_and_status() - attributes = { - "device_id": encrypted_device_id, - "id_status": device_id_status.value, - "os": { - "name": platform.system().lower(), - "version": platform.version(), - "release": platform.release(), - "arch": platform.machine(), - }, + attrs = _merge_metadata(attributes, metadata) + if self._global_metadata: + attrs = {**self._global_metadata, **attrs} + filtered = _filter_event_data(event_name, attrs) + if filtered is None or not filtered: + # Unknown/empty event: not whitelisted. + return None + event_data = dict(filtered) + event_data.update( + { + "appName": "Olive", + "LibraryVersion": VERSION, + "AppSessionGuid": self._app_session_guid, } - self.log(HEARTBEAT_EVENT_NAME, attributes, metadata) - except Exception: - # Fail silently — telemetry must never crash the host application - pass + ) + serialized_snapshots = {} + for snapshot_field in ("configOverrides", "packageConfigOverrides"): + snapshot = event_data.get(snapshot_field) + if not isinstance(snapshot, str): + continue + event_data.pop(snapshot_field) + try: + parsed_snapshot = json.loads(snapshot) + except (TypeError, ValueError): + continue + scrubbed_snapshot = scrub_config_snapshot_for_telemetry(parsed_snapshot) + serialized_snapshot = json.dumps( + scrubbed_snapshot, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + if len(serialized_snapshot.encode("utf-8")) > MAX_TELEMETRY_STRING_LENGTH: + serialized_snapshot = '{"truncated":"[truncated]"}' + serialized_snapshots[snapshot_field] = serialized_snapshot + event_data = { + key: value + for key, value in event_data.items() + if value is not None and isinstance(value, (str, bytes, bytearray, bool, int, float, datetime, os.PathLike)) + } + scrubbed = scrub_value_for_telemetry(event_data) + if not isinstance(scrubbed, dict): + return None + scrubbed.update(serialized_snapshots) + exception_message = event_data.get("exceptionMessage") + if event_name == ERROR_EVENT_NAME and isinstance(exception_message, str): + scrubbed["exceptionMessage"] = scrub_error_message_for_telemetry(exception_message) + envelope = CommonSchemaJsonSerializationHelper.create_event_envelope( + event_name=event_name, + timestamp=datetime.now(timezone.utc), + ikey=self._envelope_ikey, + data=scrubbed, + ) + return CommonSchemaJsonSerializationHelper.serialize_to_json_bytes(envelope) def disable_telemetry(self) -> None: - """Disable all telemetry logging. + """Fully disable telemetry for the remainder of this process.""" + with self._lock: + type(self)._process_disabled = True + self._disabled = True + self._enabled = False + uploader_stopped = True + if self._uploader is not None: + self._uploader.retain_queued_rows() + uploader_stopped = self._uploader.stop_loop(0) + if self._uploader is not None and uploader_stopped: + self._uploader.close() + self._uploader = None + if self._uploader is None and self._store is not None: + self._store.close() + self._store = None + + @classmethod + def disable_process_telemetry(cls) -> None: + """Latch full suppression without constructing telemetry resources.""" + with cls._lock: + cls._process_disabled = True + if cls._instance is not None: + cls._instance.disable_telemetry() + + @classmethod + def is_process_telemetry_disabled(cls) -> bool: + """Return whether environment or API state fully disables this process.""" + with cls._lock: + return cls._process_disabled or is_telemetry_disabled_by_environment() + + def shutdown( + self, + timeout_millis: float = 10_000, + callback_timeout_millis: float = 2_000, + flush_seconds: float = 0, + ) -> None: + """Stop the background uploader with bounded cleanup. - After calling this method, no telemetry events will be sent until - telemetry is explicitly re-enabled. + Delivery does not depend on a flush here: durability guarantees that any + undelivered events remain in the on-disk store and are uploaded on the + next run (or by a concurrently-running process). Synchronous network I/O + occurs only when a caller explicitly supplies ``flush_seconds`` (used by + ephemeral Docker runners). """ try: - if self._logger is None: - return - self._logger.disable_telemetry() + timeout_seconds = max(0.0, timeout_millis / 1000.0) + callback_timeout_seconds = max(0.0, callback_timeout_millis / 1000.0) + flush_seconds = max(0.0, flush_seconds) + disabled = bool(getattr(self, "_disabled", False)) + if not disabled and bool(getattr(self, "_recipe_only_ci_telemetry", False)): + flush_seconds = max(flush_seconds, callback_timeout_seconds) + deadline = time.monotonic() + max(timeout_seconds, callback_timeout_seconds, flush_seconds) + + def remaining_seconds() -> float: + return max(0.0, deadline - time.monotonic()) + + uploader_stopped = True + if self._uploader is not None: + uploader_stopped = self._uploader.stop_loop( + join_timeout_seconds=0 if disabled else min(timeout_seconds, remaining_seconds()) + ) + if uploader_stopped: + if flush_seconds > 0 and not disabled: + flush_timeout = min(flush_seconds, remaining_seconds()) + if flush_timeout > 0: + self._uploader.flush(flush_timeout) + self._uploader.close() + self._uploader = None + if self._store is not None and uploader_stopped: + self._store.close() + self._store = None + if self._uploader is None and self._store is None: + self._initialized = False except Exception: # Fail silently — telemetry must never crash the host application pass - def shutdown(self, timeout_millis: float = 10_000, callback_timeout_millis: float = 2_000) -> None: - """Shutdown telemetry and flush pending events. - - Shutdown sequence: - 1. Wait for in-flight flush to complete (up to 1 second) - 2. Wait for callbacks + signal shutdown to cache handler - 3. Shutdown logger (cleans up callbacks automatically) - """ + def __del__(self): + """Safety-net cleanup on garbage collection.""" try: - # Step 1: Wait for pending flush to complete (matches C# 1-second timeout) - start_time = time.time() - while time.time() - start_time < 1.0: - if not self._cache_handler or not self._cache_handler.is_flushing: - break - time.sleep(0.05) - - # Step 2: Wait for callbacks/flush to complete before shutting down cache handler - if self._cache_handler: - # Nothing can be done if callbacks don't complete in time, so we ignore the result - _ = self._cache_handler.wait_for_callbacks(callback_timeout_millis / 1000) - self._cache_handler.shutdown() - - # Step 3: Shutdown logger (callbacks cleaned up automatically) - if self._logger is not None: - self._logger.shutdown() + self.shutdown(timeout_millis=0, callback_timeout_millis=0, flush_seconds=0) except Exception: - # Fail silently — telemetry must never crash the host application pass - def __del__(self): - """Cleanup telemetry resources on garbage collection. - This is a safety net to ensure resources are cleaned up even if - shutdown() is not explicitly called. However, relying on __del__ - is not recommended - always call shutdown() explicitly when done. - """ - try: - self.shutdown() - except Exception: - # Silently ignore errors during cleanup - pass +def _get_logger() -> Optional[Telemetry]: + """Get or create telemetry without publishing a disabled singleton.""" + return Telemetry.get_or_create_if_enabled() -def _get_logger() -> Telemetry: - """Get or create the singleton Telemetry instance.""" - return Telemetry() +def disable_telemetry() -> None: + """Fully disable telemetry for the remainder of this process.""" + Telemetry.disable_process_telemetry() def _merge_metadata(attributes: Optional[dict[str, Any]], metadata: Optional[dict[str, Any]]) -> dict[str, Any]: - merged = dict(attributes or {}) - if metadata: - merged.update(metadata) + merged = dict(metadata or {}) + if attributes: + merged.update(attributes) return merged -def _parse_payload(payload: bytes) -> list[dict[str, Any]]: - """Parse telemetry payload into individual event entries. - - Design decisions: - - Filter events to only allowed keys (privacy/security) - - Store as minimal JSON (reduces cache size) - - Fail silently on malformed data (telemetry should be robust) - - Assumptions: - - Payload is newline-delimited JSON (OneCollector format) - - Events have "name", "time", and "data" fields - - Only whitelisted events and fields should be cached - """ - entries = [] - try: - payload_text = payload.decode("utf-8") - lines = payload_text.splitlines() - - for raw_line in lines: - line = raw_line.strip() - if not line: - continue - try: - event = json.loads(line) - event_name = event["name"] - if not event_name: - continue - # Filter to only allowed keys for privacy/security - filtered_data = _filter_event_data(event_name, event["data"]) - if not filtered_data: - continue - entries.append( - { - "ts": event["time"] or time.time(), - "event_name": event_name, - # Compact JSON to reduce cache size - "event_data": json.dumps(filtered_data, ensure_ascii=False, separators=(",", ":")), - } - ) - except Exception: - # Skip malformed lines - continue - except Exception: - # If entire payload is malformed, return empty list - return [] - - return entries - - def _filter_event_data(event_name: str, data: dict[str, Any]) -> Optional[dict[str, Any]]: """Filter event data to only allowed keys for privacy/security. - Design decisions: - - Whitelist approach (only explicitly allowed keys are included) - - Support nested keys with dot notation (e.g., "os.name") - - Return None if no allowed keys found (filters out unknown events) - - Assumptions: - - ALLOWED_KEYS dict defines all cacheable events and their fields - - Unknown events should not be cached (privacy/security) + Whitelist approach: only explicitly allowed keys (with dot-notation support + for nested values, e.g. "os.name") are kept. Returns None for unknown events + so they are neither persisted nor sent. """ if event_name not in ALLOWED_KEYS: return None @@ -721,7 +559,7 @@ def _filter_event_data(event_name: str, data: dict[str, Any]) -> Optional[dict[s value = _get_nested_value(data, key) if value is None: continue - _set_nested_value(filtered, key, value) + _set_nested_value(filtered, FIELD_NAMES.get(key, key), value) return filtered or None @@ -740,37 +578,3 @@ def _set_nested_value(data: dict[str, Any], key: str, value: Any) -> None: for part in parts[:-1]: current = current.setdefault(part, {}) current[parts[-1]] = value - - -def _read_cache_entries(cache_path: Path) -> list[dict[str, Any]]: - """Read all entries from a cache file, decoding each line. - - Design decisions: - - Use file locking for multi-process safety - - Continue reading past malformed entries (partial data recovery) - - Return empty list on complete read failure (fail gracefully) - - Each line is base64-decoded before JSON parsing. - - Assumptions: - - Cache file contains newline-delimited base64-encoded entries (one per line) - - Each line is independent (one malformed line doesn't affect others) - - Empty or whitespace-only lines are skipped - """ - entries = [] - try: - with _exclusive_file_lock(cache_path, mode="r") as cache_file: - for raw_line in cache_file: - line = raw_line.strip() - if not line: - continue - try: - line = json.loads(_decode_cache_line(line)) - if isinstance(line, dict): - entries.append(line) - except Exception: - # Malformed line, skip and continue - continue - except Exception: - # If file cannot be opened or read, return empty list - return [] - return entries diff --git a/olive/telemetry/telemetry_extensions.py b/olive/telemetry/telemetry_extensions.py index e5b13395d0..493fc07aba 100644 --- a/olive/telemetry/telemetry_extensions.py +++ b/olive/telemetry/telemetry_extensions.py @@ -6,13 +6,16 @@ import functools import inspect import time +import traceback +from contextlib import suppress from types import TracebackType -from typing import Any, Callable, Optional, TypeVar +from typing import Any, Callable, Optional, TypeVar, cast -from olive.telemetry.telemetry import ACTION_EVENT_NAME, ERROR_EVENT_NAME, _get_logger -from olive.telemetry.utils import _format_exception_message +from olive.telemetry.telemetry import ACTION_EVENT_NAME, ERROR_EVENT_NAME, RECIPE_EVENT_NAME, _get_logger +from olive.telemetry.telemetry_redaction import scrub_error_message_for_telemetry _TFunc = TypeVar("_TFunc", bound=Callable[..., Any]) +_ERROR_LOGGED_ATTR = "_olive_telemetry_logged" def log_action( @@ -22,27 +25,80 @@ def log_action( success: bool, metadata: Optional[dict[str, Any]] = None, ) -> None: - telemetry = _get_logger() - attributes = { - "invoked_from": invoked_from, - "action_name": action_name, - "duration_ms": duration_ms, - "success": success, - } - telemetry.log(ACTION_EVENT_NAME, attributes, metadata) + with suppress(Exception): + telemetry = _get_logger() + if telemetry is None: + return + attributes = { + "invoked_from": invoked_from, + "action_name": action_name, + "duration_ms": duration_ms, + "success": success, + } + telemetry.log(ACTION_EVENT_NAME, attributes, metadata) def log_error( exception_type: str, exception_message: str, metadata: Optional[dict[str, Any]] = None, +) -> None: + with suppress(Exception): + telemetry = _get_logger() + if telemetry is None: + return + attributes = { + "exception_type": exception_type, + "exception_message": _redact_error_message(exception_message), + } + telemetry.log(ERROR_EVENT_NAME, attributes, metadata) + + +def log_recipe_result( + recipe_name: str, + success: bool, + metadata: Optional[dict[str, Any]] = None, ) -> None: telemetry = _get_logger() + if telemetry is None: + return attributes = { - "exception_type": exception_type, - "exception_message": exception_message, + "recipe_name": recipe_name, + "success": success, } - telemetry.log(ERROR_EVENT_NAME, attributes, metadata) + telemetry.log(RECIPE_EVENT_NAME, attributes, metadata) + + +def _redact_error_message(text: str) -> str: + return scrub_error_message_for_telemetry(text) + + +def _is_exception_logged(exc: BaseException) -> bool: + return bool(getattr(exc, _ERROR_LOGGED_ATTR, False)) + + +def _mark_exception_logged(exc: BaseException) -> None: + try: + setattr(exc, _ERROR_LOGGED_ATTR, True) + except Exception: + # Some exception implementations do not allow custom attributes. + pass + + +def _format_exception_message(ex: BaseException, tb: Optional[TracebackType] = None) -> str: + """Format exception and frame metadata without collecting source-code lines.""" + lines = [] + if tb is not None: + lines.append("Traceback (most recent call last):") + for frame in traceback.extract_tb(tb, limit=5): + frame_name = _redact_error_message(frame.name) + lines.append(f'File "[path]", line {frame.lineno}, in {frame_name}') + try: + exception_message = str(ex) + except Exception: + exception_message = "" + lines.append(f"{type(ex).__name__}: {_redact_error_message(exception_message)}") + return _redact_error_message("\n".join(lines)) def _resolve_invoked_from(skip_frames: int = 0) -> str: @@ -69,6 +125,18 @@ def _resolve_invoked_from(skip_frames: int = 0) -> str: return "Interactive" +def _resolve_action_name(func: Callable[..., Any]) -> str: + action_name = getattr(func, "__name__", "unknown") + qualname = getattr(func, "__qualname__", action_name) + if "." not in qualname: + return action_name + owner = qualname.rsplit(".", 1)[0].rsplit(".", 1)[-1] + if owner == "": + return action_name + owner = owner[: -len("Command")] if owner.endswith("Command") else owner + return owner if action_name == "run" else f"{owner}.{action_name}" + + class ActionContext: """Context manager for recording telemetry around a block of work.""" @@ -79,7 +147,18 @@ def __init__( metadata: Optional[dict[str, Any]] = None, ): self.action_name = action_name - self.invoked_from = invoked_from if invoked_from is not None else _resolve_invoked_from() + try: + telemetry = _get_logger() + self._telemetry_enabled = bool(telemetry is not None and telemetry.accepts_detailed_events) + except Exception: + self._telemetry_enabled = False + self.invoked_from = ( + invoked_from + if invoked_from is not None + else _resolve_invoked_from() + if self._telemetry_enabled + else "disabled" + ) self.metadata = metadata or {} self._start_time: Optional[float] = None @@ -96,7 +175,11 @@ def __exit__( exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> bool: - duration_ms = int((time.perf_counter() - (self._start_time or time.perf_counter())) * 1000) + if not self._telemetry_enabled: + return False + end_time = time.perf_counter() + start_time = self._start_time if self._start_time is not None else end_time + duration_ms = int((end_time - start_time) * 1000) success = exc_type is None log_action( @@ -107,12 +190,13 @@ def __exit__( metadata=self.metadata, ) - if exc_type is not None and exc_val is not None: + if exc_type is not None and exc_val is not None and not _is_exception_logged(exc_val): log_error( exception_type=exc_type.__name__, exception_message=_format_exception_message(exc_val, exc_tb), metadata=self.metadata, ) + _mark_exception_logged(exc_val) # Do not suppress exceptions return False @@ -123,13 +207,21 @@ def action(func: _TFunc) -> _TFunc: @functools.wraps(func) def wrapper(*args: Any, **kwargs: Any): - invoked_from = _resolve_invoked_from() - action_name = func.__name__ - if args and hasattr(args[0], "__class__"): - cls_name = args[0].__class__.__name__ - cls_name = cls_name[: -len("Command")] if cls_name.endswith("Command") else cls_name - if cls_name: - action_name = cls_name if action_name == "run" else f"{cls_name}.{action_name}" + try: + telemetry = _get_logger() + if telemetry is None or not telemetry.accepts_detailed_events: + return func(*args, **kwargs) + except Exception: + return func(*args, **kwargs) + + # Resolve telemetry context defensively: instrumentation (including + # inspect.stack()) must never propagate into the wrapped call. + try: + invoked_from = _resolve_invoked_from() + action_name = _resolve_action_name(func) + except Exception: + invoked_from = "unknown" + action_name = getattr(func, "__name__", "unknown") start_time = time.perf_counter() success = True @@ -137,10 +229,12 @@ def wrapper(*args: Any, **kwargs: Any): return func(*args, **kwargs) except Exception as exc: success = False - log_error( - exception_type=type(exc).__name__, - exception_message=_format_exception_message(exc, exc.__traceback__), - ) + if not _is_exception_logged(exc): + log_error( + exception_type=type(exc).__name__, + exception_message=_format_exception_message(exc, exc.__traceback__), + ) + _mark_exception_logged(exc) raise finally: duration_ms = int((time.perf_counter() - start_time) * 1000) @@ -151,4 +245,4 @@ def wrapper(*args: Any, **kwargs: Any): success=success, ) - return wrapper # type: ignore[return-value] + return cast("_TFunc", wrapper) diff --git a/olive/telemetry/telemetry_redaction.py b/olive/telemetry/telemetry_redaction.py new file mode 100644 index 0000000000..c2f1cd4341 --- /dev/null +++ b/olive/telemetry/telemetry_redaction.py @@ -0,0 +1,434 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Free-text telemetry redaction.""" + +import os +import re +from collections.abc import Mapping +from collections.abc import Set as AbstractSet +from datetime import date, datetime, time, timedelta +from uuid import UUID + +MAX_TELEMETRY_STRING_LENGTH = 40_960 +MAX_ERROR_MESSAGE_LENGTH = MAX_TELEMETRY_STRING_LENGTH + +_SENSITIVE_COMPACT_KEYS = { + "accountkey", + "accesskey", + "accesstoken", + "apikey", + "auth", + "authorization", + "authtoken", + "clientsecret", + "connectionstring", + "credential", + "credentials", + "passwd", + "password", + "privatekey", + "pwd", + "sastoken", + "secret", + "secretkey", + "servicecredential", + "sig", + "signature", + "subscriptionkey", + "token", +} +_ENVIRONMENT_COMPACT_KEYS = {"env", "environment", "environmentvariables", "environmentvars", "envvariables", "envvars"} +_PATH_KEY_SUFFIXES = ("dir", "dirs", "file", "files", "path", "paths") + + +def normalize_config_key_for_telemetry(key) -> str: + if key is None: + return "" + value = str(key) + value = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", value) + value = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", value) + return re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_") + + +def is_sensitive_config_key_for_telemetry(key) -> bool: + normalized = normalize_config_key_for_telemetry(key) + compact = normalized.replace("_", "") + parts = set(normalized.split("_")) + if compact in _SENSITIVE_COMPACT_KEYS: + return True + if parts & {"auth", "authorization", "credential", "credentials", "passwd", "password", "pwd", "secret", "token"}: + return True + if "key" in parts and parts & {"access", "account", "api", "private", "secret", "subscription"}: + return True + if {"connection", "string"} <= parts: + return True + if compact.endswith( + ("authorization", "credential", "credentials", "passwd", "password", "pwd", "secret", "signature", "token") + ): + return True + return any( + marker in compact + for marker in ( + "accesskey", + "accountkey", + "apikey", + "connectionstring", + "privatekey", + "secretkey", + "subscriptionkey", + ) + ) + + +def is_environment_config_key_for_telemetry(key) -> bool: + normalized = normalize_config_key_for_telemetry(key) + return normalized.replace("_", "") in _ENVIRONMENT_COMPACT_KEYS or bool( + set(normalized.split("_")) & {"env", "environment"} + ) + + +def is_path_like_config_key_for_telemetry(key) -> bool: + compact = normalize_config_key_for_telemetry(key).replace("_", "") + return bool(compact) and compact.endswith(_PATH_KEY_SUFFIXES) + + +def scrub_config_snapshot_for_telemetry(value, key=None): + """Recursively scrub a JSON-compatible configuration snapshot by key and value.""" + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return "" + if is_path_like_config_key_for_telemetry(key): + return "" + if isinstance(value, Mapping): + items = {} + collisions = set() + for child_key, child_value in value.items(): + safe_key = scrub_string_for_telemetry(str(child_key)) + if not safe_key: + continue + if safe_key in items: + collisions.add(safe_key) + else: + items[safe_key] = scrub_config_snapshot_for_telemetry(child_value, child_key) + return {safe_key: items[safe_key] for safe_key in sorted(items) if safe_key not in collisions} + if isinstance(value, (list, tuple)): + return [scrub_config_snapshot_for_telemetry(item, key) for item in value] + if isinstance(value, str): + return scrub_string_for_telemetry(value) + return value + + +def _token_start(value: str, index: int) -> int: + while index > 0 and not value[index - 1].isspace() and value[index - 1] not in "\"'": + index -= 1 + return index + + +def _is_drive_path_anchor(value: str, index: int) -> bool: + char = value[index] + if not char.isascii() or not char.isalpha(): + return False + if index > 0 and value[index - 1] not in "\"' \t=([{,;": + return False + if index + 2 >= len(value): + return False + return value[index + 1] == ":" and value[index + 2] in "/\\" + + +def _is_sensitive_slash_option(value: str, index: int) -> bool: + key_start = index + 1 + if key_start >= len(value) or not value[key_start].isascii() or not value[key_start].isalpha(): + return False + key_end = key_start + while key_end < len(value) and value[key_end].isascii() and (value[key_end].isalnum() or value[key_end] in "_.-"): + key_end += 1 + if key_end == key_start or not is_sensitive_config_key_for_telemetry(value[key_start:key_end]): + return False + separator = key_end + while separator < len(value) and value[separator].isspace(): + separator += 1 + return separator < len(value) and (value[separator] in "=:" or separator > key_end) + + +def _find_path_anchor(value: str): + index = 0 + slash_token_end = 0 + slash_token_start = 0 + slash_token_analyzed = False + relative_slash_anchor = None + while index < len(value): + char = value[index] + if value.startswith(("./", "../", ".\\", "..\\"), index) and ( + index == 0 or value[index - 1].isspace() or value[index - 1] in "\"'=([{,;" + ): + return index + if char == ":" and index + 2 < len(value) and value[index + 1 : index + 3] == "//": + scheme_start = index + while scheme_start > 0 and ( + value[scheme_start - 1].isascii() + and (value[scheme_start - 1].isalnum() or value[scheme_start - 1] in "+-.") + ): + scheme_start -= 1 + return scheme_start + if char == "\\" and index + 1 < len(value) and value[index + 1] == "\\": + return index + if char == "~" and index + 1 < len(value) and value[index + 1] in "/\\": + return index + if _is_drive_path_anchor(value, index): + return index + if char == "\\": + if ( + index + 1 < len(value) + and value[index + 1] not in "\\\r\n \t" + and (index == 0 or value[index - 1] in "\"' \t=([{,;") + ): + return index + separators = 0 + for candidate in value[index:]: + if candidate in "\r\n": + break + if candidate == "\\": + separators += 1 + if separators >= 2: + return _token_start(value, index) + if char == "/": + if index >= slash_token_end: + slash_token_start = _token_start(value, index) + slash_token_end = index + while ( + slash_token_end < len(value) + and not value[slash_token_end].isspace() + and value[slash_token_end] not in "\"'" + ): + slash_token_end += 1 + slash_token_analyzed = False + if ( + index + 1 < len(value) + and value[index + 1] not in "/\r\n \t" + and (index == 0 or value[index - 1] in "\"' \t=([{,;") + and not _is_sensitive_slash_option(value, index) + ): + return index + if not slash_token_analyzed: + slash_token_analyzed = True + segments = 0 + cursor = index + while cursor < slash_token_end and value[cursor] == "/": + separator_end = cursor + 1 + while separator_end < slash_token_end and value[separator_end] == "/": + separator_end += 1 + cursor = separator_end + segment_start = cursor + while cursor < slash_token_end and value[cursor] not in "/\r\n \t": + cursor += 1 + if cursor == segment_start: + break + segments += 1 + if segments >= 2: + return relative_slash_anchor if relative_slash_anchor is not None else slash_token_start + if segments == 1 and slash_token_start < index: + if relative_slash_anchor is not None: + return relative_slash_anchor + token = value[slash_token_start:slash_token_end] + if token.lower() not in {"and/or", "n/a", "read/write"} and any(char.isalpha() for char in token): + relative_slash_anchor = slash_token_start + index += 1 + return None + + +def _is_secret_key_char(char: str) -> bool: + return char.isascii() and (char.isalnum() or char in "_.-") + + +def _is_secret_key_boundary(char: str) -> bool: + return char.isspace() or char in "?&#;,\"'([{/-" + + +def _find_sensitive_value_anchor(value: str): + index = 0 + while index < len(value): + char = value[index] + if not char.isascii() or not char.isalpha(): + index += 1 + continue + if index > 0 and not _is_secret_key_boundary(value[index - 1]): + index += 1 + continue + + key_end = index + 1 + while key_end < len(value) and _is_secret_key_char(value[key_end]): + key_end += 1 + if not is_sensitive_config_key_for_telemetry(value[index:key_end]): + index = key_end + continue + + separator = key_end + if separator < len(value) and value[separator] in "\"'": + separator += 1 + before_whitespace = separator + while separator < len(value) and value[separator].isspace(): + separator += 1 + assignment = separator < len(value) and value[separator] in "=:" + cli_option = index > 0 and value[index - 1] in "-/" + delimited_cli_value = False + if cli_option and not assignment: + separator = key_end + while separator < len(value) and (value[separator].isspace() or value[separator] in "\"',[](){}"): + if value[separator] in "\"',[](){}": + delimited_cli_value = True + separator += 1 + separated_cli_value = cli_option and separator > before_whitespace and separator < len(value) + separated_cli_value = separated_cli_value and (value[separator] != "-" or delimited_cli_value) + if not assignment and not separated_cli_value: + index = key_end + continue + + value_start = separator + 1 if assignment else separator + while value_start < len(value) and value[value_start].isspace(): + value_start += 1 + if value_start < len(value) and value[value_start] not in "&;\r\n": + return value_start + index = key_end + return None + + +def _is_user_info_terminator(char: str) -> bool: + return char.isspace() or char in '"\\/?#[]{}' + + +def _is_authority_terminator(char: str) -> bool: + return char.isspace() or char in "\"')},;/?#" + + +def _find_credential_url_anchor(value: str): + token_start = 0 + colon = None + index = 0 + while index < len(value): + char = value[index] + if colon is None: + if _is_user_info_terminator(char): + token_start = index + 1 + elif char == ":" and index > token_start: + colon = index + index += 1 + continue + + if _is_user_info_terminator(char): + token_start = index + 1 + colon = None + index += 1 + continue + if char != "@" or colon + 1 == index: + index += 1 + continue + + host_start = index + 1 + if host_start == len(value): + index += 1 + continue + if value[host_start] == "[": + host_end = value.find("]", host_start + 1) + if host_end > host_start + 1: + return token_start + index += 1 + continue + + host_end = host_start + while host_end < len(value) and not _is_authority_terminator(value[host_end]): + host_end += 1 + if host_end > host_start: + return token_start + index = host_end + 1 + token_start = index + colon = None + return None + + +def _truncate_utf8(value: str, max_bytes: int) -> str: + encoded = value.encode("utf-8") + if len(encoded) <= max_bytes: + return value + return encoded[:max_bytes].decode("utf-8", errors="ignore") + + +def _scrub_string_for_telemetry(value: str, max_bytes: int) -> str: + path_anchor = _find_path_anchor(value) + secret_value = _find_sensitive_value_anchor(value) + credential_anchor = _find_credential_url_anchor(value) + anchors = ( + (path_anchor, "[path]"), + (secret_value, "[secret]"), + (credential_anchor, "[secret]"), + ) + anchor, marker = min( + ((anchor, marker) for anchor, marker in anchors if anchor is not None), + default=(None, None), + ) + if anchor is not None: + return _truncate_utf8(value[:anchor], max_bytes - len(marker)) + marker + return _truncate_utf8(value, max_bytes) + + +def scrub_string_for_telemetry(value: str) -> str: + """Redact and cap a general telemetry string.""" + return _scrub_string_for_telemetry(value, MAX_TELEMETRY_STRING_LENGTH) + + +def scrub_error_message_for_telemetry(value: str) -> str: + """Redact and cap an error message at 40,960 UTF-8 bytes.""" + return _scrub_string_for_telemetry(value, MAX_ERROR_MESSAGE_LENGTH) + + +def scrub_value_for_telemetry(value, key=None): + """Recursively scrub strings and path-like values before serialization.""" + if is_sensitive_config_key_for_telemetry(key) or is_environment_config_key_for_telemetry(key): + return "" + if is_path_like_config_key_for_telemetry(key): + return "" + if isinstance(value, os.PathLike): + return "[path]" + if isinstance(value, str): + return scrub_string_for_telemetry(value) + if isinstance(value, (bytes, bytearray)): + try: + return scrub_string_for_telemetry(bytes(value).decode("utf-8")) + except UnicodeDecodeError: + return "[binary]" + if value is None or isinstance( + value, + (bool, int, float, datetime, date, time, timedelta, UUID), + ): + return value + if isinstance(value, Mapping): + items = {} + collisions = set() + for key, child in value.items(): + if isinstance(key, os.PathLike): + safe_key = "[path]" + elif isinstance(key, str): + safe_key = scrub_string_for_telemetry(key) + else: + try: + safe_key = scrub_string_for_telemetry(str(key)) + except Exception: + safe_key = f"[unsupported:{type(key).__name__}]" + if safe_key: + if safe_key in items: + collisions.add(safe_key) + else: + items[safe_key] = scrub_value_for_telemetry(child, key) + return {safe_key: items[safe_key] for safe_key in sorted(items) if safe_key not in collisions} + if isinstance(value, list): + return [scrub_value_for_telemetry(child, key) for child in value] + if isinstance(value, tuple): + return tuple(scrub_value_for_telemetry(child, key) for child in value) + if isinstance(value, AbstractSet): + children = [scrub_value_for_telemetry(child, key) for child in value] + return sorted(children, key=lambda child: (type(child).__name__, repr(child))) + try: + return scrub_string_for_telemetry(str(value)) + except Exception: + return f"[unsupported:{type(value).__name__}]" diff --git a/olive/telemetry/uploader.py b/olive/telemetry/uploader.py new file mode 100644 index 0000000000..6ee3cd7f60 --- /dev/null +++ b/olive/telemetry/uploader.py @@ -0,0 +1,334 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Background uploader that drains the SQLite offline store to OneCollector. + +Reads the oldest batch of events, POSTs them, and: +- deletes them on success (HTTP 2xx), +- deletes them on a permanent, non-retryable failure (e.g. HTTP 4xx) so a + poison event cannot block the queue forever, +- leaves them on a transient failure (network error / HTTP 5xx / timeout) to be + retried on the next cycle or the next process run. + +Durability is provided by the on-disk store, so the process can exit at any time +without losing events and without an exit-time flush. +""" + +import threading +import time +from enum import Enum +from typing import NamedTuple, Optional + +from olive.telemetry.library.options import CompressionType, OneCollectorTransportOptions +from olive.telemetry.library.payload_builder import PayloadBuilder +from olive.telemetry.library.transport import HttpJsonPostTransport +from olive.telemetry.offline_store import OfflineEventStore +from olive.telemetry.process_lock import ProcessDrainLock + + +class DrainOutcome(Enum): + EMPTY = "empty" + PROGRESS = "progress" + SPLIT = "split" + ACKNOWLEDGE_RETRY = "acknowledge_retry" + DELETE_RETRY = "delete_retry" + STORAGE_RETRY = "storage_retry" + TRANSPORT_RETRY = "transport_retry" + + +class DrainResult(NamedTuple): + delivered: int + left: int + outcome: DrainOutcome + + +class EventUploader: + """Drains the offline store and ships events over HTTP on a daemon thread. + + When several processes share one telemetry database, a single-holder + advisory lock ensures only one of them drains at a time, so the same event + is never uploaded by two concurrent drainers. Processes that do not hold the + lock keep writing events durably; the holder drains them. + """ + + def __init__( + self, + store: OfflineEventStore, + instrumentation_key: str, + endpoint: str = OneCollectorTransportOptions.DEFAULT_ENDPOINT, + compression: CompressionType = CompressionType.DEFLATE, + drain_interval_seconds: float = 2.0, + max_items_per_drain: int = 256, + send_timeout_seconds: float = 10.0, + idle_backoff_seconds: float = 30.0, + ): + self._store = store + self._drain_interval = drain_interval_seconds + self._max_items = max_items_per_drain + self._send_timeout = send_timeout_seconds + self._idle_backoff = idle_backoff_seconds + self._drain_lock = ProcessDrainLock(store.db_path + ".lock") + + self._transport = HttpJsonPostTransport( + endpoint=endpoint, + ikey=instrumentation_key, + compression=compression, + ) + + self._wake = threading.Event() + self._stop = threading.Event() + self._retain_rows = threading.Event() + self._mutation_lock = threading.Lock() + self._thread: Optional[threading.Thread] = None + self._pending_ack_ids: list[int] = [] + self._pending_delete_ids: list[int] = [] + self._split_batch_size: Optional[int] = None + + # ----- control ------------------------------------------------------- + + def start(self) -> None: + if self._thread is not None: + return + self._thread = threading.Thread(target=self._run, name="olive-telemetry-uploader", daemon=True) + self._thread.start() + + def request_drain(self) -> None: + """Nudge this process when it is the designated drainer.""" + if self._drain_lock.held: + self._wake.set() + + def stop_loop(self, join_timeout_seconds: float = 5.0) -> bool: + """Stop the background loop and wait briefly for it to exit. + + Returns True if the thread actually stopped (so a caller may safely drain + as the sole drainer), False if it is still alive (e.g. stuck in an + in-flight send) — in which case the caller must NOT drain, to avoid + double-sending the rows the thread is still processing. + """ + self._stop.set() + self._wake.set() + thread = self._thread + if thread is None: + return True + thread.join(join_timeout_seconds) + stopped = not thread.is_alive() + if stopped: + self._thread = None + return stopped + + def signal_stop(self) -> None: + """Ask the loop to stop without blocking the caller. + + The daemon thread winds down on its next wake; the drain lock is released + when it exits (or by the OS at process exit). Use this for the opt-out + path so disabling telemetry never blocks the host. + """ + self._stop.set() + self._wake.set() + + def retain_queued_rows(self) -> None: + """Stop new draining and prevent any later queue deletion.""" + with self._mutation_lock: + self._retain_rows.set() + self.signal_stop() + + def close(self) -> None: + """Release the single-drainer lock (urllib holds no persistent connection).""" + self._drain_lock.release() + + def stop(self, timeout_seconds: float = 12.0) -> None: + """Stop the loop and release the drain lock (convenience).""" + if self.stop_loop(timeout_seconds): + self.close() + + # ----- draining ------------------------------------------------------ + + def _delete_unless_retained(self, ids: list[int], deadline: Optional[float] = None) -> Optional[bool]: + with self._mutation_lock: + if self._retain_rows.is_set(): + return None + return self._store.delete(ids, deadline) + + def _finish_handled_rows(self, ids: list[int], deadline: Optional[float] = None) -> DrainResult: + self._pending_ack_ids = ids + with self._mutation_lock: + if self._retain_rows.is_set(): + return DrainResult(0, len(ids), DrainOutcome.TRANSPORT_RETRY) + if not self._store.acknowledge(ids, deadline): + return DrainResult(0, len(ids), DrainOutcome.ACKNOWLEDGE_RETRY) + self._pending_ack_ids = [] + self._pending_delete_ids = ids + deleted = self._store.delete(ids, deadline) + if not deleted: + return DrainResult(0, len(ids), DrainOutcome.DELETE_RETRY) + self._pending_delete_ids = [] + return DrainResult(len(ids), 0, DrainOutcome.PROGRESS) + + def drain_once(self, deadline: Optional[float] = None) -> DrainResult: + """Attempt to upload one batch. + + The outcome distinguishes local work that should continue immediately + (poison isolation or acknowledged deletion) from a retryable transport + failure that should back off. + """ + if self._retain_rows.is_set(): + return DrainResult( + 0, + len(self._pending_ack_ids) + len(self._pending_delete_ids), + DrainOutcome.TRANSPORT_RETRY, + ) + if self._pending_ack_ids: + return self._finish_handled_rows(self._pending_ack_ids, deadline) + if self._pending_delete_ids: + pending_ids = self._pending_delete_ids + deleted = self._delete_unless_retained(pending_ids, deadline) + if deleted is None: + return DrainResult(0, len(pending_ids), DrainOutcome.TRANSPORT_RETRY) + if not deleted: + return DrainResult(0, len(pending_ids), DrainOutcome.DELETE_RETRY) + self._pending_delete_ids = [] + return DrainResult(len(pending_ids), 0, DrainOutcome.PROGRESS) + + batch_limit = min(self._max_items, self._split_batch_size or self._max_items) + acknowledged_ids = self._store.get_acknowledged_ids(batch_limit, deadline) + if acknowledged_ids is None: + return DrainResult(0, 0, DrainOutcome.STORAGE_RETRY) + if acknowledged_ids: + self._pending_delete_ids = acknowledged_ids + deleted = self._delete_unless_retained(acknowledged_ids, deadline) + if deleted is None: + return DrainResult(0, len(acknowledged_ids), DrainOutcome.TRANSPORT_RETRY) + if not deleted: + return DrainResult(0, len(acknowledged_ids), DrainOutcome.DELETE_RETRY) + self._pending_delete_ids = [] + return DrainResult(len(acknowledged_ids), 0, DrainOutcome.PROGRESS) + + batch = self._store.get_batch_for_upload(batch_limit, deadline) + if batch is None: + return DrainResult(0, 0, DrainOutcome.STORAGE_RETRY) + if not batch: + return DrainResult(0, 0, DrainOutcome.EMPTY) + + builder = PayloadBuilder( + max_size_bytes=OneCollectorTransportOptions.DEFAULT_MAX_PAYLOAD_SIZE_BYTES, + max_items=OneCollectorTransportOptions.DEFAULT_MAX_ITEMS_PER_PAYLOAD, + ) + included: list[int] = [] + for row_id, payload in batch: + if not builder.can_add(payload): + if builder.is_empty: + return self._finish_handled_rows([row_id], deadline) + break + builder.add(payload) + included.append(row_id) + payload_bytes = builder.build() + + timeout = self._send_timeout + if deadline is not None: + timeout = min(timeout, max(0.0, deadline - time.monotonic())) + if timeout <= 0.0: + return DrainResult(0, len(included), DrainOutcome.TRANSPORT_RETRY) + admission_released = False + + def release_admission() -> None: + nonlocal admission_released + if not admission_released: + admission_released = True + self._mutation_lock.release() + + self._mutation_lock.acquire() + try: + if self._retain_rows.is_set(): + return DrainResult(0, len(included), DrainOutcome.TRANSPORT_RETRY) + success, status = self._transport.send( + payload_bytes, + timeout, + item_count=len(included), + on_send_admitted=release_admission, + ) + except Exception: + success, status = (False, None) + finally: + release_admission() + + if success: + self._split_batch_size = None + return self._finish_handled_rows(included, deadline) + if not HttpJsonPostTransport.is_retryable(status): + if status in {400, 413, 422} and len(included) > 1: + self._split_batch_size = max(1, len(included) // 2) + return DrainResult(0, len(included), DrainOutcome.SPLIT) + # Permanent rejection (e.g. 4xx): drop so it can't block the queue. + self._split_batch_size = None + return self._finish_handled_rows(included, deadline) + # Transient failure: leave the rows for the next attempt. + return DrainResult(0, len(included), DrainOutcome.TRANSPORT_RETRY) + + def flush(self, max_seconds: float = 5.0) -> None: + """Best-effort drain of all pending events, bounded by max_seconds. + + Only drains after the background loop stops and if this process can hold + the single-drainer lock; otherwise events remain durable for a later run. + """ + if self._thread is not None and self._thread.is_alive(): + return + if not self._drain_lock.acquire(): + return + try: + deadline = time.monotonic() + max(0.0, max_seconds) + delete_retries = 0 + while time.monotonic() < deadline: + result = self.drain_once(deadline) + if result.outcome is DrainOutcome.EMPTY: + return # queue empty + if result.outcome in {DrainOutcome.STORAGE_RETRY, DrainOutcome.TRANSPORT_RETRY}: + return # transient failure; leave the rest for next run + if result.outcome in {DrainOutcome.ACKNOWLEDGE_RETRY, DrainOutcome.DELETE_RETRY}: + delete_retries += 1 + if delete_retries >= 2: + return + else: + delete_retries = 0 + finally: + self._drain_lock.release() + + def _run(self) -> None: + try: + while not self._stop.is_set(): + self._wake.clear() + transient_failure = False + # Only one process drains at a time. If another holds the lock, skip + # draining this cycle; our events remain durable for the holder. + if self._drain_lock.acquire(): + try: + delete_retries = 0 + while not self._stop.is_set(): + result = self.drain_once() + if result.outcome is DrainOutcome.EMPTY: + break + if result.outcome is DrainOutcome.TRANSPORT_RETRY: + transient_failure = True + break + if result.outcome is DrainOutcome.STORAGE_RETRY: + transient_failure = True + break + if result.outcome in {DrainOutcome.ACKNOWLEDGE_RETRY, DrainOutcome.DELETE_RETRY}: + delete_retries += 1 + if delete_retries >= 2: + transient_failure = True + break + else: + delete_retries = 0 + except Exception: + transient_failure = True + else: + transient_failure = True + + wait = self._idle_backoff if transient_failure else self._drain_interval + self._wake.wait(wait) + finally: + # Release the single-drainer lock when the loop exits so another + # process can take over (also released by close()/OS on exit). + self._drain_lock.release() diff --git a/olive/telemetry/utils.py b/olive/telemetry/utils.py index 52a39acded..934288bc8a 100644 --- a/olive/telemetry/utils.py +++ b/olive/telemetry/utils.py @@ -2,15 +2,10 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -import base64 import functools import os import platform -import tempfile -import traceback from pathlib import Path -from types import TracebackType -from typing import Optional ORT_SUPPORT_DIR = r"Microsoft/DeveloperTools/.onnxruntime" @@ -18,15 +13,24 @@ def _resolve_home_dir() -> Path: """Resolve the user home directory with fallbacks for container environments.""" home = os.getenv("HOME") - if home: - return Path(home).expanduser() + if home and Path(home).is_absolute(): + return Path(home) + if platform.system() != "Windows": + try: + import pwd + + passwd_home = Path(pwd.getpwuid(os.getuid()).pw_dir) + if passwd_home.is_absolute(): + return passwd_home + except (AttributeError, ImportError, KeyError, OSError): + pass try: - return Path.home() + fallback_home = Path.home() + if fallback_home.is_absolute(): + return fallback_home except (RuntimeError, KeyError): - # /var/tmp persists across reboots unlike /tmp (FHS spec) - if platform.system() != "Windows": - return Path("/var/tmp") - return Path(tempfile.gettempdir()) + pass + raise RuntimeError("No absolute per-user telemetry storage directory is available") @functools.lru_cache(maxsize=1) @@ -34,9 +38,9 @@ def get_telemetry_base_dir() -> Path: os_name = platform.system() if os_name == "Windows": base_dir = os.environ.get("LOCALAPPDATA") or os.environ.get("APPDATA") - if not base_dir: - base_dir = str(Path.home() / "AppData" / "Local") - return Path(base_dir) / "Microsoft" / ".onnxruntime" + if base_dir and Path(base_dir).is_absolute(): + return Path(base_dir) / ORT_SUPPORT_DIR + return _resolve_home_dir() / "AppData" / "Local" / ORT_SUPPORT_DIR if os_name == "Darwin": home = _resolve_home_dir() @@ -44,102 +48,6 @@ def get_telemetry_base_dir() -> Path: # Use XDG_CACHE_HOME if set, otherwise fall back to $HOME/.cache cache_dir = os.getenv("XDG_CACHE_HOME") - if not cache_dir: - cache_dir = str(_resolve_home_dir() / ".cache") - - return Path(cache_dir).expanduser() / ORT_SUPPORT_DIR - - -def _format_exception_message(ex: BaseException, tb: Optional[TracebackType] = None) -> str: - """Format an exception and trim local paths for readability.""" - folder = "Olive" - file_line = 'File "' - formatted = traceback.format_exception(type(ex), ex, tb, limit=5) - lines = [] - for line in formatted: - line_trunc = line.strip() - if line_trunc.startswith(file_line) and folder in line_trunc: - idx = line_trunc.find(folder) - if idx != -1: - line_trunc = line_trunc[idx + len(folder) :] - elif line_trunc.startswith(file_line): - idx = line_trunc[len(file_line) :].find('"') - line_trunc = line_trunc[idx + len(file_line) :] - lines.append(line_trunc) - return "\n".join(lines) - - -class _ExclusiveFileLock: - """Cross-platform exclusive file lock context manager. - - Uses fcntl on Unix/Linux/macOS, msvcrt on Windows. - Prevents cache corruption when multiple processes access the same file. - - Design decisions: - - Lock is held for the entire duration of file access (prevents partial reads/writes) - - Lock is released automatically on close (even on exceptions) - - Platform-specific implementation (fcntl for POSIX, msvcrt for Windows) - - Assumptions: - - File locking is supported on the platform - - Lock is advisory on some systems (cooperative locking) - """ - - def __init__(self, file_path: Path, mode: str): - self.file_path = file_path - self.mode = mode - self.file = None - - def __enter__(self): - self.file = open(self.file_path, self.mode, encoding="utf-8") - - try: - # Platform-specific locking - if os.name == "posix": - import fcntl - - fcntl.flock(self.file.fileno(), fcntl.LOCK_EX) - elif os.name == "nt": - import msvcrt - - # Lock 1 byte at position 0 - msvcrt.locking(self.file.fileno(), msvcrt.LK_LOCK, 1) - except Exception: - self.file.close() - self.file = None - raise - - return self.file - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.file: - # Unlock happens automatically on close - self.file.close() - - -def _exclusive_file_lock(file_path: Path, mode: str): - """Create an exclusive file lock context manager. - - :param file_path: Path to the file to lock. - :param mode: File open mode ('r', 'a', 'w', etc.). - :return: Context manager that returns an open file handle. - """ - return _ExclusiveFileLock(file_path, mode) - - -def _encode_cache_line(plaintext: str) -> str: - """Encode a single cache line using base64. - - :param plaintext: The plaintext string to encode. - :return: Base64-encoded string (safe for a single text line). - """ - return base64.b64encode(plaintext.encode("utf-8")).decode("ascii") - - -def _decode_cache_line(encoded: str) -> str: - """Decode a single base64-encoded cache line. - - :param encoded: The base64-encoded string. - :return: The decoded plaintext string. - """ - return base64.b64decode(encoded.encode("ascii")).decode("utf-8") + if cache_dir and Path(cache_dir).is_absolute(): + return Path(cache_dir) / ORT_SUPPORT_DIR + return _resolve_home_dir() / ".cache" / ORT_SUPPORT_DIR diff --git a/olive/workflows/run/run.py b/olive/workflows/run/run.py index 89100e1c1c..8ecb74952c 100644 --- a/olive/workflows/run/run.py +++ b/olive/workflows/run/run.py @@ -5,7 +5,7 @@ import logging from copy import deepcopy from pathlib import Path -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union from olive.common.utils import set_tempdir from olive.hardware.constants import ExecutionProvider @@ -13,6 +13,15 @@ from olive.package_config import OlivePackageConfig from olive.systems.accelerator_creator import create_accelerator from olive.systems.common import SystemType +from olive.telemetry.recipe_telemetry import _build_recipe_result_metadata, _load_config_input_for_telemetry +from olive.telemetry.telemetry import Telemetry, is_ci_environment +from olive.telemetry.telemetry_extensions import ( + _format_exception_message, + _is_exception_logged, + _mark_exception_logged, + log_error, + log_recipe_result, +) from olive.workflows.run.config import RunConfig if TYPE_CHECKING: @@ -152,30 +161,81 @@ def run( list_required_packages: bool = False, package_config: Optional[Union[str, Path, dict]] = None, tempdir: Optional[Union[str, Path]] = None, + recipe_telemetry_metadata: Optional[dict[str, Any]] = None, + emit_recipe_telemetry: bool = True, + emit_error_telemetry: bool = True, ): # set tempdir set_tempdir(tempdir) + package_config_input = package_config + try: + package_config_telemetry_input = ( + _load_config_input_for_telemetry(package_config_input) if package_config_input is not None else None + ) + except Exception: + package_config_telemetry_input = None + + package_config_provided = package_config is not None if package_config is None: package_config = OlivePackageConfig.get_default_config_path() - package_config = OlivePackageConfig.parse_file_or_obj(package_config) - run_config: RunConfig = RunConfig.parse_file_or_obj(run_config) - - if list_required_packages: - # set the log level to INFO for packages - set_verbosity_info() - required_packages = get_required_packages(package_config, run_config) - generate_files_from_packages(required_packages, "olive_requirements.txt") - return None - - if run_config.engine.host and run_config.engine.host.type == SystemType.Docker: - docker_system = run_config.engine.host.create_system() - return docker_system.run_workflow(run_config) - - # set log level for olive - set_default_logger_severity(run_config.engine.log_severity_level) - return run_engine(package_config, run_config) + parsed_run_config = None + success = False + exception = None + try: + package_config = OlivePackageConfig.parse_file_or_obj(package_config) + parsed_run_config = RunConfig.parse_file_or_obj(run_config) + + if list_required_packages: + # set the log level to INFO for packages + set_verbosity_info() + required_packages = get_required_packages(package_config, parsed_run_config) + generate_files_from_packages(required_packages, "olive_requirements.txt") + success = True + return None + + if parsed_run_config.engine.host and parsed_run_config.engine.host.type == SystemType.Docker: + docker_system = parsed_run_config.engine.host.create_system() + workflow_output = docker_system.run_workflow(deepcopy(parsed_run_config)) + success = True + return workflow_output + + # set log level for olive + set_default_logger_severity(parsed_run_config.engine.log_severity_level) + workflow_output = run_engine(package_config, parsed_run_config) + success = True + return workflow_output + except Exception as exc: + exception = exc + raise + finally: + if exception is not None and emit_error_telemetry and not _is_exception_logged(exception): + log_error( + exception_type=type(exception).__name__, + exception_message=_format_exception_message(exception, exception.__traceback__), + ) + _mark_exception_logged(exception) + if emit_recipe_telemetry: + try: + metadata = _build_recipe_result_metadata( + run_config, + None, + parsed_run_config, + recipe_telemetry_metadata, + list_required_packages=list_required_packages, + package_config_input=package_config_telemetry_input, + package_config_provided=package_config_provided, + ) + recipe_name = metadata.pop("recipe_name", None) + if recipe_name: + log_recipe_result(recipe_name, success=success, metadata=metadata) + except Exception: + logger.debug("Failed to emit recipe result telemetry.", exc_info=True) + if is_ci_environment(): + telemetry = Telemetry.get_existing_instance() + if telemetry is not None: + telemetry.shutdown(timeout_millis=2_000, callback_timeout_millis=2_000) def generate_files_from_packages(packages, file_name): diff --git a/requirements.txt b/requirements.txt index 1032c72f97..cfb5a1b9de 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,6 @@ numpy onnx onnx_ir>=0.1.2 onnxscript>=0.5.3 -opentelemetry-sdk>=1.39.1 optuna pandas pydantic>=2.0 diff --git a/scripts/quantize_and_compare_perplexity.py b/scripts/quantize_and_compare_perplexity.py new file mode 100644 index 0000000000..79b3d46b7b --- /dev/null +++ b/scripts/quantize_and_compare_perplexity.py @@ -0,0 +1,492 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Manual verification helper: compare wikitext perplexity before/after quantizing a real HF model. + +This is a local, one-off validation tool -- it is NOT wired into any Olive workflow config or +CI. It exists to answer "does quantizing this real (downloaded) checkpoint produce a sane +quality regression, or does it blow up / silently corrupt the model?" -- the same class of +question that surfaced real bugs in #2584 (RTN) after synthetic-model unit tests had already +passed. + +Usage: + python scripts/quantize_and_compare_perplexity.py \ + --model_id ibm-granite/granite-3.0-1b-a400m-base \ + --pass_name Gptq \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' + +By default this evaluates perplexity over the *entire* wikitext-2 test split (use +``--num_samples`` to restrict to a prefix for a quicker/dirtier check). The model can be any HF +model id or local path; the pass can be any registered Olive pytorch *quantization* pass (Gptq, +Rtn, KQuant, AutoAWQQuantizer, GptqQuantizer, GptqModel, ...) with any config -- nothing here is +hardcoded to GPTQ or to MoE. Non-quantization passes (e.g. SparseGPT, which prunes rather than +quantizes) can technically be loaded too, but the size/perplexity comparison this script prints is +framed around quantization and may be misleading for a pruning pass. + +Fairness notes (read before trusting a delta from this script): + * Baseline and quantized models are loaded with the *same* ``--dtype`` (default "auto", i.e. + each checkpoint's native dtype) and the same ``--experts_implementation``, so a measured + perplexity delta is attributable to quantization and not to an incidental dtype/backend + mismatch between the two loads. + * "Weights size" reports two *different* metrics, both labeled explicitly in the summary: an + in-memory figure (parameter count x element size at the loaded dtype), computed identically + for baseline and quantized so that pair is directly comparable, and a separately-labeled + on-disk figure (actual saved weight-file bytes, quantized side only) that reflects real + storage compression. Do not compare the in-memory baseline number against the on-disk + quantized number -- they measure different things. + * Calibration sample/token counts are captured by intercepting the *actual* dataset the pass + builds internally (not recomputed separately by this script), so they are structurally + guaranteed to match what was really used to calibrate -- and are skipped/reported as n/a for + data-free passes (e.g. RTN) that don't consume a data_config at all. + * "Quantization time" is wall-clock for the whole ``pass.run()`` call (load + calibrate + + quantize + save), not a pure inference/perf benchmark -- a data-free pass like RTN being much + faster than a calibration-based pass like GPTQ is an expected algorithmic trade-off, not a + regression. + * ``--device`` controls where the baseline model loads and where perplexity is evaluated for + both models. It does NOT control where quantization/calibration itself runs -- some passes + (e.g. Gptq's layerwise calibration) pick their own device internally (cuda when available, + else cpu) independent of this flag. On a single-GPU machine this is usually moot; on a + multi-GPU machine, quantization may run on a different GPU than the one named here. + * A single run against one (possibly non-random) slice of one dataset, on one machine/GPU, is a + smoke test, not a statistically rigorous benchmark -- treat deltas smaller than the + run-to-run/sample-to-sample noise floor with caution, especially for small ``--num_samples`` + overrides, and do not treat cited timing numbers as reproducible to more than roughly + plus-or-minus 20% run-to-run without re-measuring on your own hardware. + * Calibration data defaults to Olive's own built-in default (wikitext-2 train), matching what a + real user gets out of the box -- this is what the primary results table should be based on. + Wikitext-2 train and the (also wikitext-2, by default) eval split are disjoint but same-domain, + so as a *secondary*, opt-in cross-domain sanity check (not the headline number), evaluate on + C4 instead: ``--dataset allenai/c4 --dataset_config en --split validation --eval_streaming + --num_samples 200`` (C4 is sharded across 1000+ files, so non-streaming slicing is + prohibitively slow -- ``--eval_streaming`` is required for it). +""" + +import argparse +import gc +import json +import shutil +import statistics +import tempfile +import time +from pathlib import Path + +import torch + +# ruff: noqa: T201 # this is a CLI tool; print() output is the point + +_WEIGHT_FILE_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth") + + +def weights_size_gb(path: Path) -> float: + """On-disk size of only the model-weight files under ``path``, in GiB. + + Excludes config/tokenizer/vocab/metadata files so this is comparable across checkpoints that + may differ in how much non-weight metadata they carry. + """ + return sum(f.stat().st_size for f in path.rglob("*") if f.is_file() and f.suffix in _WEIGHT_FILE_SUFFIXES) / ( + 1024**3 + ) + + +def param_memory_gb(model: torch.nn.Module) -> float: + """In-memory footprint of a model's parameters at their current dtype, in GiB.""" + return sum(p.numel() * p.element_size() for p in model.parameters()) / (1024**3) + + +def capture_moe_fallback_counts() -> tuple[list, callable]: + """Monkeypatch ``CoverageReport.add`` to capture per-layer MoE fallback (RTN) stats. + + ``Gptq``'s MoE calibration only *logs* its coverage/fallback report -- it doesn't return it + to the caller. Since the fallback experts (too few calibration tokens for a useful Hessian) + are exactly the "which experts did NOT get real GPTQ treatment" signal this script wants, + intercept the structured ``LayerCoverage`` objects before they're formatted into a log + string, instead of parsing log text. + + Returns (captured_layers, restore_fn). Stays empty for non-MoE passes: the module-level patch + is applied unconditionally (import is cheap and side-effect-free until ``CoverageReport.add`` + is actually called), but data-free/non-MoE passes never call it, so ``captured`` stays empty. + """ + from olive.passes.pytorch.moe_calib import CoverageReport + + captured = [] + original_add = CoverageReport.add + + def patched_add(self, coverage): + captured.append(coverage) + return original_add(self, coverage) + + CoverageReport.add = patched_add + + def restore(): + CoverageReport.add = original_add + + return captured, restore + + +def capture_calibration_dataset() -> tuple[dict, callable]: + """Monkeypatch the calibration-dataset builder actually called by ``run_layerwise_quantization``. + + ``olive.passes.pytorch.quant_utils`` does ``from .train_utils import get_calibration_dataset``, + binding its own module-level name -- patching ``train_utils.get_calibration_dataset`` would + not affect that call site, so this patches ``quant_utils.get_calibration_dataset`` directly. + This captures the *exact* dataset the pass consumes (not a separately-recomputed copy), so + the reported sample/token counts are structurally guaranteed to match, and for data-free + passes (e.g. RTN) that never call this function, ``captured`` stays empty. + """ + import olive.passes.pytorch.quant_utils as quant_utils_mod + + captured: dict = {} + original = quant_utils_mod.get_calibration_dataset + + def patched(model, data_config): + dataset = original(model, data_config) + captured["dataset"] = dataset + return dataset + + quant_utils_mod.get_calibration_dataset = patched + + def restore(): + quant_utils_mod.get_calibration_dataset = original + + return captured, restore + + +def compute_perplexity( + model, tokenizer, text: str, device: str, stride: int = 512, max_len: int = 2048 +) -> tuple[float, int]: + """Sliding-window perplexity over ``text``, per the standard HF perplexity recipe. + + Uses overlapping windows (``max_len`` context, ``stride`` step) so every scored token has a + long enough context, without requiring the whole document to fit in one forward pass. + ``max_len`` is a fixed evaluation-window size (matching the common HF perplexity recipe, + which typically also uses a fixed window rather than each model's full context length) -- + it is applied identically to baseline and quantized so the comparison stays fair even though + it may be smaller than a given model's actual ``max_position_embeddings``. Override via + ``--max_len`` if you need a specific window size for your model. + + Returns ``(perplexity, n_scored_tokens)``. ``n_scored_tokens`` is the number of positions + that actually contributed to the loss -- not ``trg_len``, since a causal-LM's internal + label-shift means the first position of every window's target has no prediction to score. + """ + model.eval() + input_ids = tokenizer(text, return_tensors="pt").input_ids + seq_len = input_ids.size(1) + if seq_len < 2: + raise ValueError( + f"Eval text tokenized to only {seq_len} token(s) -- need at least 2 for a scoreable " + "causal-LM window. Check --dataset/--dataset_config/--split/--num_samples." + ) + + nlls = [] + n_tokens = 0 + prev_end = 0 + for begin in range(0, seq_len, stride): + end = min(begin + max_len, seq_len) + trg_len = end - prev_end + ids = input_ids[:, begin:end].to(device) + target_ids = ids.clone() + target_ids[:, :-trg_len] = -100 # ignore_index: don't double-count the overlapped prefix + with torch.no_grad(): + loss = model(ids, labels=target_ids).loss + # the model shifts labels internally (logits[:-1] vs labels[1:]), so one more position + # than "trg_len" is unscored; weight by what was actually scored, not the raw window size. + n_valid = (target_ids[:, 1:] != -100).sum().item() + nlls.append(loss * n_valid) + n_tokens += n_valid + prev_end = end + if end == seq_len: + break + return torch.exp(torch.stack(nlls).sum() / n_tokens).item(), n_tokens + + +def load_eval_text( + dataset: str, dataset_config: str, split: str, num_samples: int | None, streaming: bool = False +) -> str: + """Load and concatenate non-empty rows of ``split`` for perplexity evaluation. + + ``num_samples=None`` (the default) uses the *entire* split -- wikitext-2's test split has + many blank/heading-only rows, so a small fixed prefix (e.g. the first 200 rows) can end up + representing only a small, non-random, unrepresentative slice of the actual text. + + ``streaming=True`` is required for datasets sharded across many files (e.g. ``allenai/c4``, + used as an optional cross-domain eval set): a non-streaming row-slice like ``split[:50]`` + still triggers enumerating/downloading every shard's metadata to compute slice boundaries, + which is far slower than just streaming the first ``num_samples`` rows directly. Streaming + mode requires ``num_samples`` to be set (no notion of "the entire split" for an unbounded + stream). + """ + from datasets import load_dataset + + if streaming: + if num_samples is None: + raise ValueError("--num_samples is required when using a streaming eval dataset (e.g. C4).") + ds = load_dataset(dataset, dataset_config, split=split, streaming=True) + rows = [] + for row in ds: + if row["text"].strip(): + rows.append(row["text"]) + if len(rows) >= num_samples: + break + return "\n\n".join(rows) + + ds = load_dataset(dataset, dataset_config, split=split) + rows = [row for row in ds["text"] if row.strip()] + if num_samples is not None: + rows = rows[:num_samples] + return "\n\n".join(rows) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model_id", required=True, help="HF model id or local model path.") + parser.add_argument("--pass_name", default="Gptq", help="Olive pytorch pass class name, e.g. Gptq, Rtn, KQuant.") + parser.add_argument( + "--pass_config", + default='{"bits": 4, "group_size": 128, "sym": true}', + help="JSON dict of pass config kwargs (e.g. bits/group_size/sym/moe/...). Pass an explicit " + "'data_config' key here to override the calibration dataset (see --calib_* flags for a shortcut).", + ) + parser.add_argument("--dataset", default="wikitext", help="HF dataset id for the perplexity eval text.") + parser.add_argument("--dataset_config", default="wikitext-2-raw-v1", help="HF dataset config name.") + parser.add_argument("--split", default="test", help="Dataset split to use for the eval text.") + parser.add_argument( + "--num_samples", + type=int, + default=None, + help="Restrict the eval split to the first N (non-empty) rows. Default: use the entire split. " + "Required (must be set) when --eval_streaming is used.", + ) + parser.add_argument( + "--eval_streaming", + action="store_true", + help="Stream the eval dataset instead of a plain row-slice load. Required for datasets sharded across " + "many files, e.g. '--dataset allenai/c4 --dataset_config en --split validation --eval_streaming " + "--num_samples 200' to use C4 as a cross-domain eval set (checking whether same-domain " + "wikitext-calibration-on-wikitext-eval inflates quality vs. a genuinely held-out domain).", + ) + parser.add_argument( + "--dtype", + default="auto", + help="torch_dtype used to load BOTH the baseline and the (pre-quantization) input model, so a measured " + "perplexity delta is attributable to quantization and not to a dtype mismatch between the two loads. " + "'auto' (default) preserves each checkpoint's native dtype, matching Olive's own default load behavior.", + ) + parser.add_argument( + "--calib_dataset", + default=None, + help="HF dataset id for calibration data, e.g. 'allenai/c4' for a cross-domain calibration set (the " + "convention used by the GPTQ/AWQ papers) instead of Olive's wikitext-2 default, which shares a domain " + "with the default --dataset/--dataset_config eval text above (train/test are disjoint, but same-domain " + "calibration can still inflate quality on a same-domain eval set relative to true generalization). " + "Default: use Olive's built-in default (Salesforce/wikitext, wikitext-2-raw-v1). Ignored if --pass_config " + "already specifies 'data_config', or if the target pass doesn't declare a data_config parameter at all " + "(e.g. Rtn, which is data-free).", + ) + parser.add_argument("--calib_dataset_config", default=None, help="HF dataset config name for --calib_dataset.") + parser.add_argument( + "--calib_split", + default=None, + help="Dataset split to use for calibration data. Default: None, meaning don't override -- inherit " + "whatever Olive's own current default is (get_calibration_data_config's own 'split' default), so this " + "script always reflects real Olive behavior even if that default changes in the future, instead of " + "silently drifting out of sync with a value hardcoded here.", + ) + parser.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu") + parser.add_argument( + "--experts_implementation", + default="eager", + help="Set via model.set_experts_implementation(...) on BOTH the baseline and the quantized model before " + "evaluating a MoE model, so the two loads use the same MoE compute backend. Quantized MoE weights are " + "storage-only QuantTensors that only the eager experts loop can consume.", + ) + parser.add_argument( + "--keep_output", + action="store_true", + help="Keep the quantized model directory instead of deleting it when the script exits.", + ) + parser.add_argument( + "--max_len", + type=int, + default=2048, + help="Sliding-window context length used for perplexity evaluation (applied identically to " + "baseline and quantized). Default 2048, matching the common HF perplexity recipe -- override " + "if you specifically need a window matching a given model's max_position_embeddings.", + ) + args = parser.parse_args() + + if args.num_samples is not None and args.num_samples <= 0: + parser.error("--num_samples must be a positive integer.") + + pass_config = {} + try: + pass_config = json.loads(args.pass_config) + except json.JSONDecodeError as e: + parser.error(f"--pass_config is not valid JSON: {e}") + if not isinstance(pass_config, dict): + parser.error("--pass_config must be a JSON object (dict), e.g. '{\"bits\": 4}'.") + + from olive.hardware import DEFAULT_CPU_ACCELERATOR + from olive.model import HfModelHandler + from olive.package_config import OlivePackageConfig + from olive.passes.olive_pass import create_pass_from_dict + from olive.passes.pytorch.train_utils import get_calibration_data_config + + # Resolve the pass class through Olive's own package registry (olive_config.json) rather than + # guessing a module path from the class name: several registered passes (e.g. AutoAWQQuantizer + # in autoawq.py, GptqQuantizer in autogptq.py) do not live in a module named after the + # lowercased class name, so a naive `olive.passes.pytorch.{name.lower()}` import silently + # breaks for them. + pass_cls = OlivePackageConfig.load_default_config().import_pass_module(args.pass_name) + pass_accepts_data_config = "data_config" in pass_cls.default_config(DEFAULT_CPU_ACCELERATOR) + + print( + f"=== Loading eval text: {args.dataset}/{args.dataset_config} [{args.split}][:{args.num_samples}] " + f"(streaming={args.eval_streaming}) ===" + ) + text = load_eval_text(args.dataset, args.dataset_config, args.split, args.num_samples, args.eval_streaming) + + dtype_kwargs = {} if args.dtype == "none" else {"torch_dtype": args.dtype} + + print(f"=== Loading baseline model: {args.model_id} (dtype={args.dtype}) ===") + from transformers import AutoModelForCausalLM, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.model_id) + baseline = AutoModelForCausalLM.from_pretrained(args.model_id, **dtype_kwargs).to(args.device) + if hasattr(baseline, "set_experts_implementation"): + baseline.set_experts_implementation(args.experts_implementation) + baseline_size_gb = param_memory_gb(baseline) + baseline_dtype = next(baseline.parameters()).dtype + + print("=== Baseline perplexity ===") + baseline_ppl, baseline_n_tokens = compute_perplexity(baseline, tokenizer, text, args.device, max_len=args.max_len) + print(f"Baseline perplexity: {baseline_ppl:.4f} (scored {baseline_n_tokens} tokens)") + del baseline + gc.collect() + torch.cuda.empty_cache() + + # Load the pre-quantization input model with the SAME dtype kwargs as the baseline above, so the + # only difference between "baseline" and "quantized" is the quantization itself. + input_model = HfModelHandler(model_path=args.model_id, load_kwargs=dtype_kwargs) + + if pass_accepts_data_config and "data_config" not in pass_config: + calib_dataset_kwargs = {} + if args.calib_split: + calib_dataset_kwargs["split"] = args.calib_split + if args.calib_dataset: + calib_dataset_kwargs["data_name"] = args.calib_dataset + if args.calib_dataset_config: + calib_dataset_kwargs["subset"] = args.calib_dataset_config + data_config = get_calibration_data_config( + args.model_id, + trust_remote_code=input_model.get_load_kwargs().get("trust_remote_code", False), + **calib_dataset_kwargs, + ) + pass_config = {**pass_config, "data_config": data_config} + elif not pass_accepts_data_config and (args.calib_dataset or "data_config" in pass_config): + print( + f"NOTE: {args.pass_name} does not declare a 'data_config' parameter (data-free quantization); " + "--calib_dataset / pass_config['data_config'] will be ignored." + ) + + # Capture the calibration dataset actually built and consumed *inside* run_layerwise_quantization + # (rather than recomputing a separate copy here), so the reported sample/token counts are + # structurally guaranteed to match what was really used -- and stay empty for data-free passes. + captured_calib, restore_calib_capture = capture_calibration_dataset() + fallback_layers, restore_coverage_capture = capture_moe_fallback_counts() + + printable_pass_config = {k: v for k, v in pass_config.items() if k != "data_config"} + out_dir = Path(tempfile.mkdtemp(prefix="olive_quant_demo_")) + try: + print(f"=== Running {args.pass_name} pass with config={printable_pass_config} ===") + quant_pass = create_pass_from_dict(pass_cls, pass_config, disable_search=True) + quant_start = time.time() + output_model = quant_pass.run(input_model, str(out_dir)) + quant_duration_s = time.time() - quant_start + quantized_weights_size_gb = weights_size_gb(out_dir) + + if "dataset" in captured_calib: + calib_num_samples = len(captured_calib["dataset"]) + # Sum tokens across the full batch dimension of every row (not just row["input_ids"][0]) + # so this stays correct if a --calib_* override or a future default ever uses batch_size > 1. + calib_num_tokens = sum(row["input_ids"].numel() for row in captured_calib["dataset"]) + calib_summary = f"{calib_num_samples} samples, {calib_num_tokens} tokens" + else: + calib_summary = "n/a (data-free pass)" + + print("=== Loading quantized model ===") + quantized = output_model.load_model() + if hasattr(quantized, "set_experts_implementation"): + quantized.set_experts_implementation(args.experts_implementation) + quantized = quantized.to(args.device) + + print("=== Quantized perplexity ===") + quant_ppl, quant_n_tokens = compute_perplexity(quantized, tokenizer, text, args.device, max_len=args.max_len) + print(f"Quantized perplexity: {quant_ppl:.4f} (scored {quant_n_tokens} tokens)") + + # Apples-to-apples with baseline_size_gb: same measurement method (in-memory param + # bytes at the loaded dtype) on both sides. weights_size_gb(out_dir) above is a + # DIFFERENT metric (on-disk artifact size, includes any save-time packing/compression) + # and is reported separately -- conflating the two was Major finding #3 from review. + quantized_size_gb = param_memory_gb(quantized) + + total_experts = sum(lc.num_experts for lc in fallback_layers) + total_starved = sum(lc.starved for lc in fallback_layers) + total_unseen = sum(lc.unseen for lc in fallback_layers) + + print("\n=== SUMMARY ===") + print(f"Model: {args.model_id}") + print(f"Pass: {args.pass_name}({printable_pass_config})") + print(f"Dtype (baseline & input): {baseline_dtype} (--dtype={args.dtype})") + print(f"Experts implementation: {args.experts_implementation} (applied to both models)") + print(f"Calibration set: {calib_summary}") + print( + f"Quantization time: {quant_duration_s:.1f}s " + "(end-to-end: load + calibrate + quantize + save, NOT a pure inference benchmark)" + ) + print(f"Baseline weights size: {baseline_size_gb:.3f} GiB (in-memory params, {baseline_dtype})") + quantized_dtype = next(quantized.parameters()).dtype + print(f"Quantized weights size: {quantized_size_gb:.3f} GiB (in-memory params, {quantized_dtype})") + print( + f"Quantized on-disk size: {quantized_weights_size_gb:.3f} GiB (saved weight files only, different metric)" + ) + print( + f"Eval text: {args.dataset}/{args.dataset_config}[{args.split}]" + f"{f'[:{args.num_samples}]' if args.num_samples is not None else ' (full split)'}" + ) + print(f"Baseline perplexity: {baseline_ppl:.4f} ({baseline_n_tokens} tokens scored)") + print(f"Quantized perplexity: {quant_ppl:.4f} ({quant_n_tokens} tokens scored)") + print(f"Delta: {quant_ppl - baseline_ppl:+.4f}") + if fallback_layers: + fallback_label = ( + f"{args.pass_name} + RTN fallback" + if total_starved + total_unseen + else f"{args.pass_name} (no fallback)" + ) + print( + f"MoE fallback experts: {total_starved + total_unseen}/{total_experts} " + f"({total_starved} starved, {total_unseen} unseen) across {len(fallback_layers)} MoE layers " + f"-- quantized with RTN instead of {args.pass_name} due to insufficient calibration coverage " + f"[label: {fallback_label}]" + ) + print("MoE per-layer token/expert coverage:") + for lc in fallback_layers: + counts = sorted(lc.token_counts) + print( + f" {lc.layer_name}: {lc.num_experts} experts, tokens/expert " + f"min={counts[0]} median={statistics.median(counts):.0f} max={counts[-1]}, " + f"{lc.starved} starved, {lc.unseen} unseen" + ) + else: + print("MoE fallback experts: n/a (not a MoE calibration run)") + finally: + restore_coverage_capture() + restore_calib_capture() + if args.keep_output: + print(f"Quantized model kept at: {out_dir}") + else: + shutil.rmtree(out_dir, ignore_errors=True) + + +if __name__ == "__main__": + main() diff --git a/skills/olive/SKILL.md b/skills/olive/SKILL.md index efcfb0cb28..5b0aa8e803 100644 --- a/skills/olive/SKILL.md +++ b/skills/olive/SKILL.md @@ -5,7 +5,7 @@ license: MIT compatibility: Requires Python 3.10 or later and the olive-ai package. Model downloads and some dependency installations require network access; GPU, NPU, and vendor-specific workflows require matching hardware and runtimes. metadata: author: microsoft - version: "2.0.0" + version: "2.3.0" --- # Microsoft Olive @@ -24,6 +24,52 @@ Do not invent command flags, pass names, pass parameters, model types, or execut existing project already has an Olive config, preserve its conventions and make the smallest necessary change. +## Required target questions + +Before searching recipes, generating a workflow, inspecting candidate pass chains, or running an Olive +command, confirm the target requirements with the user. Do not infer them from the model, provider defaults, +installed hardware, or the source checkpoint: + +- Always confirm the requested output precision. For quantized workflows, confirm weight precision and + activation precision separately when both apply. +- Confirm the execution provider if the user has not specified it. +- For QNN, confirm the target device or backend, such as QNN GPU versus HTP/NPU, and the target Qualcomm SoC or product. + Derive and confirm the corresponding ORT `soc_model` setting when the workflow requires it. +- For OpenVINO, confirm the target device, such as CPU, GPU, or NPU, and the target OpenVINO runtime or + toolkit version. + +Ask for every missing or ambiguous value before continuing. These values are recipe-search constraints: +include the model or architecture, provider, device, precision, and, when applicable, SoC model or runtime +version in the search. A recipe for the wrong device, SoC, precision, or runtime version is reference +material, not a drop-in starting point. + +### Guide users who do not know the target details + +Do not require the user to understand Olive precision names, execution-provider internals, QNN backends, +SoC IDs, or runtime versioning. If the user does not know a required value, switch to guided discovery +before searching for a final recipe: + +1. Ask one plain-language question at a time. For precision, first ask whether the priority is quality, + balanced quality and resource use, or minimum model size and highest feasible speed. +2. Determine whether the target is the current machine or another device. Inspect local hardware, available + execution providers, SDKs, and runtime versions when the actual target is accessible; otherwise ask for + the target product or chip name. +3. Inspect the source model configuration and existing weight format so unsupported or redundant precision + conversions are not offered. +4. Map the user's goal and detected target to only the combinations supported by the installed Olive + version, exporter, provider, and relevant model family. +5. Present a small set of valid choices with short tradeoffs, mark a recommended choice, and get explicit + confirmation before searching recipes or generating a workflow. + +For QNN, translate a product or chip name into QNN GPU versus HTP/NPU and the required SoC setting; do not +ask a non-expert for a numeric `soc_model` value when it can be derived from authoritative QNN or ONNX +Runtime documentation. For OpenVINO, detect the installed version and available target devices when +possible, then ask the user to confirm the intended device. + +If a required target fact cannot be detected and the user cannot provide it, do not invent it. Explain which +decision remains unresolved. Produce a portable non-AOT or experimental scaffold only if the user explicitly +chooses that fallback, and do not label it as a hardware-specific final recipe. + ## Choose the right interface | User goal | Preferred interface | @@ -48,7 +94,8 @@ provider selection. 1. Identify the input model format and path or Hugging Face ID. 2. Identify the desired output: optimized ONNX, quantized model, adapter, benchmark, or reusable workflow. -3. Identify the target device and execution provider only when the user has not already specified them. +3. Complete the required target questions above. Do not search recipes or generate a config while a required + value is missing or ambiguous. 4. Check `olive --help` in the active environment. 5. Use an explicit output directory and `--log_level 1` for meaningful progress logs. 6. For expensive or unfamiliar high-level commands, add `--dry_run`. Inspect the generated @@ -69,12 +116,33 @@ Use a YAML or JSON workflow when the user needs multiple passes, reusable config evaluation, search, custom scripts, remote systems, or settings not exposed by a high-level command. Read [the workflow configuration guide](references/workflow-config.md) before creating or editing a -workflow. For a model- and provider-specific workflow, first look for the same model or a close architecture -in [microsoft/olive-recipes](https://github.com/microsoft/olive-recipes). Read that recipe's README and use -the executable workflow JSON named in its command; `info.yml` and `info.yaml` are recipe catalog metadata, -not files for `olive run --config`. - -If there is no close recipe, generate a version-specific config with a high-level command: +workflow. For a model- and provider-specific workflow, search +[microsoft/olive-recipes](https://github.com/microsoft/olive-recipes) before generating a generic config. +Read each selected recipe's README, executable workflow JSON, requirements, and version or commit pins; +`info.yml` and `info.yaml` are recipe catalog metadata, not files for `olive run --config`. + +If the exact model and provider are absent, derive a candidate from prior recipes instead of stopping at an +exact-name search or merely replacing `model_path` in one recipe. Triangulate from: + +- The same architecture or model family on any provider for model type, exporter, component layout, and + architecture-specific graph transformations. +- The same provider and device for systems, lowering, static-shape, compilation or AOT passes, environment + boundaries, and provider options. +- The same source weight format and quantization scheme for compatible quantization, calibration, and data + settings. + +Inspect the target model's Hugging Face configuration and repository metadata before merging those +references. Check its architecture, task, context and cache design, modality, parameter and checkpoint size, +and existing `quantization_config`. Do not schedule dequantization or a second quantizer for a pre-quantized +checkpoint unless the selected exporter explicitly supports that conversion. + +Preserve multi-stage workflows from the provider recipe, such as separate quantization and QNN AOT +environments. Copy a pass or setting only when its model format, graph, precision, device, and runtime +preconditions still hold. Never infer that two models are compatible from repository names alone; a +distilled model can use a different base architecture than its name suggests. + +Use a high-level dry run as an installed-version compatibility scaffold after studying the reference +recipes, not as the final inferred recipe: ```shell olive optimize \ @@ -85,6 +153,11 @@ olive optimize \ --dry_run ``` +Compare its output with the reference pass chains, inspect every nontrivial pass schema, and explain which +recipe supplied each model-specific or provider-specific decision. If the installed exporter or a required +pass does not support the target architecture, report the recipe as blocked rather than presenting a +structurally valid config as runnable. + When authoring a workflow: - Keep `input_model`, `passes`, and output settings explicit. @@ -124,6 +197,21 @@ Structural validation cannot prove that remote models are accessible, local data the model is supported by every pass, or the target hardware has enough memory. Surface those constraints instead of presenting validation as execution success. +## Deep dives for quantization work + +For PyTorch/Hugging Face weight quantization (RTN, GPTQ, and related passes), read these before +making nontrivial changes or answering detailed questions: + +- [Quantization onboarding](references/quantization-onboarding.md) -- pass overview, shared config + surface, RTN vs. GPTQ trade-offs, calibration data and split-hygiene notes. +- [MoE GPTQ onboarding](references/moe-gptq.md) -- per-expert calibration mechanics, the K-last + layout requirement and architecture allow-list, the dual fallback-threshold design (routing + skew vs. statistical sufficiency), and what a three-model benchmark showed about fallback rates + and quantization wall-time. +- [Profiling/benchmark example](references/profiling-benchmark-example.md) -- how to use + `scripts/quantize_and_compare_perplexity.py` to compare a pass against baseline on a real model, + with a worked three-model example table. + ## Dependency and hardware rules - Reuse the user's active environment when it already contains the required Olive and runtime packages. diff --git a/skills/olive/references/cli.md b/skills/olive/references/cli.md index ef489bb2ec..07691a75c9 100644 --- a/skills/olive/references/cli.md +++ b/skills/olive/references/cli.md @@ -131,7 +131,7 @@ Exporter choices are mutually exclusive: - `--use_model_builder` for supported generative models - `--use_dynamo_exporter` for PyTorch Dynamo export -- `--use_mobius_builder` for supported multi-component or multimodal models; requires `mobius-ai` +- `--use_mobius_builder` for supported multi-component or multimodal models; requires `mobius-onnx` - No exporter flag uses the command's default PyTorch/Optimum route Use `--target_opset`, `--torch_dtype`, and `--fixed_param_dict` only when the exporter supports them. diff --git a/skills/olive/references/moe-gptq.md b/skills/olive/references/moe-gptq.md new file mode 100644 index 0000000000..174e780c50 --- /dev/null +++ b/skills/olive/references/moe-gptq.md @@ -0,0 +1,242 @@ +# MoE GPTQ Onboarding + +How Olive's `Gptq` pass (`olive/passes/pytorch/gptq.py`) quantizes Mixture-of-Experts (MoE) +models, why it needs a dedicated calibration path (`olive/passes/pytorch/moe_calib.py`), the +dual fallback-threshold design, and what the empirical validation from a three-model benchmark +showed. Read [`quantization-onboarding.md`](quantization-onboarding.md) first for the general +RTN/GPTQ context if you haven't already. For how to reproduce the benchmark numbers cited below, +see [`profiling-benchmark-example.md`](profiling-benchmark-example.md). + +## Why MoE needs its own calibration path + +Standard (dense-model) GPTQ calibration hooks a single `nn.Linear` per layer and accumulates one +Hessian from whatever activations flow through it. Fused MoE architectures don't have a per-expert +`nn.Linear` — routing happens *inside* a single "experts" module's forward call, so a plain +`register_forward_hook` on that module sees one undifferentiated activation batch with no way to +attribute individual rows to the expert that actually processed them. + +`moe_calib.py` solves this by hooking into transformers' own experts-implementation registry +(`ALL_EXPERTS_FUNCTIONS` / `@use_experts_implementation`, requires `transformers >= 5.0`): +Olive registers one generic recording implementation and points the model at it for the duration +of calibration. Every decorated experts class dispatches through it uniformly, so this file +contains no per-architecture branching for the recording logic itself — only an allow-list (see +below) of which architectures it has been verified safe for. + +## The K-last layout requirement + +GPTQ's `(K, K)` Hessian math assumes the weight's contraction dimension (`K`, the input-feature +dimension) is the last dimension of the fused weight tensor: `(num_experts, out_features, K)`. +Olive's MoE support is **allow-listed**, not best-effort, because getting this wrong silently +mis-quantizes the model rather than erroring: + +- `SUPPORTED_MOE_EXPERTS_CLASSES` in `moe_calib.py` lists the exact experts class name (not just + `model_type`) verified to (a) use `(num_experts, out, in)` — K-last — fused weight storage, and + (b) carry the `@use_experts_implementation` decorator. Currently verified: DeepSeek-V3, + GraniteMoE, Jamba, Mixtral, OLMoE, Phi-MoE, Qwen2-MoE, Qwen3-MoE (against transformers 5.14.1). +- The class name is checked *in addition to* `model_type` specifically so a spoofed or mismatched + `config.model_type` string cannot smuggle an unverified experts module past the allow-list. +- **Transposed-layout architectures** (`gpt_oss`, `llama4`, `aria`) store `(num_experts, in, out)` + — K is not last. These are refused with a clear error rather than silently mis-quantized. + Supporting them requires a layout-normalization step that is deliberately out of scope for the + current implementation — if you need this, that normalization (transpose before Hessian + accumulation, transpose back before save) is the right place to start, not extending the + allow-list to include them as-is. + +If you're adding support for a new MoE architecture: verify its fused weight layout and decorator +support against the currently-installed `transformers` version before adding it to +`SUPPORTED_MOE_EXPERTS_CLASSES` — do not assume a new architecture matches by analogy to an +existing allow-listed one. + +## Per-expert fallback: two independent conditions + +Not every expert in a calibration run receives enough routed tokens to compute a well-formed +Hessian. When an expert's calibration signal is inadequate, `Gptq` falls back to quantizing that +expert with RTN instead of GPTQ. The design doc's rationale (`gptq_moe_design.md`) is that this +should never be *worse* than plain RTN for that expert, since GPTQ's correction degenerates +toward the damping prior once the Hessian is under-determined — but note this is a working design +policy, not a formally proven guarantee: the design doc itself flags that a damped, rank-deficient +Hessian does not mathematically reduce to exactly RTN (damping reweights the null directions +rather than eliminating their influence, and the layerwise corrections stay coupled across +columns). Treat "never worse than RTN" as the intended, empirically-motivated behavior rather than +a theorem. + +An expert falls back if it fails **either** of two independent conditions (an OR-gate, +implemented as `LayerCoverage.threshold = max(skew_threshold, k_threshold)` in `moe_calib.py` +and mirrored as `N < skew_threshold or N < sufficiency_threshold` in `Gptq._process_moe_param` in +`gptq.py`): + +1. **Routing skew** (`moe_fallback_threshold`, default `0.005` = 0.5%, matching GPTQModel's + default): is this expert under-served *relative to its peers* in this same calibration run? + Computed as `fallback_threshold * tokens_seen_by_this_layer`. This condition is + **scale-invariant** — if you 10x the calibration set, this expert's absolute token count also + scales ~10x, so the skew ratio is roughly unchanged. It catches routing imbalance but not + under-calibration of the whole run. Note this threshold is only ever binding when + `fallback_threshold * tokens_seen < K` for the parameter being checked — at Olive's default + 262,144-token calibration budget, that means the skew condition can only fire for `K <~ 1311`; + for larger `K` values it is structurally dominated by the sufficiency condition below (see the + OLMoE example) and never actually gates anything on its own. +2. **Statistical sufficiency** (`moe_fallback_min_k_multiple`, default `2.0`): does this expert + have *enough absolute samples* to produce a well-formed `(K, K)` Hessian at all? An expert's + Hessian `H = sum(x xT)` accumulated from `N` routed tokens satisfies `rank(H) <= N`, so + `N < K` is *necessary* for `H` to be singular but not by itself sufficient to prove it (`N` + linearly-dependent samples at `N >= K` can still leave `H` rank-deficient) — practically, + `N < K` (`k=1`) is only the bare-minimum-rank floor, not the point at which the Hessian is + actually well-conditioned. A full-team review of this pass measured the real RTN-vs-GPTQ MSE + crossover under anisotropic activations at 4-bit to sit closer to `N` in the `1.5x-2x K` range: + at the `k=1` floor, GPTQ was measured ~4-6% *worse* than RTN on average for experts in that + band. `moe_fallback_min_k_multiple` was therefore raised from the original `k=1` default to + `k=2` (`N = 2K`) as a conservative, empirically-motivated policy choice past that crossover — + not a proven sufficient threshold. No external MoE-quantization literature search turned up a + specific "safe multiple of K" recommendation, only qualitative guidance; `k=2` is the repo's + own measured choice. This condition is **absolute**: more calibration data always helps it + directly, unlike the routing-skew condition. + +Note also that fallback is decided **per (expert, parameter)**, not per expert as a whole: a fused +MoE layer typically has two quantizable parameters (`gate_up_proj`, `down_proj`), and they can +have different `K` (input-feature dimension), so the *same* expert can pass the sufficiency check +for one parameter and fail it for the other. "N/M fallback experts" in a coverage summary counts +starved-or-unseen occurrences for the parameter(s) actually tracked in the coverage report, not +necessarily every quantizable parameter for that expert — check the per-layer log for the +specific parameter dimension if you need to know exactly which weights fell back. + +### Why both conditions, and why this was a design-vs-implementation gap + +The original design doc (`gptq_moe_design.md`) settled on K-multiple sufficiency as the *sole* +fallback trigger, treating the routing-skew percentage as diagnostic/reporting-only ("skew = +diagnostic (report); sufficiency = trigger (fallback)"). The first shipped implementation, +however, only implemented the skew-percentage condition (`DEFAULT_MOE_FALLBACK_THRESHOLD = +0.005`) from its very first commit — the design's settled K-multiple decision was never actually +implemented. This was not a later regression; it was a design decision that was reconciled only +during later benchmark validation work, which added `moe_fallback_min_k_multiple` and made +*both* conditions gate the fallback via an OR-gate — a stricter combination than the design doc's +own final call, which gated on sufficiency alone and used skew only for reporting. + +**Empirical evidence the two conditions are not redundant** (OLMoE-1B-7B-0924, full WikiText-2 +`train` calibration, 262,144 calibration tokens reaching the layer, `k=1` i.e. +`moe_fallback_min_k_multiple=1.0`, the value in effect when this specific case was first found): +layer 2, expert 5 received `N=1331` tokens for its `gate_up_proj` parameter (K=2048). +- Skew threshold: `0.005 * 262,144 = 1310.7` — `1331 > 1310.7`, so the **skew-only** check would + **not** flag this expert. This is not because the expert is genuinely well-served: OLMoE routes + top-8-of-64, so a "fair share" of tokens for one expert is `262,144 * 8/64 = 32,768` — the + 1,331 tokens this expert saw is only ~4% of fair share (~24x under-routed). The skew threshold + simply doesn't bind at this `K`, for the structural reason noted above (0.5% of 262,144 = + 1310.7 < K=2048), not because the expert looks reasonably calibrated. +- Sufficiency threshold at `k=1`: `1.0 * K = 1.0 * 2048 = 2048.0` — `1331 < 2048.0`, so the + **sufficiency** check **does** flag it: this expert's Hessian is necessarily rank-deficient + (`N < K`). + +This confirms the two conditions are not redundant in general (there exist real experts the +sufficiency check catches that skew alone would not), though this specific benchmark only produced +a case in the "sufficiency fires, skew doesn't" direction — it does not exercise a case where skew +would fire but sufficiency wouldn't, so it validates the OR-gate's usefulness over +skew-alone, but does not by itself demonstrate a case where dropping the sufficiency-alone design +(skew as report-only) would have been wrong in the opposite direction. + +With the default later raised to `k=2` (`moe_fallback_min_k_multiple=2.0`, see the note above), +this same OLMoE layer/expert is caught even more clearly: the sufficiency threshold becomes +`2.0 * 2048 = 4096.0`, so `1331 < 4096.0` by a wider margin, and the same rerun at `k=2` also +newly catches several additional experts across other layers whose `N` fell in the `1x-2x K` +"no man's land" (e.g. layer 8 expert 18/48/55 with `N` in the 2445-3772 range against +`K=2048` — these pass at `k=1` but fail at `k=2`). See "What the three-model benchmark showed" +below for the aggregate before/after `k=1` vs `k=2` comparison. + +## Coverage logging and diagnostics + +Every MoE layer's calibration coverage is logged via `CoverageReport.log_summary()` (built from +per-layer `LayerCoverage.format()` strings; see `moe_calib.py`), which reports, per layer: number +of experts covered/starved/unseen, the combined effective threshold (`max(skew, sufficiency)`), +and the min/median/max observed token count per expert. Read this log when investigating an +unexpectedly high fallback count for a given model — it will tell you whether the fallback is +driven by routing imbalance across experts or raw insufficiency (the whole calibration set too +small for this model's K), and at which specific `K` (i.e. which parameter) the threshold bound. + +## What the three-model benchmark showed about fallback rates + +Fallback rate does **not** scale with the total number of experts in a model. From a benchmark +across three MoE models (bits=4, group_size=128, sym=true, full WikiText-2 `train` calibration, +`k=2` i.e. `moe_fallback_min_k_multiple=2.0` — see `profiling-benchmark-example.md` for the full +table and reproduction steps): + +| Model | Experts/layer x layers | Total experts | Fallback experts (k=2) | +| --- | --- | --- | --- | +| granite-3.0-1b-a400m-base | 32 x 24 | 768 | 3 (0.4%) | +| OLMoE-1B-7B-0924 | 64 x 16 | 1024 | 21 (2.1%) | +| Qwen1.5-MoE-A2.7B | 60 x 24 | 1440 | 0 (0.0%) | + +Qwen1.5-MoE has the *most* total experts (1440, driven by having more layers, not more +experts-per-layer than OLMoE) but *zero* fallbacks, while OLMoE has fewer total experts but the +*highest* fallback rate. Since calibration tokens are split **per layer** among that layer's +experts, the relevant comparison is experts-per-layer (~60-64, comparable across both models), +not the total. The proximate cause is genuine routing imbalance for the affected OLMoE experts +(e.g. the layer-2-expert-5 case above, ~24x under fair share), but note the *fallback condition +that actually fires* for these K=2048 cases is sufficiency, not skew (skew is structurally inert +at this K, per the note above) — "routing skew" describes the underlying router behavior, while +"sufficiency" is the specific check that catches it in the current implementation. Qwen1.5-MoE's +per-layer logs showed min tokens/expert consistently well above the `k=2` sufficiency threshold +across all 24 layers, with no layer ever showing a significantly under-routed expert. + +Raising the sufficiency multiplier from `k=1` to `k=2` roughly doubled the fallback count for +granite (2->3) and OLMoE (10->21) as expected — more experts in the `1x-2x K` "no man's land" are +now caught — while Qwen1.5-MoE's fallback count stayed at 0 in both runs (its router is well +enough balanced on this calibration set that no expert falls below even the `k=2` threshold). +Despite the higher fallback count, **GPTQ perplexity did not get measurably worse at `k=2` vs +`k=1`** (see `profiling-benchmark-example.md`'s `k=1` vs `k=2` comparison table) — consistent with +the `1x-2x K` band being a region where GPTQ wasn't reliably beating RTN in the first place, so +routing those experts to RTN removes downside risk at effectively no aggregate cost. + +**Takeaway**: read fallback count as a diagnostic of *that specific model's router balance on +that specific calibration domain*, not as a function of how many experts the model has. Two +models with similar experts-per-layer counts can have very different fallback rates purely +because their trained routers distribute tokens differently (this is plausibly related to +architecture differences like shared-expert-plus-top-k gating and the load-balancing auxiliary +loss used at pretraining time, but this benchmark did not directly inspect router logits to +confirm that hypothesis). + +## Quantization wall-time: calibration-set size matters less than you'd expect + +A common intuition is "bigger calibration set → proportionally longer GPTQ quantization." For +MoE models this is **not** the dominant effect. GPTQ per-layer cost splits into two phases with +very different scaling behavior: + +1. **Hessian accumulation** (forward passes over calibration data) — scales ~linearly with the + number of calibration tokens. +2. **Per-expert GPTQ solve** (Cholesky decomposition + blockwise quantization on each expert's + `(K, K)` Hessian) — a **fixed cost** independent of how many tokens built that Hessian; it + depends only on `K` and the number of (expert, parameter) pairs to solve. + +On granite-3.0-1b-a400m-base, increasing the calibration set from a stale `train[:1000]` slice +(~61,816 tokens, a single unreplicated prior run, not independently re-verified) to the full +`train` split (262,144 tokens, ~4.24x more) only increased quantization wall-time by ~18% +(658.7s at `k=1` vs. ~558s in that single comparison) — consistent with phase 2 (1,536 +fixed-cost Cholesky/quantize solves for this model: 24 layers x 32 experts x 2 params) dominating +total time for a model with a modest active-parameter count and many experts. This ratio is +model-specific and based on a single before/after comparison, not repeated runs — for a model +with a much larger active-parameter count (slower forward pass) or fewer experts (solve phase +less dominant), the balance could shift back toward calibration-size-linear scaling. Treat the +two-phase *mechanism* as the durable takeaway; treat the specific "~18% for ~4.24x more data" +ratio as one illustrative data point, not a general formula. + +The three-model table's wall-times at `k=1` (658.7s / 1499.3s / 2475.6s for 1,536 / 2,048 / 2,880 +total solves) are directionally consistent with solve-count-driven scaling, but with only three +data points where solve count, `K`, and total parameter volume all increase together, this +benchmark cannot cleanly separate "cost driven by number of solves" from "cost driven by total +active parameter volume" as the dominant covariate — both plausibly contribute. Don't treat +`num_layers x num_experts` as a validated predictive formula from this data alone; treat it as the +mechanistically-motivated hypothesis the two-phase model above suggests. Rerunning the same three +models at `k=2` (653.4s / 1484.9s / 2588.9s) showed wall-time essentially unchanged from `k=1` +despite the higher fallback count — consistent with the fixed-cost-per-solve model, since raising +the sufficiency multiplier slightly *reduces* the number of GPTQ solves actually performed (more +experts skip straight to RTN), so total time trends flat-to-down rather than up. + +## Before touching MoE GPTQ code + +1. Read `moe_calib.py`'s module docstring and `SUPPORTED_MOE_EXPERTS_CLASSES` first — the + allow-list is a deliberate fail-closed safety mechanism, not an oversight to work around. +2. If you're changing fallback-threshold behavior, keep both conditions in mind — they measure + different things (relative skew vs. absolute sufficiency) and are not interchangeable. +3. Check `test/passes/pytorch/test_gptq.py` and any MoE-specific test file for the existing test + conventions (parametrization over architectures, fixture patterns) before adding new tests. +4. Validate any threshold or calibration change against a real model, not just synthetic/tiny + test fixtures — the dual-condition disagreement case above only showed up on a real model + (OLMoE) with real routing behavior; tiny random-weight test models are unlikely to reproduce + realistic routing skew. diff --git a/skills/olive/references/profiling-benchmark-example.md b/skills/olive/references/profiling-benchmark-example.md new file mode 100644 index 0000000000..0d8d20bf69 --- /dev/null +++ b/skills/olive/references/profiling-benchmark-example.md @@ -0,0 +1,161 @@ +# Profiling / Benchmark Example: `quantize_and_compare_perplexity.py` + +A worked example for `scripts/quantize_and_compare_perplexity.py`, a local validation script that +quantizes a real Hugging Face model with a given Olive pass and reports the perplexity +regression, quantization wall-time, model size, and (for MoE calibration) per-expert fallback +coverage. Read [`quantization-onboarding.md`](quantization-onboarding.md) for background on the +passes being compared and [`moe-gptq.md`](moe-gptq.md) for the MoE-specific fallback mechanism +whose coverage output this script surfaces. + +## What the script is (and isn't) + +This is a **local, one-off validation tool** — it is not wired into any Olive workflow config or +CI. It exists to answer "does quantizing this real (downloaded) checkpoint produce a sane quality +regression, or does it blow up / silently corrupt the model?", the same class of question that +has previously surfaced real bugs after synthetic-model unit tests had already passed. It is not +a substitute for unit tests, and a single run is a smoke test, not a statistically rigorous +benchmark — see the "Fairness notes" in the script's own module docstring before trusting a +delta from it. + +The script is generic: `--pass_name` can be any Olive PyTorch quantization pass (`Gptq`, `Rtn`, +`KQuant`, ...) with any `--pass_config`, and `--model_id` can be any HF model id or local path — +nothing is hardcoded to GPTQ or to MoE. + +## Basic usage + +```shell +python scripts/quantize_and_compare_perplexity.py \ + --model_id ibm-granite/granite-3.0-1b-a400m-base \ + --pass_name Gptq \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' \ + --device cuda:0 +``` + +By default this evaluates perplexity over the **entire** WikiText-2 `test` split and calibrates +GPTQ on Olive's default WikiText-2 `train` split (use `--num_samples` to restrict the eval split +to a prefix for a quicker/dirtier check; see `--help` for calibration-side overrides such as +`--calib_dataset`). + +Key flags: + +| Flag | Purpose | +| --- | --- | +| `--model_id` | HF model id or local path (required). | +| `--pass_name` | Olive pass class name, e.g. `Gptq`, `Rtn` (default `Gptq`). | +| `--pass_config` | JSON dict of pass config overrides. | +| `--dataset` / `--dataset_config` / `--split` | Eval dataset (default WikiText-2 `test`). | +| `--calib_dataset` / `--calib_dataset_config` | Calibration dataset override for calibrated passes. | +| `--dtype` | Load dtype for both baseline and quantized models (default `auto`, i.e. each checkpoint's native dtype). | +| `--experts_implementation` | Forces a specific MoE experts implementation on both loads for a fair comparison. | +| `--device` | Device for both quantization and eval. | + +## Fairness guarantees baked into the script (read before trusting a delta) + +- Baseline and quantized models are loaded with the *same* `--dtype` and + `--experts_implementation`, so a measured perplexity delta is attributable to quantization, not + to an incidental dtype/backend mismatch between the two loads. +- Reported "weights size" includes **both** an in-memory (parameter-count x dtype-size) figure — + computed identically for baseline and quantized so the two are apples-to-apples — and a + separately labeled on-disk size (actual saved weight-file bytes) for the quantized model. Do + not compare the in-memory baseline figure against the on-disk quantized figure; they are + different metrics reported side-by-side, not the same metric. +- Calibration sample/token counts are captured by intercepting the actual dataset the pass builds + internally, not recomputed separately by the script, so they are guaranteed to match what was + really used to calibrate — and are reported as `n/a` for data-free passes (e.g. RTN). +- "Quantization time" is wall-clock for the entire pass run (load + calibrate + quantize + save), + not a pure inference benchmark — RTN being much faster than GPTQ is an expected algorithmic + trade-off, not a regression. + +## Worked example: three-way MoE model comparison + +The following table was produced by running the script once per (model, pass) pair — six runs +total — with `bits=4, group_size=128, sym=true, moe=true`, full WikiText-2 `train` calibration +(Olive's default post-#2609), and the full WikiText-2 `test` split for eval: + +```shell +# Baseline + RTN +python scripts/quantize_and_compare_perplexity.py \ + --model_id --pass_name Rtn \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' \ + --device cuda:0 + +# Baseline + GPTQ +python scripts/quantize_and_compare_perplexity.py \ + --model_id --pass_name Gptq \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' \ + --device cuda:0 +``` + +| Model | Baseline PPL | RTN PPL (Δ) | GPTQ PPL (Δ) | Quant time RTN / GPTQ | Fallback experts | +| --- | --- | --- | --- | --- | --- | +| ibm-granite/granite-3.0-1b-a400m-base | 6.2877 (354,564 tok) | 7.5861 (+1.2984) | 6.9492 (+0.6615) | 8.0s / 653.4s | 3/768 (0.4%) | +| allenai/OLMoE-1B-7B-0924 | 6.6182 (288,720 tok) | 7.1091 (+0.4909) | 6.8937 (+0.2755) | 52.6s / 1484.9s | 21/1024 (2.1%) | +| Qwen/Qwen1.5-MoE-A2.7B | 6.4246 (298,937 tok) | 6.9251 (+0.5005) | 6.6117 (+0.1872) | 100.0s / 2588.9s | 0/1440 (0.0%) | + +Calibration set for all GPTQ runs: 128 samples / 262,144 tokens (full WikiText-2 `train` split). +`moe_fallback_min_k_multiple=2.0` (`k=2`, the current default; see `moe-gptq.md`). All +baseline/quantized in-memory weight sizes are equal within each model (fake-quantization +dequantizes back to the original dtype for `transformers` compatibility) — only the *on-disk* +saved size actually shrinks; see the script's own summary output for per-run on-disk figures. + +**What this table demonstrates**: + +- GPTQ beat RTN on perplexity delta for every model tested (not just on average) — e.g. granite: + +0.6615 vs. +1.2984; Qwen1.5-MoE: +0.1872 vs. +0.5005 — at the cost of substantially longer + quantization time (minutes vs. seconds). +- MoE fallback rate is **not** proportional to total expert count — see `moe-gptq.md` for the + detailed explanation (Qwen1.5-MoE has the most total experts of the three but zero fallbacks; + the driver is per-model routing skew, not raw expert count). +- Quantization time scales roughly with `num_layers x num_experts` (the per-expert Cholesky solve + count), not with calibration token count — see `moe-gptq.md` for why a ~4x increase in + calibration tokens (from a stale `train[:1000]` slice to the full `train` split) only produced + an ~18% wall-time increase on granite, rather than the naively-expected ~4x. + +### `k=1` vs `k=2`: effect of the sufficiency-threshold multiplier + +`moe_fallback_min_k_multiple` was raised from `k=1` (the original shipped default, `N=K`) to +`k=2` (`N=2K`, the current default) after a full-team review found the real RTN-vs-GPTQ MSE +crossover under anisotropic activations at 4-bit sits closer to `1.5x-2x K`, not `1x K` (see +`moe-gptq.md`). All three models were rerun with the identical methodology at both settings to +quantify the actual effect: + +| Model | GPTQ PPL (Δ) at k=1 | GPTQ PPL (Δ) at k=2 | Fallback experts k=1 | Fallback experts k=2 | GPTQ time k=1 | GPTQ time k=2 | +| --- | --- | --- | --- | --- | --- | --- | +| granite-3.0-1b-a400m-base | 6.9560 (+0.6683) | 6.9492 (+0.6615) | 2/768 (0.3%) | 3/768 (0.4%) | 658.7s | 653.4s | +| OLMoE-1B-7B-0924 | 6.8966 (+0.2784) | 6.8937 (+0.2755) | 10/1024 (1.0%) | 21/1024 (2.1%) | 1499.3s | 1484.9s | +| Qwen1.5-MoE-A2.7B | 6.6117 (+0.1872) | 6.6117 (+0.1872) | 0/1440 (0.0%) | 0/1440 (0.0%) | 2475.6s | 2588.9s | + +Baseline and RTN-only numbers are identical between the two runs (RTN never reads +`moe_fallback_min_k_multiple`), confirming no other environment drift between the two benchmark +sessions. Despite roughly doubling the fallback count for granite and OLMoE, **perplexity did not +get measurably worse at `k=2` — it stayed flat or improved slightly**, and quantization time did +not increase (granite/OLMoE were slightly faster; Qwen1.5-MoE's small increase is within normal +run-to-run variance for a ~2,500s job, and it has 0 fallback at both settings so there is no +solve-count difference to explain it). The likely explanation: experts in the `1x-2x K` band have +technically-full-rank but severely ill-conditioned Hessians, so GPTQ's correction there was +already dominated by the damping prior rather than real signal — routing those borderline experts +to RTN at `k=2` removes the risk of an unlucky bad correction without giving up much upside, so +raising the threshold is close to a free win on the models tested here. + + +## Interpreting a run's console output + +Each run prints a `=== SUMMARY ===` block. For a MoE calibration run, also read the per-layer +`MoE coverage [...]` log lines emitted during quantization — each reports the effective fallback +threshold (`max(skew, sufficiency)`), how many experts were starved/unseen, and the observed +min/median/max token counts per expert for that layer. This is the fastest way to tell whether an +unexpectedly high fallback count is caused by routing imbalance in a specific layer or by an +under-sized calibration set overall. + +## Tips for running your own comparison + +- Start with a small/tiny model (e.g. a `*-tiny-random` HF checkpoint) to smoke-test your + `--pass_config` before committing GPU time to a multi-billion-parameter model — GPTQ + quantization time for a large MoE model can run into the tens of minutes. +- Pin `--dtype` explicitly if you need bit-for-bit reproducible baseline numbers across machines; + `auto` (the default) picks each checkpoint's native dtype, which is normally what you want for + a realistic baseline but can differ between environments with different default dtype handling. +- If you see a `transformers` warning about token sequence length exceeding + `max_position_embeddings` during eval-text tokenization, this is expected and harmless — the + perplexity computation uses a sliding window (`stride`/`max_len`), not a single full-sequence + forward pass; it does not indicate truncation or a scoring bug. diff --git a/skills/olive/references/quantization-onboarding.md b/skills/olive/references/quantization-onboarding.md new file mode 100644 index 0000000000..7b65a48799 --- /dev/null +++ b/skills/olive/references/quantization-onboarding.md @@ -0,0 +1,111 @@ +# Olive Quantization Onboarding + +An orientation to Olive's weight-quantization passes for PyTorch/Hugging Face models: what each +pass does, when to reach for it, the config knobs they share, and where to look in the codebase +before making changes. For MoE-specific GPTQ details (per-expert Hessians, fallback thresholds), +see [`moe-gptq.md`](moe-gptq.md). For a worked benchmark example comparing passes end-to-end, see +[`profiling-benchmark-example.md`](profiling-benchmark-example.md). + +## Where quantization passes live + +PyTorch/Hugging Face weight-quantization passes live in `olive/passes/pytorch/`: + +| Pass | Real class / registry name | Module | Calibration data? | Notes | +| --- | --- | --- | --- | --- | +| `Rtn` | `Rtn` | `rtn.py` | No | Round-to-nearest; fastest, data-free, weakest accuracy recovery. | +| `Gptq` | `Gptq` | `gptq.py` | Yes | Layerwise, Hessian-based weight correction using calibration data. | +| AutoGPTQ wrapper | `GptqQuantizer` | `autogptq.py` | Yes | Thin wrapper delegating to the third-party `auto-gptq` library. | +| GPTQModel wrapper | `GptqModel` | `gptqmodel.py` | Yes | Thin wrapper delegating to the third-party `gptqmodel` library. | +| AutoAWQ wrapper | `AutoAWQQuantizer` | `autoawq.py` | Yes | Wraps the third-party `autoawq` library (activation-aware weight quantization). | +| K-quant | `KQuant` | `kquant.py` | Varies | K-quant style block quantization. | + +The registry/class name (not a guessed lowercase-of-class-name module path) is what you must pass +to `--pass_name` for `scripts/quantize_and_compare_perplexity.py`, to `olive_config.json`'s +`"passes"` map, and to workflow configs — check `olive/olive_config.json` if you're ever unsure of +the exact registered name for a pass. + +ONNX-side quantization passes (post-export, operate on ONNX graphs rather than PyTorch modules) +live separately in `olive/passes/onnx/` (e.g. `rtn_quantization.py`, `hqq_quantization.py`, +`nvmo_quantization.py`, `inc_quantization.py`) — those are a different code path and are out of +scope for this doc, which focuses on the PyTorch-side `Rtn`/`Gptq` family. + +Shared logic (quantizer construction, model wrapping/unwrapping, layerwise iteration, save/load) +lives in `olive/passes/pytorch/quant_utils.py`. Read `get_quantizer_config()`, +`prepare_model()`, `run_layerwise_quantization()`, and `finalize()` there before modifying any +pass — almost every pass calls into these four functions and duplicating their logic in a new +pass is very rarely the right move. + +## Shared config surface + +Every weight-quantization pass accepts a common set of parameters from +`get_quantizer_config()` in `quant_utils.py`: + +| Param | Meaning | +| --- | --- | +| `bits` | Quantization bit-width (`PrecisionBits.BITS2/4/8`). | +| `group_size` | Block size for per-group scale/zero-point (`-1` means per-channel/whole-row). | +| `sym` | Symmetric (zero-point fixed at the bit-width's midpoint) vs. asymmetric quantization. | +| `lm_head` | Whether to also quantize the language-model head. | +| `embeds` | Whether to also quantize input embeddings (`Rtn` and `KQuant` only). | +| `overrides` | Per-module overrides for any of the above, keyed by module name pattern. | + +`Gptq` additionally exposes `damp_percent` (Hessian damping factor), `desc_act` (activation-order +column permutation, only valid for `group_size=-1`), and `data_config` (calibration dataset — see +below). When `moe=True` it also exposes `moe_fallback_threshold` and +`moe_fallback_min_k_multiple` — see `moe-gptq.md`. + +## RTN vs. GPTQ: when to use which + +- **RTN** is data-free and much faster than GPTQ (single-digit seconds for a ~1.3B active-param + model, tens of seconds to ~1-2 minutes for larger multi-billion-parameter MoE models in + practice — see `profiling-benchmark-example.md` for exact figures), but has the largest + accuracy regression of the two. Use it as a fast baseline, for environments without a + calibration dataset, or when the accuracy loss is acceptable for the target use case. +- **GPTQ** uses a calibration dataset to compute per-layer Hessians (`H = sum(x xT)` over + observed activations) and applies second-order error correction while quantizing each column, + substantially reducing the accuracy regression relative to RTN at the cost of a much longer + quantization pass (calibration forward passes + per-layer Cholesky/blockwise-quantize solves). + +Empirically (see `profiling-benchmark-example.md` for the full three-model MoE comparison), GPTQ +consistently produced a smaller perplexity regression than RTN on every model tested — but at +roughly 30-80x the quantization wall-time of RTN on the models measured there (single-run, +single-machine timings; treat the exact ratio as illustrative, not a guaranteed multiplier for +every model/hardware combination). Choose RTN when turnaround time matters more than the last bit +of accuracy; choose GPTQ when accuracy matters more and you can afford a longer one-time +quantization pass. + +## Calibration data (`data_config`) + +`Gptq` (and the other calibration-based passes) accept a `data_config` pointing at a Hugging +Face dataset config, or fall back to Olive's default WikiText-2 calibration set for +`HfModelHandler` inputs if none is given. See `olive/data/config.py` and +`configure-workflows/how-to-configure-data` in the Sphinx docs for how to point at a custom +dataset. + +Two split-hygiene facts worth internalizing: + +- WikiText-2's `train`/`validation`/`test` are official, pre-defined, non-overlapping Hugging + Face dataset splits (split by source Wikipedia article) — Olive does not construct or dedupe + them itself. A handful of near-identical rows across splits (e.g. repeated section-header + strings like `" = = Career = = "`) are not real leakage; they are trivial boilerplate that + recurs across unrelated articles. +- Use the `train` split (or a data-appropriate calibration set) for GPTQ calibration and a + disjoint split (e.g. `test`) for perplexity evaluation. Do not calibrate and evaluate on the + same rows. + +## Before modifying a quantization pass + +1. Read `quant_utils.py` fully — most behavior you might think belongs in a specific pass file is + actually implemented once, shared across passes. +2. Check `test/passes/pytorch/` for the existing test pattern for the pass you're touching (e.g. + `test_gptq.py`, `test_rtn.py`) before writing new tests; follow the existing + parametrization/fixture conventions rather than introducing a new test style. +3. If your change affects calibration or fallback behavior for Mixture-of-Experts models + specifically, read `moe-gptq.md` first — MoE calibration has its own module + (`moe_calib.py`) and several subtleties (per-expert Hessians, routing skew vs. statistical + sufficiency) that are easy to get wrong by analogy with the dense-model code path. + +If you're adding a brand-new pass rather than modifying an existing one, see the Sphinx guide +`docs/source/how-to/extending/how-to-add-optimization-pass.md` for how to register it in +`olive/olive_config.json` and wire it into the pass-discovery system — that step is not covered +here. diff --git a/skills/olive/references/workflow-config.md b/skills/olive/references/workflow-config.md index 2b9b7ecee1..66f1c122e6 100644 --- a/skills/olive/references/workflow-config.md +++ b/skills/olive/references/workflow-config.md @@ -10,18 +10,95 @@ Olive accepts both YAML and JSON. Hand-maintained workflows may use YAML for com [microsoft/olive-recipes](https://github.com/microsoft/olive-recipes) normally uses JSON for executable workflows. JSON does not allow comments or trailing commas. -## Choose a starting point - -For a model- and provider-specific workflow, first find the same model or a close architecture in -`microsoft/olive-recipes`. Read the recipe README and use the JSON file named in its `olive run --config` -command. Files named `info.yml` or `info.yaml` describe the recipe for catalog and automation purposes; they -are not Olive workflow configs. - -Reuse the closest recipe's model type, pass chain, system, provider, and data shape, then change only the -fields required for the user's model and output. Recipes can require backend-specific packages or target a -different Olive version, so validate the result against the active installation. - -When no close recipe exists, generate a workflow for the installed Olive version: +## Collect search constraints first + +Do not search for or synthesize a recipe until the target constraints are known: + +- Precision: always confirm the requested output precision. For quantized workflows, distinguish weight + precision from activation precision. +- Provider: confirm the intended execution provider. +- QNN: confirm QNN GPU versus HTP/NPU, the target Qualcomm SoC or product, and whether the workflow must produce an AOT context binary (derive and confirm the ORT `soc_model` value when required). +- OpenVINO: confirm the CPU, GPU, or NPU target and the exact OpenVINO runtime or toolkit version. + +Ask the user for missing values rather than selecting defaults. Use the resulting model or architecture, +provider, device, precision, SoC, runtime version, and output form as recipe-search terms. Provider-only +searches are insufficient because QNN GPU and HTP/NPU workflows differ, as can OpenVINO workflows across +devices and releases. + +### When the user does not know the constraints + +Translate technical constraints into guided choices instead of guessing: + +1. Ask one question at a time, starting with a plain-language optimization goal: quality first, balanced, or + minimum size and highest feasible speed. +2. Establish whether the target is local or remote. Detect hardware, available execution providers, SDKs, + and runtime versions on an accessible target; otherwise ask for its product or chip name. +3. Inspect the source model's architecture and current weight format before offering precision choices. +4. Eliminate combinations unsupported by the installed Olive version, exporter, provider, or model family. +5. Offer a small set of valid choices, explain the tradeoffs, recommend one, and require explicit + confirmation before the final recipe search. + +For QNN, derive the backend and SoC setting from an authoritative product-to-SoC mapping when possible. For +OpenVINO, detect the installed runtime version and available devices when possible. Hardware and runtime +inspection is constraint discovery, not permission to silently choose a target. + +If the target is remote or unavailable and the device, SoC, or runtime version remains unknown, stop before +hardware-specific lowering or AOT compilation. A portable non-AOT scaffold is acceptable only after the +user explicitly selects it and it is labeled experimental. + +## Choose and synthesize a starting point + +For a model- and provider-specific workflow, search `microsoft/olive-recipes` before generating a generic +workflow. An exact model-and-provider recipe is the strongest starting point. Read its README and use the +JSON file named in its `olive run --config` command. Files named `info.yml` or `info.yaml` describe the +recipe for catalog and automation purposes; they are not Olive workflow configs. + +An exact-name miss does not mean there is no useful prior art. Build a reference set from the closest +available recipes: + +- Same architecture or base model on any provider: learn the input model type, exporter, component layout, + cache I/O, and graph surgeries. +- Same provider and device with the closest compatible architecture: learn the systems, execution provider, + static shapes, compilation or AOT stages, and provider options. +- Same source weight format or quantization: learn compatible precision conversions, quantizers, + calibration data, and block or group settings. +- Similar model scale or component count: learn splitting strategies, memory constraints, and staged + outputs. + +Repository names are not architecture evidence. For example, a distilled model can use Qwen or Llama +internals while a newer model from the same publisher can introduce an unsupported architecture. Confirm +the target model's `architectures`, `model_type`, task, context/cache design, modality, parameter count, +checkpoint size, and `quantization_config` from its model configuration and repository metadata. + +Read the README, all executable configs used by its commands, requirements, and Olive/runtime version pins +for every selected reference. Provider recipes may intentionally separate quantization, export, and AOT +compilation into different configs and Python environments. Preserve those stage boundaries unless the +installed pass documentation explicitly supports combining them. + +Synthesize the candidate by assigning each concern to the most relevant reference: + +1. Start with the exporter and graph structure from the architecture reference. +2. Apply systems, provider lowering, static-shape, and compilation stages from the provider reference. +3. Apply quantization only from a reference with a compatible input weight format and exporter path. +4. Carry over data shapes and calibration only when the target model and pass require the same semantics. +5. Keep pass ordering, intermediate model types, and environment handoffs intact. +6. Record which recipe supports each inherited pass or non-obvious setting. + +Architecture-specific graph surgeries may be required. Retain or adapt them when the target uses the same +architecture and exporter and the exported graph satisfies the surgery's pattern and pass preconditions. +Do not transfer surgeries, cache names, tensor shapes, or quantization settings solely because they appear +in the closest provider recipe. Check each pass's input model type, output model type, architecture +assumptions, precision support, device, and execution provider. If the source checkpoint is already FP8, +FP4, GPTQ, AWQ, or another quantized format, do not blindly add a second quantizer; first verify that the +exporter can consume or intentionally convert that representation. + +For QNN, distinguish QNN GPU from HTP/NPU targets by the recipe's accelerator and provider options rather +than assuming that `QNNExecutionProvider` always means NPU. Preserve a recipe's separate host quantization +and QNN compilation environments, including its `PythonEnvironment`, intermediate model path, and context +binary stage. + +After deriving the candidate, generate a workflow for the installed Olive version as a compatibility +scaffold: ```shell olive optimize \ @@ -32,8 +109,14 @@ olive optimize \ --dry_run ``` -Olive writes `generated-workflow/config.json`. Edit that file rather than rebuilding a complex provider -recipe from memory. +Olive writes `generated-workflow/config.json`. Compare it with the reference set instead of accepting it as +the final recipe. Keep installed-version field names and defaults where they are compatible, then restore +recipe-specific stages that the high-level command cannot express. A dry run validates argument handling; +it is not evidence that the exporter, passes, model architecture, or target runtime support the model. + +When no adequate architecture or provider reference exists, use the generated workflow only as a clearly +labeled experimental scaffold. Do not present it as a runnable inferred recipe until the required exporter +and pass compatibility has been established. For a hand-authored starting point, copy `assets/workflow.yaml` from this skill. It is a classic Hugging Face-to-ONNX conversion and graph-optimization example, not a universal template for current generative diff --git a/test/cli/init/test_init_command.py b/test/cli/init/test_init_command.py index 257f8c6dad..7ae505aeec 100644 --- a/test/cli/init/test_init_command.py +++ b/test/cli/init/test_init_command.py @@ -2,7 +2,9 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # ------------------------------------------------------------------------- -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +import pytest class TestInitCommand: @@ -31,7 +33,46 @@ def test_run(self, mock_wizard_cls): args = parser.parse_args(["init", "-o", "./my-output"]) cmd = InitCommand(parser, args, []) - cmd.run() + telemetry = MagicMock(accepts_detailed_events=True) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.log_action") as log_action, + patch("olive.telemetry.telemetry_extensions.log_error") as log_error, + ): + cmd.run() mock_wizard_cls.assert_called_once_with(default_output_path="./my-output") mock_wizard_cls.return_value.start.assert_called_once() + log_action.assert_called_once() + assert log_action.call_args.kwargs["action_name"] == "Init" + assert log_action.call_args.kwargs["success"] is True + log_error.assert_not_called() + + @patch("olive.cli.init.wizard.InitWizard") + def test_run_failure_emits_one_action_and_error(self, mock_wizard_cls): + from argparse import ArgumentParser + + from olive.cli.init import InitCommand + + parser = ArgumentParser() + sub_parsers = parser.add_subparsers() + InitCommand.register_subcommand(sub_parsers) + args = parser.parse_args(["init"]) + cmd = InitCommand(parser, args, []) + mock_wizard_cls.return_value.start.side_effect = RuntimeError("boom") + + telemetry = MagicMock(accepts_detailed_events=True) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.log_action") as log_action, + patch("olive.telemetry.telemetry_extensions.log_error") as log_error, + pytest.raises(RuntimeError, match="boom"), + ): + cmd.run() + + log_action.assert_called_once() + assert log_action.call_args.kwargs["action_name"] == "Init" + assert log_action.call_args.kwargs["success"] is False + log_error.assert_called_once() diff --git a/test/cli/test_api.py b/test/cli/test_api.py index f6bfd151f8..f65c79f4f4 100644 --- a/test/cli/test_api.py +++ b/test/cli/test_api.py @@ -91,6 +91,18 @@ def test_optimize_function_basic(self, mock_cmd_cls): mock_cmd.run.assert_called_once() assert result is mock_output + def test_optimize_function_latches_disable_telemetry(self): + from olive import optimize + from olive.cli.optimize import OptimizeCommand + + with ( + patch("olive.cli.api.disable_telemetry") as disable, + patch.object(OptimizeCommand, "run", return_value=MagicMock()), + ): + optimize("test_model", disable_telemetry=True) + + disable.assert_called_once_with() + @patch("olive.cli.api.CaptureOnnxGraphCommand") def test_capture_cmd_basic(self, mock_cmd_cls): from olive import capture_onnx_graph diff --git a/test/cli/test_base.py b/test/cli/test_base.py index 428bf74095..640259e074 100644 --- a/test/cli/test_base.py +++ b/test/cli/test_base.py @@ -2,14 +2,40 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # ------------------------------------------------------------------------- +# pylint: disable=protected-access import json +import sys +import types from pathlib import Path from types import SimpleNamespace from unittest.mock import patch import pytest -from olive.cli.base import get_input_model_config +from olive.cli.base import BaseOliveCLICommand, get_input_model_config + + +def _mock_genai_version(monkeypatch): + """Install a stub onnxruntime_genai so the version guard passes without the real package.""" + genai_module = types.ModuleType("onnxruntime_genai") + genai_module.__version__ = "0.15.0" + monkeypatch.setitem(sys.modules, "onnxruntime_genai", genai_module) + + +def test_parse_extra_options_splits_key_value_pairs(monkeypatch): + _mock_genai_version(monkeypatch) + + assert BaseOliveCLICommand._parse_extra_options(["exclude_embeds=true", "filename=model=part.onnx"]) == { + "exclude_embeds": "true", + "filename": "model=part.onnx", + } + + +def test_parse_extra_options_rejects_missing_value_separator(monkeypatch): + _mock_genai_version(monkeypatch) + + with pytest.raises(ValueError, match="Expected key=value"): + BaseOliveCLICommand._parse_extra_options(["exclude_embeds"]) @pytest.mark.parametrize( diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 999e189132..ed5d39e890 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -3,17 +3,119 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import json +import os import subprocess import sys +from argparse import Namespace from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from olive.cli.base import TEST_OUTPUT_MARKER_FILE +from olive.cli.launcher import get_cli_parser from olive.cli.launcher import main as cli_main +def test_launcher_handles_commands_without_disable_telemetry(): + parser = MagicMock() + service = MagicMock() + telemetry = MagicMock() + parser.parse_known_args.return_value = (Namespace(func=lambda *_: service), []) + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.Telemetry") as mock_telemetry, + ): + mock_telemetry.get_or_create_if_enabled.return_value = telemetry + cli_main([]) + + service.run.assert_called_once() + telemetry.shutdown.assert_called_once() + + +def test_launcher_without_subcommand_does_not_initialize_telemetry(): + parser = MagicMock() + parser.parse_known_args.return_value = (Namespace(), []) + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.Telemetry") as mock_telemetry, + pytest.raises(SystemExit), + ): + cli_main([]) + + parser.print_help.assert_called_once() + mock_telemetry.assert_not_called() + + +def test_launcher_shuts_down_telemetry_on_command_failure(): + parser = MagicMock() + service = MagicMock() + telemetry = MagicMock() + service.run.side_effect = RuntimeError("boom") + parser.parse_known_args.return_value = ( + Namespace(func=lambda *_: service, disable_telemetry=False), + [], + ) + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.Telemetry") as mock_telemetry, + ): + mock_telemetry.get_or_create_if_enabled.return_value = telemetry + with pytest.raises(RuntimeError, match="boom"): + cli_main([]) + + telemetry.shutdown.assert_called_once() + + +def test_launcher_latches_opt_out_before_constructing_telemetry(monkeypatch): + monkeypatch.delenv("OLIVE_DISABLE_TELEMETRY", raising=False) + parser = MagicMock() + service = MagicMock() + parser.parse_known_args.side_effect = [ + (Namespace(func=lambda *_: service, disable_telemetry=True), []), + (Namespace(func=lambda *_: service, disable_telemetry=False), []), + ] + call_order = [] + observed_during_run = [] + disabled = False + service.run.side_effect = lambda: observed_during_run.append(os.environ.get("OLIVE_DISABLE_TELEMETRY")) + + def latch_disable(): + nonlocal disabled + disabled = True + call_order.append("disable") + + def get_or_create(): + if disabled: + return None + call_order.append("construct") + return MagicMock() + + with ( + patch("olive.cli.launcher.get_cli_parser", return_value=parser), + patch("olive.cli.launcher.disable_telemetry", side_effect=latch_disable), + patch("olive.cli.launcher.Telemetry.get_or_create_if_enabled", side_effect=get_or_create), + ): + cli_main([]) + cli_main([]) + + assert call_order == ["disable"] + assert observed_during_run == [None, None] + assert "OLIVE_DISABLE_TELEMETRY" not in os.environ + + +def test_init_command_parses_full_telemetry_opt_out(): + parser = get_cli_parser() + + args, unknown_args = parser.parse_known_args(["init", "--disable_telemetry"]) + + assert args.disable_telemetry is True + assert unknown_args == [] + + @pytest.mark.parametrize("console_script", [True, False]) @pytest.mark.parametrize( "command", @@ -108,7 +210,18 @@ def test_workflow_run_command(mock_run, tempdir, list_required_packages, tmp_pat # assert mock_run.assert_called_once_with( - {"key": "value"}, package_config=None, tempdir=tempdir, list_required_packages=list_required_packages + {"key": "value"}, + package_config=None, + tempdir=tempdir, + list_required_packages=list_required_packages, + recipe_telemetry_metadata={ + "recipe_command": "WorkflowRun", + "recipe_source": "config_file", + "recipe_format": "json", + "execution_mode": "list_required_packages" if list_required_packages else "run", + "package_config_provided": False, + }, + emit_error_telemetry=False, ) @@ -150,6 +263,23 @@ def test_workflow_run_command_with_overrides(mock_repo_exists, mock_run, tmp_pat list_required_packages=False, package_config=None, tempdir=None, + recipe_telemetry_metadata={ + "recipe_command": "WorkflowRun", + "recipe_source": "config_file", + "recipe_format": "json", + "execution_mode": "run", + "package_config_provided": False, + "config_overrides": { + "input_model": { + "type": "HfModel", + "model_path": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "load_kwargs": {"attn_implementation": "sdpa", "trust_remote_code": False}, + }, + "output_dir": str(Path("new_output_path").resolve()), + "log_severity_level": 2, + }, + }, + emit_error_telemetry=False, ) @@ -200,6 +330,14 @@ def test_workflow_run_command_with_test_override(mock_run, tmp_path): list_required_packages=False, package_config=None, tempdir=None, + recipe_telemetry_metadata={ + "recipe_command": "WorkflowRun", + "recipe_source": "config_file", + "recipe_format": "json", + "execution_mode": "run", + "package_config_provided": False, + }, + emit_error_telemetry=False, ) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index 7b492474ee..fd568b2a82 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: @@ -387,7 +463,7 @@ def _assert_discrepancy(self, tmp_path: Path): for exporter in self.exporters: if exporter == EXPORTER_MOBIUS and not _HAS_MOBIUS: self.fail( - "Requested exporter 'mobius' but mobius-ai is not installed. Install mobius-ai or remove '--exporter mobius'." + "Requested exporter 'mobius' but mobius-onnx is not installed. Install mobius-onnx or remove '--exporter mobius'." ) for model_id in self.model_ids: with self.subTest(model_id=model_id, exporter=exporter): diff --git a/test/cli/test_model_package.py b/test/cli/test_model_package.py index ec8ca9c944..c1d11a47ae 100644 --- a/test/cli/test_model_package.py +++ b/test/cli/test_model_package.py @@ -31,6 +31,13 @@ # --------------------------------------------------------------------------- +def _shared_asset_dir(package_root: Path) -> Path: + """Return the package's single ``shared_assets/sha256-/`` directory.""" + assets = sorted((package_root / "shared_assets").glob("sha256-*")) + assert len(assets) == 1, f"expected exactly one shared asset, found {[p.name for p in assets]}" + return assets[0] + + def _make_onnx_inline(onnx_path: Path, metadata_props: dict[str, str] | None = None) -> Path: """Write a minimal ONNX file with no external data.""" onnx_path.parent.mkdir(parents=True, exist_ok=True) @@ -243,17 +250,17 @@ def test_writes_proposal_layout(self, tmp_path): assert (out / "models").is_dir() manifest = json.loads((out / "manifest.json").read_text()) - assert manifest["schema_version"] == 1 - # ``decoder`` (not ``model``) — the genai_config role is ``decoder``, - # so _extract_task -> ``text_generation`` -> component dir ``decoder``. - assert manifest["components"] == ["decoder"] - assert manifest["producer"]["model_name"] == "test_model" - assert manifest["producer"]["model_version"] == "2.0" + assert manifest["schema_version"] == "1.0" + # One component holds every genai_config role, because ORT-GenAI + # selects a single component and loads one complete config from it. + assert manifest["components"] == {"model": "models/model"} + assert manifest["additional_metadata"]["producer"]["model_name"] == "test_model" + assert manifest["additional_metadata"]["producer"]["model_version"] == "2.0" # metadata uses inline EP - metadata = json.loads((out / "models" / "decoder" / "metadata.json").read_text()) - assert metadata["schema_version"] == 1 - assert metadata["component_name"] == "decoder" + metadata = json.loads((out / "models" / "model" / "component.json").read_text()) + assert "schema_version" not in metadata + assert metadata["component_name"] == "model" assert set(metadata["variants"]) == {"soc_60", "soc_73"} for variant_payload in metadata["variants"].values(): assert variant_payload == {"ep": "QNNExecutionProvider"} @@ -261,8 +268,8 @@ def test_writes_proposal_layout(self, tmp_path): # No variant.json is emitted; the ONNX file lands in the variant # directory. for v in ("soc_60", "soc_73"): - assert not (out / "models" / "decoder" / v / "variant.json").exists() - assert (out / "models" / "decoder" / v / "model.onnx").is_file() + assert not (out / "models" / "model" / v / "variant.json").exists() + assert (out / "models" / "model" / v / "model.onnx").is_file() class TestGeneratePackageSingleSource: @@ -274,12 +281,12 @@ def test_single_source_is_valid_package(self, tmp_path): cmd.run() manifest = json.loads((out / "manifest.json").read_text()) - assert manifest["components"] == ["decoder"] - metadata = json.loads((out / "models" / "decoder" / "metadata.json").read_text()) + assert manifest["components"] == {"model": "models/model"} + metadata = json.loads((out / "models" / "model" / "component.json").read_text()) assert "cpu_x64" in metadata["variants"] assert metadata["variants"]["cpu_x64"] == {"ep": "CPUExecutionProvider"} # No shared_weights because nothing to dedup. - assert not (out / "models" / "decoder" / "shared_weights").exists() + assert not (out / "models" / "model" / "shared_weights").exists() # --------------------------------------------------------------------------- @@ -307,7 +314,7 @@ def test_writes_proposal_shape_for_single_variant(self, tmp_path): ) assert (out / "manifest.json").is_file() - assert (out / "models" / "decoder" / "metadata.json").is_file() + assert (out / "models" / "decoder" / "component.json").is_file() # No variant.json is emitted. assert not (out / "models" / "decoder" / "cpu" / "variant.json").exists() assert (out / "models" / "decoder" / "cpu" / "model.onnx").is_file() @@ -330,16 +337,27 @@ def test_manifest_uses_proposal_schema(self, tmp_path): ) manifest = json.loads((out / "manifest.json").read_text()) - assert manifest["schema_version"] == 1 - assert manifest["components"] == ["decoder"] + assert manifest["schema_version"] == "1.0" + assert manifest["components"] == {"decoder": "models/decoder"} assert manifest["package_name"] == "package" assert manifest["package_version"] == "1.0" - assert manifest["configs_dir"] == "configs" - assert manifest["producer"] == { + assert manifest["additional_metadata"]["producer"] == { "tool": "olive-ai", "tool_version": "1.2.3", "model_name": "demo", } + # The ORT schema rejects unknown top-level manifest keys, so nothing + # outside its vocabulary may be emitted. + assert set(manifest) <= { + "schema_version", + "package_name", + "package_version", + "description", + "layout", + "components", + "shared_assets", + "additional_metadata", + } # No legacy fields assert "name" not in manifest assert "component_models" not in manifest @@ -363,8 +381,8 @@ def test_metadata_uses_inline_ep(self, tmp_path): ], ) - metadata = json.loads((out / "models" / "decoder" / "metadata.json").read_text()) - assert metadata["schema_version"] == 1 + metadata = json.loads((out / "models" / "decoder" / "component.json").read_text()) + assert "schema_version" not in metadata assert metadata["component_name"] == "decoder" assert metadata["variants"]["qnn-npu"] == { "ep": "QNNExecutionProvider", @@ -389,7 +407,7 @@ def test_metadata_omits_optional_fields_when_unset(self, tmp_path): ], ) - metadata = json.loads((out / "models" / "decoder" / "metadata.json").read_text()) + metadata = json.loads((out / "models" / "decoder" / "component.json").read_text()) assert metadata["variants"]["cpu"] == {"ep": "CPUExecutionProvider"} def test_overlay_carries_session_and_provider_options(self, tmp_path): @@ -415,10 +433,10 @@ def test_overlay_carries_session_and_provider_options(self, tmp_path): ], ) - # Runtime fields go to genai_config_overlay.json, not variant.json. + # Runtime fields live in the variant's own genai_config.json. assert not (out / "models" / "decoder" / "cuda" / "variant.json").exists() - overlay = json.loads((out / "models" / "decoder" / "cuda" / "genai_config_overlay.json").read_text()) - assert overlay == { + config = json.loads((out / "models" / "decoder" / "cuda" / "genai_config.json").read_text()) + assert config == { "model": { "decoder": { "filename": "model.onnx", @@ -454,8 +472,8 @@ def test_overlay_provider_options_match_ep_by_name(self, tmp_path): ], ) - overlay = json.loads((out / "models" / "decoder" / "qnn" / "genai_config_overlay.json").read_text()) - assert overlay["model"]["decoder"]["session_options"]["provider_options"] == [ + config = json.loads((out / "models" / "decoder" / "qnn" / "genai_config.json").read_text()) + assert config["model"]["decoder"]["session_options"]["provider_options"] == [ {"qnn": {"backend_path": "QnnHtp.so"}} ] @@ -484,8 +502,8 @@ def test_overlay_emits_empty_provider_options_for_cpu(self, tmp_path): ], ) - overlay = json.loads((out / "models" / "decoder" / "cpu" / "genai_config_overlay.json").read_text()) - assert overlay == { + config = json.loads((out / "models" / "decoder" / "cpu" / "genai_config.json").read_text()) + assert config == { "model": { "decoder": { "filename": "model.onnx", @@ -532,7 +550,7 @@ def test_overlay_lifts_per_variant_model_level_fields(self, tmp_path): ], ) - overlay = json.loads((out / "models" / "decoder" / "npu" / "genai_config_overlay.json").read_text()) + overlay = json.loads((out / "models" / "decoder" / "npu" / "genai_config.json").read_text()) model_patch = overlay["model"] assert model_patch["context_length"] == 4224 assert model_patch["pad_token_id"] == 200020 @@ -544,15 +562,14 @@ def test_overlay_lifts_per_variant_model_level_fields(self, tmp_path): # the overlay — otherwise it would duplicate the base copy. assert "vocab_size" not in model_patch - def test_base_genai_strips_per_variant_model_fields(self, tmp_path): - """The base ``configs/genai_config.json`` must not carry per-variant fields. + def test_variant_config_is_complete_and_self_contained(self, tmp_path): + """Each variant carries a full ``genai_config.json``, not a merge patch. - If ``context_length`` (or similar) lived in the base, GenAI's overlay - merge would still honor the per-variant value (overlay scalar wins), - but ``_VARIANT_LEVEL_MODEL_KEYS`` includes arrays (``eos_token_id``) - whose presence in the base would trigger GenAI's array-append merge - semantics — the merged result would duplicate the array. So the base - must be free of every variant-level model key. + ORT-GenAI loads ``/genai_config.json`` directly + and never merges a package-level base, so every field it needs — + structural (``vocab_size``) and per-variant (``context_length``, + ``eos_token_id``, ``filename``) alike — must be present in that one + file. """ onnx_path = _make_onnx_inline(tmp_path / "src" / "model.onnx") out = tmp_path / "package" @@ -591,16 +608,23 @@ def test_base_genai_strips_per_variant_model_fields(self, tmp_path): config_files={"genai_config.json": cfg}, ) - base = json.loads((out / "configs" / "genai_config.json").read_text()) - model = base["model"] - for stripped in ("context_length", "pad_token_id", "eos_token_id", "bos_token_id", "type"): - assert stripped not in model, f"base genai_config must not contain {stripped!r}" - # Variant-specific decoder fields also stripped. - assert "filename" not in model["decoder"] - assert "session_options" not in model["decoder"] - # Structural shared fields remain. + config = json.loads((out / "models" / "decoder" / "cpu" / "genai_config.json").read_text()) + model = config["model"] + # The variant config is complete: variant-level keys come from this + # variant's own source rather than leaking in from a shared base. + assert model["context_length"] == 131072 + assert model["pad_token_id"] == 199999 + assert model["eos_token_id"] == [200020, 199999] + assert model["bos_token_id"] == 199999 + assert model["type"] == "phi3" + assert model["decoder"]["filename"] == "model.onnx" + # Structural shared fields survive the base strip. assert model["vocab_size"] == 200064 assert model["decoder"]["head_size"] == 128 + # No package-level base config is emitted; ORT-GenAI only ever reads + # the selected variant's genai_config.json. + assert not (out / "configs").exists() + assert not (out / "genai_config.json").exists() # --------------------------------------------------------------------------- @@ -765,13 +789,52 @@ def test_sidecar_sweep_does_not_overwrite_external_data(self, tmp_path): class TestConfigsAndSafety: - def test_copies_config_files_into_configs_dir(self, tmp_path): + def test_shares_tokenizer_assets_by_content_address(self, tmp_path): onnx_path = _make_onnx_inline(tmp_path / "src" / "model.onnx") cfg_a = tmp_path / "configs_src" / "tokenizer.json" cfg_a.parent.mkdir(parents=True) cfg_a.write_text("{}") cfg_b = tmp_path / "configs_src" / "genai_config.json" cfg_b.write_text("{}") + cfg_c = tmp_path / "configs_src" / "tokenizer_config.json" + cfg_c.write_text("{}") + out = tmp_path / "package" + + write_model_package( + output_dir=out, + variants=[ + VariantSpec( + component_name="decoder", + variant_name="cpu", + onnx_files=[onnx_path], + ep="CPUExecutionProvider", + ) + ], + config_files={"tokenizer.json": cfg_a, "tokenizer_config.json": cfg_c, "genai_config.json": cfg_b}, + ) + + asset_dir = _shared_asset_dir(out) + assert (asset_dir / "tokenizer.json").is_file() + assert (asset_dir / "tokenizer_config.json").is_file() + # genai_config.json is the base each variant config is derived from, + # never a shared asset of its own. + assert not (asset_dir / "genai_config.json").exists() + assert not (out / "configs").exists() + + config = json.loads((out / "models" / "decoder" / "cpu" / "genai_config.json").read_text()) + assert config["model"]["tokenizer_dir"] == f"sha256:{asset_dir.name.removeprefix('sha256-')}" + + def test_no_tokenizer_dir_without_tokenizer_config(self, tmp_path): + """Shared assets that hold no tokenizer must not be advertised as one. + + onnxruntime-extensions unconditionally opens + ``/tokenizer_config.json``, so pointing ``tokenizer_dir`` + at an asset without it cannot help and only misreports where the + tokenizer was expected. + """ + onnx_path = _make_onnx_inline(tmp_path / "src" / "model.onnx") + stray = tmp_path / "src" / "notes.txt" + stray.write_text("not a tokenizer") out = tmp_path / "package" write_model_package( @@ -784,11 +847,13 @@ def test_copies_config_files_into_configs_dir(self, tmp_path): ep="CPUExecutionProvider", ) ], - config_files={"tokenizer.json": cfg_a, "genai_config.json": cfg_b}, + config_files={"notes.txt": stray}, ) - assert (out / "configs" / "tokenizer.json").is_file() - assert (out / "configs" / "genai_config.json").is_file() + # The file is still staged; only the tokenizer_dir pointer is withheld. + assert (_shared_asset_dir(out) / "notes.txt").is_file() + config = json.loads((out / "models" / "decoder" / "cpu" / "genai_config.json").read_text()) + assert "tokenizer_dir" not in config.get("model", {}) def test_rejects_non_empty_output_dir(self, tmp_path): onnx_path = _make_onnx_inline(tmp_path / "src" / "model.onnx") @@ -893,12 +958,12 @@ def test_skips_config_file_with_unsafe_key(self, tmp_path): ) # assert: unsafe keys are dropped, safe key copied + asset_dir = _shared_asset_dir(out) assert not (out.parent / "escape.txt").exists() - assert not (out / "configs" / "subdir").exists() - assert not (out / "configs" / "..").is_dir() or not (out / ".." / "escape.txt").exists() - assert (out / "configs" / "ok.txt").exists() - # configs/ should contain only the one safe entry - assert sorted(p.name for p in (out / "configs").iterdir()) == ["ok.txt"] + assert not (asset_dir / "subdir").exists() + assert (asset_dir / "ok.txt").exists() + # the shared asset should contain only the one safe entry + assert sorted(p.name for p in asset_dir.iterdir()) == ["ok.txt"] # --------------------------------------------------------------------------- @@ -952,7 +1017,7 @@ def test_passes_through_comma_delimited_metadata(self, tmp_path): cmd.run() # assert: compatibility_string passes the raw opaque string through verbatim - metadata = json.loads((out / "models" / "decoder" / "metadata.json").read_text()) + metadata = json.loads((out / "models" / "model" / "component.json").read_text()) variant = metadata["variants"]["soc_60"] assert variant["ep"] == "QNNExecutionProvider" assert variant["compatibility_string"] == "soc_60,soc_69,soc_73" @@ -1090,7 +1155,7 @@ def test_packs_pipeline_with_all_stage_onnx_files(self, tmp_path): cmd.run() - variant_dir = out.with_suffix(".ortpackage") / "models" / "decoder" / "qnn_npu" + variant_dir = out.with_suffix(".ortpackage") / "models" / "model" / "qnn_npu" assert variant_dir.is_dir() for fname in stage_files: assert (variant_dir / fname).is_file(), f"missing stage file {fname}" @@ -1116,7 +1181,7 @@ def test_pipeline_overlay_lifts_full_stage_structure_from_source(self, tmp_path) cmd.run() - overlay_path = out.with_suffix(".ortpackage") / "models" / "decoder" / "qnn_npu" / "genai_config_overlay.json" + overlay_path = out.with_suffix(".ortpackage") / "models" / "model" / "qnn_npu" / "genai_config.json" overlay = json.loads(overlay_path.read_text()) decoder = overlay["model"]["decoder"] assert "pipeline" in decoder @@ -1130,13 +1195,11 @@ def test_pipeline_overlay_lifts_full_stage_structure_from_source(self, tmp_path) # decoder-level session_options also lifted from source so log_id etc. survive. assert decoder["session_options"]["log_id"] == "onnxruntime-genai" - def test_base_genai_strips_pipeline_field(self, tmp_path): - """``pipeline`` lives only in the overlay; base must not duplicate it. + def test_variant_config_carries_pipeline_exactly_once(self, tmp_path): + """The variant config holds the pipeline array exactly once. - GenAI's overlay parser appends arrays rather than replacing them - (``src/config.cpp:PipelineModelObject_Element``), so a ``pipeline`` - in both base and overlay would double every stage. The strip is the - guard. + The shared base strips ``pipeline`` and each variant re-applies its + own, so the stage list can never be duplicated by a merge. """ src = _create_pipeline_source( tmp_path, @@ -1151,9 +1214,12 @@ def test_base_genai_strips_pipeline_field(self, tmp_path): cmd.run() - base = json.loads((out.with_suffix(".ortpackage") / "configs" / "genai_config.json").read_text()) - decoder = base["model"]["decoder"] - assert "pipeline" not in decoder, "base genai_config must not retain the pipeline array" + pkg = out.with_suffix(".ortpackage") + config = json.loads((pkg / "models" / "model" / "qnn_npu" / "genai_config.json").read_text()) + decoder = config["model"]["decoder"] + stage_files = [next(iter(stage.values()))["filename"] for stage in decoder["pipeline"]] + assert stage_files == ["e.onnx", "c.onnx", "i.onnx", "h.onnx"] + assert not (pkg / "configs").exists() def test_flat_source_ep_derived_from_source_genai_when_attrs_missing(self, tmp_path): """For flat sources, source genai's ``provider_options`` overrules name guess. @@ -1186,12 +1252,10 @@ def test_flat_source_ep_derived_from_source_genai_when_attrs_missing(self, tmp_p cmd.run() - metadata = json.loads((out.with_suffix(".ortpackage") / "models" / "decoder" / "metadata.json").read_text()) + metadata = json.loads((out.with_suffix(".ortpackage") / "models" / "model" / "component.json").read_text()) assert metadata["variants"]["vitia_npu"]["ep"] == "VitisAIExecutionProvider" overlay = json.loads( - ( - out.with_suffix(".ortpackage") / "models" / "decoder" / "vitia_npu" / "genai_config_overlay.json" - ).read_text() + (out.with_suffix(".ortpackage") / "models" / "model" / "vitia_npu" / "genai_config.json").read_text() ) assert overlay["model"]["decoder"]["session_options"]["provider_options"] == [{"VitisAI": {}}] @@ -1245,33 +1309,37 @@ class TestVLMMultiRoleOverlay: """Multi-role (vision + embedding + decoder) VLM packaging. A flat VLM source dir packs >1 ONNX file referenced by >1 role in the - same ``genai_config.json``. Each role becomes its own component - (``models/vision/``, ``models/embedding/``, ``models/decoder/``); each - component's overlay restores only that role's ``filename`` / - ``session_options``. The base genai_config strips every role's - filename / session_options and injects a ``component=`` marker - so the loader can map each role back to its on-disk directory. + same ``genai_config.json``. All of those roles land in ONE component + (``models/model/``) because ORT-GenAI selects exactly one component + per package and then resolves every role's ``filename`` against that + single selected variant directory (``src/models/multi_modal.cpp`` + builds the vision / embedding / decoder sessions from one + ``config_path``). Splitting roles across components would produce a + package ORT-GenAI refuses to open ("declares N components; + onnxruntime-genai requires exactly one"). """ - def test_each_role_becomes_its_own_component(self, tmp_path): + def test_all_roles_share_one_component(self, tmp_path): src = _create_vlm_source(tmp_path, "cpu_and_mobile") out = tmp_path / "out" cmd = _make_command(["generate-model-package", "-s", str(src), "-o", str(out)]) cmd.run() - models_dir = out.with_suffix(".ortpackage") / "models" - assert (models_dir / "vision" / "metadata.json").is_file() - assert (models_dir / "embedding" / "metadata.json").is_file() - assert (models_dir / "decoder" / "metadata.json").is_file() + pkg = out.with_suffix(".ortpackage") + manifest = json.loads((pkg / "manifest.json").read_text()) + assert manifest["components"] == {"model": "models/model"} + assert (pkg / "models" / "model" / "component.json").is_file() + # No per-role components: those would make the package unloadable. + for role in ("vision", "embedding", "decoder"): + assert not (pkg / "models" / role).exists(), f"role {role} leaked into its own component" - def test_each_components_overlay_lifts_only_its_role(self, tmp_path): - """Each per-role overlay carries exactly that role's filename, no others. + def test_single_config_restores_every_role_filename(self, tmp_path): + """The one variant config restores EVERY role's filename. - The base config strips every role's filename, but each per-role - overlay should only restore its own — duplicating the lift across - all overlays would corrupt the loader's view (and trigger array - append-merge problems for list-valued scalars). + The shared base strips per-role ``filename`` / ``session_options``; + the variant re-applies all of them, because the single component + owns every role and GenAI reads one complete config from it. """ src = _create_vlm_source(tmp_path, "cpu_and_mobile") out = tmp_path / "out" @@ -1280,33 +1348,21 @@ def test_each_components_overlay_lifts_only_its_role(self, tmp_path): cmd.run() models = out.with_suffix(".ortpackage") / "models" - decoder_overlay = json.loads((models / "decoder" / "cpu_and_mobile" / "genai_config_overlay.json").read_text()) - vision_overlay = json.loads((models / "vision" / "cpu_and_mobile" / "genai_config_overlay.json").read_text()) - embedding_overlay = json.loads( - (models / "embedding" / "cpu_and_mobile" / "genai_config_overlay.json").read_text() - ) + cfg = json.loads((models / "model" / "cpu_and_mobile" / "genai_config.json").read_text()) - assert "decoder" in decoder_overlay["model"] - assert decoder_overlay["model"]["decoder"]["filename"] == "text.onnx" - assert "vision" not in decoder_overlay["model"] - assert "embedding" not in decoder_overlay["model"] + with_filename = {name for name, body in cfg["model"].items() if isinstance(body, dict) and "filename" in body} + assert with_filename == {"vision", "embedding", "decoder"} - assert "vision" in vision_overlay["model"] - assert vision_overlay["model"]["vision"]["filename"] == "vision.onnx" - assert "decoder" not in vision_overlay["model"] - assert "embedding" not in vision_overlay["model"] + assert cfg["model"]["decoder"]["filename"] == "text.onnx" + assert cfg["model"]["vision"]["filename"] == "vision.onnx" + assert cfg["model"]["embedding"]["filename"] == "embedding.onnx" - assert "embedding" in embedding_overlay["model"] - assert embedding_overlay["model"]["embedding"]["filename"] == "embedding.onnx" - assert "decoder" not in embedding_overlay["model"] - assert "vision" not in embedding_overlay["model"] + def test_variant_dir_holds_every_roles_onnx(self, tmp_path): + """The single variant dir holds all three roles' ONNX files, flat. - def test_variant_dirs_are_flat_one_onnx_per_role(self, tmp_path): - """Each per-role variant dir holds exactly its own ONNX, flat. - - With one role per component there's no sibling-role disambiguation - to do, so the writer drops any subdir prefixes from the source - layout and writes each ONNX at the variant root. + GenAI resolves each role's ``filename`` relative to the selected + variant directory, so every role's graph must be reachable from + that one directory. """ src = _create_vlm_source(tmp_path, "cpu_and_mobile") out = tmp_path / "out" @@ -1314,18 +1370,18 @@ def test_variant_dirs_are_flat_one_onnx_per_role(self, tmp_path): cmd.run() - models = out.with_suffix(".ortpackage") / "models" - assert (models / "decoder" / "cpu_and_mobile" / "text.onnx").is_file() - assert not (models / "decoder" / "cpu_and_mobile" / "vision.onnx").exists() - assert (models / "vision" / "cpu_and_mobile" / "vision.onnx").is_file() - assert (models / "embedding" / "cpu_and_mobile" / "embedding.onnx").is_file() + variant = out.with_suffix(".ortpackage") / "models" / "model" / "cpu_and_mobile" + assert (variant / "text.onnx").is_file() + assert (variant / "vision.onnx").is_file() + assert (variant / "embedding.onnx").is_file() - def test_base_genai_injects_component_marker_for_every_role(self, tmp_path): - """Every role gets a ``component=`` marker in the base config. + def test_variant_config_omits_unknown_component_marker(self, tmp_path): + """No ``component`` marker is emitted into any role block. - The merged config the loader sees must know which component - directory each role lives in. With per-role components the - component name equals the role name. + ORT-GenAI's config parser rejects unknown fields + (``src/config.cpp`` throws ``JSON::unknown_value_error``) and has no + ``component`` field, so emitting one would make the package + unloadable. """ src = _create_vlm_source(tmp_path, "cpu_and_mobile") out = tmp_path / "out" @@ -1333,10 +1389,48 @@ def test_base_genai_injects_component_marker_for_every_role(self, tmp_path): cmd.run() - base = json.loads((out.with_suffix(".ortpackage") / "configs" / "genai_config.json").read_text()) - model = base["model"] - for role in ("vision", "embedding", "decoder"): - assert model[role]["component"] == role, f"role {role} missing self-named component marker" + config = json.loads( + (out.with_suffix(".ortpackage") / "models" / "model" / "cpu_and_mobile" / "genai_config.json").read_text() + ) + for name, body in config["model"].items(): + if isinstance(body, dict): + assert "component" not in body, f"config leaked a component marker under {name!r}" + + def test_conflicting_non_cpu_eps_across_roles_raises(self, tmp_path): + """Two roles demanding different non-CPU EPs is rejected up front. + + All roles share one component and therefore one variant EP. + ORT-GenAI additionally has a single model-level device + (``src/models/model.cpp``: "Running a model with multiple + providers is not supported"), so such a package could never load. + Failing during packaging gives a far clearer diagnostic. + """ + src = tmp_path / "mixed" + src.mkdir() + for fname in ("vision.onnx", "text.onnx"): + _make_onnx_inline(src / fname) + (src / "genai_config.json").write_text( + json.dumps( + { + "model": { + "type": "qwen3vl", + "vision": { + "filename": "vision.onnx", + "session_options": {"provider_options": [{"qnn": {}}]}, + }, + "decoder": { + "head_size": 128, + "filename": "text.onnx", + "session_options": {"provider_options": [{"cuda": {}}]}, + }, + } + } + ) + ) + cmd = _make_command(["generate-model-package", "-s", str(src), "-o", str(tmp_path / "out")]) + + with pytest.raises(ValueError, match="more than one non-CPU execution provider"): + cmd.run() # --------------------------------------------------------------------------- @@ -1361,13 +1455,13 @@ def _create_mobius_vlm_source( references each role by its full subdirectory-prefixed path (``"filename": "decoder/model.onnx"``). - With per-role components each role gets its own - ``models///`` directory, so the packager safely - flattens the packaged filename to the basename (``model.onnx``) and - rewrites the overlay to match the on-disk layout — no sibling-role - disambiguation is needed inside any one variant dir. The source-side - subdirectory prefix is used only to locate the file on disk in the - source and does not propagate into the package. + All roles land in ONE package component and therefore share one + variant directory, so the packager preserves each role's source-side + subdirectory prefix verbatim. Flattening to the basename would make + all three roles collide at ``/model.onnx``; keeping the + prefix is inherently collision-free (the files already coexisted in + the source) and means the genai_config ``filename`` values need no + rewriting at all. """ source_dir = tmp_path / name source_dir.mkdir(parents=True) @@ -1433,54 +1527,52 @@ def _create_mobius_vlm_source( class TestMobiusHierarchicalLayout: - """End-to-end packaging of Mobius-style multi-component VLM sources. - - Each Mobius role (``decoder``/``embedding``/``vision``) becomes its - own component in the package (``models/decoder/``, - ``models/embedding/``, ``models/vision/``). Per the ORT model-package - proposal, ``models//`` is the top-level grouping where one - component == one inference session. Each component's variant - directory is flat: the source-side subdir prefix is dropped because - there's no longer a sibling role to disambiguate against. + """End-to-end packaging of Mobius-style multi-role VLM sources. + + Every Mobius role (``decoder``/``embedding``/``vision``) goes into the + single package component ``models/model/``, because ORT-GenAI opens + exactly one component and resolves all of a model's role filenames + against that component's selected variant directory. The variant + directory therefore mirrors the source layout, keeping each role's + subdirectory (``decoder/model.onnx`` etc.) so same-named graphs don't + collide. """ - def test_each_role_becomes_top_level_component(self, tmp_path): + def test_all_roles_share_one_component(self, tmp_path): src = _create_mobius_vlm_source(tmp_path, "cpu") out = tmp_path / "out" cmd = _make_command(["generate-model-package", "-s", str(src), "-o", str(out)]) cmd.run() - models = out.with_suffix(".ortpackage") / "models" + pkg = out.with_suffix(".ortpackage") + manifest = json.loads((pkg / "manifest.json").read_text()) + assert manifest["components"] == {"model": "models/model"} + assert (pkg / "models" / "model" / "component.json").is_file() for role in ("decoder", "embedding", "vision"): - assert (models / role / "metadata.json").is_file(), f"missing component dir for role {role}" + assert not (pkg / "models" / role).exists(), f"role {role} leaked into its own component" - def test_variant_dir_is_flat_under_each_component(self, tmp_path): + def test_variant_dir_preserves_source_role_subdirs(self, tmp_path): src = _create_mobius_vlm_source(tmp_path, "cpu") out = tmp_path / "out" cmd = _make_command(["generate-model-package", "-s", str(src), "-o", str(out)]) cmd.run() - models = out.with_suffix(".ortpackage") / "models" - # Each per-role variant dir holds the ONNX flat (basename only). - for role in ("decoder", "embedding", "vision"): - assert (models / role / "cpu" / "model.onnx").is_file(), ( - f"missing flat model.onnx under models/{role}/cpu (writer kept subdir?)" - ) - # The source-side subdir should NOT propagate into the - # variant dir — there's only one role here, no sibling to - # disambiguate against. - assert not (models / role / "cpu" / role / "model.onnx").exists() + variant = out.with_suffix(".ortpackage") / "models" / "model" / "cpu" + # Sibling roles share one variant dir, so the source-side subdir + # must survive — flattening to the basename would collide. + for subdir in ("decoder", "embedding", "vision_encoder"): + assert (variant / subdir / "model.onnx").is_file(), f"missing {subdir}/model.onnx under the variant dir" - def test_external_data_lands_next_to_its_onnx_in_flat_layout(self, tmp_path): - """Each role's external-data blob lives flat next to its ONNX in that role's variant dir. + def test_external_data_lands_next_to_its_onnx_in_role_subdir(self, tmp_path): + """Each role's external-data blob lives next to its own ONNX. The ONNX file references ``model.onnx.data`` relative to its own - directory; the loader resolves the same way. With one role per - component each role's blob has its own dedicated directory, so - no collision is possible even when sibling source-side roles - shared the same basename. + directory and the loader resolves the same way, so the blob must + follow the ONNX into its role subdirectory. Routing blobs to the + variant root would collide all three at + ``/model.onnx.data``. """ src = _create_mobius_vlm_source(tmp_path, "cpu", with_external_data=True) out = tmp_path / "out" @@ -1488,18 +1580,23 @@ def test_external_data_lands_next_to_its_onnx_in_flat_layout(self, tmp_path): cmd.run() - models = out.with_suffix(".ortpackage") / "models" - for role in ("decoder", "embedding", "vision"): - blob = models / role / "cpu" / "model.onnx.data" - assert blob.is_file(), f"external-data blob missing under models/{role}/cpu/" - - def test_overlay_filename_is_basename(self, tmp_path): - """The overlay's ``filename`` is the basename, not the source subdir-prefixed path. - - Per-role variant dirs are flat; the overlay must match the on-disk - layout. The source-side ``decoder/model.onnx`` prefix is purely a - sibling-disambiguation device that no longer applies once the - roles are separated into top-level components. + variant = out.with_suffix(".ortpackage") / "models" / "model" / "cpu" + seen = set() + for subdir, role in (("decoder", "decoder"), ("embedding", "embedding"), ("vision_encoder", "vision_encoder")): + blob = variant / subdir / "model.onnx.data" + assert blob.is_file(), f"external-data blob missing under {subdir}/" + payload = blob.read_bytes() + assert payload == f"role-{role}".encode() * 16, f"{subdir} blob holds another role's weights" + seen.add(payload) + assert len(seen) == 3, "role blobs were deduped/overwritten into each other" + + def test_overlay_filename_keeps_role_subdir_prefix(self, tmp_path): + """The overlay's ``filename`` keeps the source's subdir-prefixed path. + + All roles share one variant dir, whose layout mirrors the source, + so the packaged filenames are byte-identical to the source's. That + also means GenAI's ``config_path / filename`` resolution finds each + graph without any rewriting. """ src = _create_mobius_vlm_source(tmp_path, "cpu") out = tmp_path / "out" @@ -1507,20 +1604,19 @@ def test_overlay_filename_is_basename(self, tmp_path): cmd.run() - models = out.with_suffix(".ortpackage") / "models" - decoder = json.loads((models / "decoder" / "cpu" / "genai_config_overlay.json").read_text()) - embedding = json.loads((models / "embedding" / "cpu" / "genai_config_overlay.json").read_text()) - vision = json.loads((models / "vision" / "cpu" / "genai_config_overlay.json").read_text()) - assert decoder["model"]["decoder"]["filename"] == "model.onnx" - assert embedding["model"]["embedding"]["filename"] == "model.onnx" - assert vision["model"]["vision"]["filename"] == "model.onnx" + model = json.loads( + (out.with_suffix(".ortpackage") / "models" / "model" / "cpu" / "genai_config.json").read_text() + )["model"] + assert model["decoder"]["filename"] == "decoder/model.onnx" + assert model["embedding"]["filename"] == "embedding/model.onnx" + assert model["vision"]["filename"] == "vision_encoder/model.onnx" - def test_configs_dir_excludes_model_artifact_subdirs(self, tmp_path): - """``decoder/``/``embedding/``/``vision_encoder/`` must not leak into ``configs/``. + def test_shared_assets_exclude_model_artifact_subdirs(self, tmp_path): + """``decoder/``/``embedding/``/``vision_encoder/`` must not leak into shared assets. Without explicit exclusion the config-file sweep would copy every source-root directory (including the model-artifact subdirs), so the - package would carry duplicate ONNXs under ``configs/`` and bloat + package would carry duplicate ONNXs in the shared asset and bloat the deliverable. The sweep recognizes model-artifact subdirs via the genai_config's role filenames and skips them. """ @@ -1530,18 +1626,22 @@ def test_configs_dir_excludes_model_artifact_subdirs(self, tmp_path): cmd.run() - configs_dir = out.with_suffix(".ortpackage") / "configs" + pkg = out.with_suffix(".ortpackage") + asset_dir = _shared_asset_dir(pkg) for excluded in ("decoder", "embedding", "vision_encoder"): - assert not (configs_dir / excluded).exists(), f"{excluded}/ leaked into configs/" - assert (configs_dir / "tokenizer_config.json").is_file() - assert (configs_dir / "model_config.json").is_file() - assert (configs_dir / "genai_config.json").is_file() - - def test_base_genai_strips_filename_and_marks_self_named_components(self, tmp_path): - """Base genai_config strips per-role ``filename`` and injects a self-named ``component`` marker. - - With per-role components the role name equals the component name, - so every role's ``component`` field is its own name. + assert not (asset_dir / excluded).exists(), f"{excluded}/ leaked into the shared asset" + assert (asset_dir / "tokenizer_config.json").is_file() + assert (asset_dir / "model_config.json").is_file() + # genai_config.json is never a shared asset: each variant owns a copy. + assert not (asset_dir / "genai_config.json").exists() + assert not (pkg / "configs").exists() + + def test_variant_config_restores_every_filename_and_omits_component_marker(self, tmp_path): + """The variant config restores every role's ``filename`` and adds no marker. + + The shared base strips per-role ``filename``; the single variant + re-applies all of them. No ``component`` field is emitted because + ORT-GenAI's config parser rejects unknown fields. """ src = _create_mobius_vlm_source(tmp_path, "cpu") out = tmp_path / "out" @@ -1549,20 +1649,23 @@ def test_base_genai_strips_filename_and_marks_self_named_components(self, tmp_pa cmd.run() - base = json.loads((out.with_suffix(".ortpackage") / "configs" / "genai_config.json").read_text()) - model = base["model"] - for role in ("decoder", "embedding", "vision"): - assert "filename" not in model[role], f"{role}.filename should be stripped from base" - assert model[role]["component"] == role, f"{role} component marker should equal role name" + model = json.loads( + (out.with_suffix(".ortpackage") / "models" / "model" / "cpu" / "genai_config.json").read_text() + )["model"] + with_filename = {name for name, body in model.items() if isinstance(body, dict) and "filename" in body} + assert with_filename == {"decoder", "embedding", "vision"} + for name, body in model.items(): + if isinstance(body, dict): + assert "component" not in body, f"config leaked a component marker under {name!r}" - def test_two_sources_each_produce_per_role_variants(self, tmp_path): - """CPU + GPU Mobius sources both contribute one variant per role to each component. + def test_two_sources_become_two_variants_of_one_component(self, tmp_path): + """CPU + GPU Mobius sources contribute one variant each to the single component. This is the user-reported scenario: - ``olive generate-model-package -s cpu -s gpu -o cpu_gpu``. Each - role component (``decoder``/``embedding``/``vision``) ends up - with two variants — ``cpu`` (CPUExecutionProvider) and ``gpu`` - (CUDAExecutionProvider). + ``olive generate-model-package -s cpu -s gpu -o cpu_gpu``. The + package holds one component with two variants — ``cpu`` + (CPUExecutionProvider) and ``gpu`` (CUDAExecutionProvider) — each + carrying the full three-role model. """ cpu = _create_mobius_vlm_source(tmp_path, "cpu", ep="CPUExecutionProvider") gpu = _create_mobius_vlm_source(tmp_path, "gpu", ep="CUDAExecutionProvider") @@ -1571,48 +1674,47 @@ def test_two_sources_each_produce_per_role_variants(self, tmp_path): cmd.run() - models = out.with_suffix(".ortpackage") / "models" - for role in ("decoder", "embedding", "vision"): - for variant in ("cpu", "gpu"): - assert (models / role / variant / "model.onnx").is_file(), f"missing models/{role}/{variant}/model.onnx" - metadata = json.loads((models / role / "metadata.json").read_text()) - assert metadata["variants"]["cpu"]["ep"] == "CPUExecutionProvider" - assert metadata["variants"]["gpu"]["ep"] == "CUDAExecutionProvider" - - def test_variant_level_scalars_lift_only_into_primary_role_overlay(self, tmp_path): - """Variant-level scalars (eos_token_id, context_length, ...) appear in exactly one overlay. - - GenAI's overlay parser append-merges arrays. If - ``eos_token_id`` (often a list) ended up in three different - per-role overlays the merged config would triple every entry. - Only the primary role per source (``_pick_primary_role`` — - ``decoder`` here) carries these scalars. + component = out.with_suffix(".ortpackage") / "models" / "model" + for variant in ("cpu", "gpu"): + for subdir in ("decoder", "embedding", "vision_encoder"): + assert (component / variant / subdir / "model.onnx").is_file(), ( + f"missing models/model/{variant}/{subdir}/model.onnx" + ) + metadata = json.loads((component / "component.json").read_text()) + assert metadata["variants"]["cpu"]["ep"] == "CPUExecutionProvider" + assert metadata["variants"]["gpu"]["ep"] == "CUDAExecutionProvider" + + def test_variant_level_scalars_present_in_every_variant_config(self, tmp_path): + """Variant-level scalars appear in every variant config, at their exact value. + + Each variant config is standalone, so it must carry the model-level + scalars in full. Because the writer replaces rather than appends, + a list-valued ``eos_token_id`` keeps exactly the source's entries + instead of accumulating copies. """ - src = _create_mobius_vlm_source(tmp_path, "cpu") + cpu = _create_mobius_vlm_source(tmp_path, "cpu", ep="CPUExecutionProvider") + gpu = _create_mobius_vlm_source(tmp_path, "gpu", ep="CUDAExecutionProvider") out = tmp_path / "out" - cmd = _make_command(["generate-model-package", "-s", str(src), "-o", str(out)]) + cmd = _make_command(["generate-model-package", "-s", str(cpu), "-s", str(gpu), "-o", str(out)]) cmd.run() - models = out.with_suffix(".ortpackage") / "models" - decoder_model = json.loads((models / "decoder" / "cpu" / "genai_config_overlay.json").read_text())["model"] - embedding_model = json.loads((models / "embedding" / "cpu" / "genai_config_overlay.json").read_text())["model"] - vision_model = json.loads((models / "vision" / "cpu" / "genai_config_overlay.json").read_text())["model"] - # context_length and type are seeded on the Mobius fixture under - # ``model``; they belong only to the primary role's overlay. - assert "context_length" in decoder_model - assert "type" in decoder_model - for non_primary in (embedding_model, vision_model): - assert "context_length" not in non_primary - assert "type" not in non_primary - - def test_explicit_cpu_role_in_gpu_source_kept_as_cpu(self, tmp_path): - """A role with explicit CPU ``provider_options`` keeps CPU even when the source dir is named like a GPU build. - - Variant-name heuristics must not override a producer's explicit - per-role provider choice — Mobius outputs sometimes mark a - helper role (e.g. ``embedding``) as CPU even inside a - predominantly-GPU build. + component = out.with_suffix(".ortpackage") / "models" / "model" + source_model = json.loads((cpu / "genai_config.json").read_text())["model"] + for variant in ("cpu", "gpu"): + model = json.loads((component / variant / "genai_config.json").read_text())["model"] + for key in ("context_length", "type", "eos_token_id", "pad_token_id", "bos_token_id"): + if key in source_model: + assert model[key] == source_model[key], f"{variant}/{key} diverged from the source" + + def test_explicit_cpu_role_keeps_cpu_provider_options(self, tmp_path): + """A role with explicit CPU ``provider_options`` keeps them, and doesn't sway the variant EP. + + Mobius outputs sometimes mark a helper role (e.g. ``embedding``) as + CPU even inside a predominantly-GPU build. The variant EP is the + single non-CPU EP across roles (CUDA here) — CPU roles are exempt + because ORT-GenAI registers CPU implicitly and a CPU role never + claims the model's device. """ src = tmp_path / "gpu" src.mkdir() @@ -1647,30 +1749,25 @@ def test_explicit_cpu_role_in_gpu_source_kept_as_cpu(self, tmp_path): cmd = _make_command(["generate-model-package", "-s", str(src), "-o", str(out)]) cmd.run() - models = out.with_suffix(".ortpackage") / "models" - assert ( - json.loads((models / "decoder" / "metadata.json").read_text())["variants"]["gpu"]["ep"] - == "CUDAExecutionProvider" - ) - assert ( - json.loads((models / "vision" / "metadata.json").read_text())["variants"]["gpu"]["ep"] - == "CUDAExecutionProvider" - ) - # Critical: explicit CPU role must NOT be promoted to CUDA via - # the variant-name "gpu" heuristic. - assert ( - json.loads((models / "embedding" / "metadata.json").read_text())["variants"]["gpu"]["ep"] - == "CPUExecutionProvider" + component = out.with_suffix(".ortpackage") / "models" / "model" + assert json.loads((component / "component.json").read_text())["variants"]["gpu"]["ep"] == ( + "CUDAExecutionProvider" ) + model = json.loads((component / "gpu" / "genai_config.json").read_text())["model"] + # Critical: the CPU role's explicit (empty) provider_options survive + # so the loader doesn't put the embedding on CUDA. + assert model["embedding"]["session_options"]["provider_options"] == [] + assert model["decoder"]["session_options"]["provider_options"] == [{"cuda": {}}] + assert model["vision"]["session_options"]["provider_options"] == [{"cuda": {}}] def test_base_config_source_picks_richest_role_set(self, tmp_path): """When sources expose different role sets, the base config is taken from the source with the most roles. - Otherwise the package's base ``configs/genai_config.json`` could - miss role blocks that downstream components rely on. Example: - gpu source only has decoder; cpu source has all three. The base - must come from cpu so embedding/vision components have role - markers in the base config. + Every variant config is derived from that base, so a base missing + role blocks would produce variant configs that cannot describe the + whole model. Example: gpu source only has decoder; cpu source has + all three. The base must come from cpu so each variant config + still carries embedding/vision. """ # cpu source: full three-role VLM. cpu = _create_mobius_vlm_source(tmp_path, "cpu") @@ -1695,12 +1792,14 @@ def test_base_config_source_picks_richest_role_set(self, tmp_path): cmd = _make_command(["generate-model-package", "-s", str(gpu_dir), "-s", str(cpu), "-o", str(out)]) cmd.run() - # Base must carry all three role blocks (so the embedding/vision - # components are findable). If first-source-wins ran, only - # decoder would appear. - base = json.loads((out.with_suffix(".ortpackage") / "configs" / "genai_config.json").read_text()) + # Every variant config must carry all three role blocks (inherited + # from the richest base). If first-source-wins ran, only decoder + # would appear. + config = json.loads( + (out.with_suffix(".ortpackage") / "models" / "model" / "cpu" / "genai_config.json").read_text() + ) for role in ("decoder", "embedding", "vision"): - assert role in base["model"], f"base config missing {role} block; wrong source selected" + assert role in config["model"], f"variant config missing {role} block; wrong source selected" class TestUnsafeGenaiFilenamesRejected: @@ -1847,12 +1946,13 @@ def test_copies_when_destination_missing(self, tmp_path): class TestRoleToComponentConflictDetection: """Two variants mapping the same role to different components must raise. - Per the per-role-component layout, each genai_config role belongs to - exactly one package component. A direct caller that constructs - variants by hand could violate this invariant (e.g. by reusing the - same source_genai under two component names); ``write_model_package`` - detects the conflict at the role_to_component build step and raises - rather than silently keep one mapping and drop the other. + A direct ``write_model_package`` caller may split roles across + components (the plain ORT model-package spec allows it), but each + genai_config role must still belong to exactly one component. A caller + could violate that by hand (e.g. by reusing the same source_genai under + two component names); ``write_model_package`` detects the conflict at + the role_to_component build step and raises rather than silently keep + one mapping and drop the other. """ def test_same_role_mapped_to_two_components_raises(self, tmp_path): diff --git a/test/common/hf/io_config/test_task_config.py b/test/common/hf/io_config/test_task_config.py index 91f18f9364..fd22d5180e 100644 --- a/test/common/hf/io_config/test_task_config.py +++ b/test/common/hf/io_config/test_task_config.py @@ -6,7 +6,7 @@ import pytest -from olive.common.hf.io_config.io_resolver import get_task_template, resolve_alias +from olive.common.hf.io_config.io_resolver import _get_nested_attr, get_task_template, resolve_alias from olive.common.hf.io_config.task_config import ( _build_inputs, generate_dummy_inputs, @@ -35,6 +35,83 @@ class Config: assert resolve_alias(Config(), "num_layers") is None + def test_top_level_value_wins_over_text_config_alias(self): + """Regression: a composite config's own top-level value must not be overridden. + + Some VL-family configs (e.g. ovis2) carry both a top-level ``hidden_size`` and a + differing ``text_config.hidden_size``; the top-level value is authoritative and the + ``text_config.*`` alias is only a fallback for configs that lack it. + """ + + class TextConfig: + hidden_size = 4096 + num_hidden_layers = 32 + num_attention_heads = 32 + + class Config: + hidden_size = 1536 + num_hidden_layers = 24 + num_attention_heads = 12 + text_config = TextConfig() + + config = Config() + assert resolve_alias(config, "hidden_size") == 1536 + assert resolve_alias(config, "num_layers") == 24 + assert resolve_alias(config, "num_attention_heads") == 12 + + def test_text_config_alias_used_when_top_level_missing(self): + class TextConfig: + hidden_size = 4096 + num_hidden_layers = 32 + + class Config: + text_config = TextConfig() + + config = Config() + assert resolve_alias(config, "hidden_size") == 4096 + assert resolve_alias(config, "num_layers") == 32 + + def test_ambiguous_per_layer_attribute_resolves_to_none(self): + """Regression: transformers 5.x per-layer configs raise on ambiguous attribute access. + + Gemma4-family (heterogeneous) configs raise ``AmbiguousGlobalPerLayerAttributeError`` + -- not ``AttributeError`` -- when a per-layer attribute such as ``head_dim`` is read + off the top-level config. Alias resolution must degrade to ``None`` instead of + propagating the error. + """ + + class AmbiguousGlobalPerLayerAttributeError(Exception): + pass + + class Config: + hidden_size = 1024 + + def __getattr__(self, name): + if name == "head_dim": + raise AmbiguousGlobalPerLayerAttributeError(name) + raise AttributeError(name) + + config = Config() + assert _get_nested_attr(config, "head_dim") is None + assert resolve_alias(config, "head_dim") is None + # unrelated attributes still resolve normally + assert resolve_alias(config, "hidden_size") == 1024 + + def test_ambiguous_per_layer_attribute_real_transformers_exception(self): + """Same as above, but with the real transformers exception type when available.""" + heterogeneity = pytest.importorskip("transformers.integrations.heterogeneity.configuration_utils") + error_cls = getattr(heterogeneity, "AmbiguousGlobalPerLayerAttributeError", None) + if error_cls is None: + pytest.skip("transformers version has no AmbiguousGlobalPerLayerAttributeError") + + class Config: + def __getattr__(self, name): + if name == "head_dim": + raise error_cls(name) + raise AttributeError(name) + + assert resolve_alias(Config(), "head_dim") is None + class TestGetIOConfig: def test_text_classification(self): diff --git a/test/common/hf/test_quant.py b/test/common/hf/test_quant.py new file mode 100644 index 0000000000..e6108f60fb --- /dev/null +++ b/test/common/hf/test_quant.py @@ -0,0 +1,89 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for the ONNX-export wrapper modules in ``olive.common.hf.quant``.""" + +from __future__ import annotations + +import pytest +import torch + +from olive.common.hf.quant import QuantEmbeddingNbit, QuantLinearNbit +from olive.common.quant.tensor import QuantTensor + + +class TestOnnxBlockSizeValidation: + """``MatMulNBits`` / ``GatherBlockQuantized`` require block_size to be a power of 2, >= 16. + + Without an explicit check, a non-positive ``group_size`` silently falls back to the full + ``embedding_dim`` / ``in_features``, which for many real models (Qwen2.5: 1536, + Phi-3-mini: 3072) is not a power of 2 — surfacing only as an opaque native ORT + session-initialization crash. + """ + + @pytest.mark.parametrize("embedding_dim", [1536, 3072, 100]) + def test_embedding_non_power_of_two_fallback_raises(self, embedding_dim: int): + with pytest.raises(ValueError, match="GatherBlockQuantized requires block_size"): + QuantEmbeddingNbit(num_embeddings=8, embedding_dim=embedding_dim, group_size=-1) + + def test_embedding_fallback_below_minimum_raises(self): + with pytest.raises(ValueError, match=r">= 16, got 8"): + QuantEmbeddingNbit(num_embeddings=8, embedding_dim=8, group_size=0) + + def test_embedding_explicit_non_power_of_two_group_size_raises(self): + with pytest.raises(ValueError, match="GatherBlockQuantized requires block_size"): + QuantEmbeddingNbit(num_embeddings=8, embedding_dim=1536, group_size=96) + + @pytest.mark.parametrize("group_size", [-1, 32, 64]) + def test_embedding_valid_block_sizes_are_accepted(self, group_size: int): + module = QuantEmbeddingNbit(num_embeddings=8, embedding_dim=64, group_size=group_size) + assert module.group_size == (64 if group_size <= 0 else group_size) + + @pytest.mark.parametrize("in_features", [1536, 3072, 100]) + def test_linear_non_power_of_two_fallback_raises(self, in_features: int): + with pytest.raises(ValueError, match="MatMulNBits requires block_size"): + QuantLinearNbit(group_size=-1, in_features=in_features, out_features=8) + + @pytest.mark.parametrize("group_size", [-1, 32, 128]) + def test_linear_valid_block_sizes_are_accepted(self, group_size: int): + module = QuantLinearNbit(group_size=group_size, in_features=128, out_features=8) + assert module.group_size == (128 if group_size <= 0 else group_size) + + +class TestQuantEmbeddingNbitFromQuantTensor: + """``from_quant_tensor`` must reshape (and therefore shape-validate) like the Linear one.""" + + def test_per_tensor_quant_tensor_fails_immediately(self): + """``group_size == 0`` (per-tensor) scales are ``(1, 1)``, not ``(num_embeddings, 1)``. + + The bare ``detach().clone()`` used to install that ``(1, 1)`` tensor over a buffer + declared ``(num_embeddings, n_groups)`` without complaint, deferring the failure to a + confusing ``reshape`` error somewhere in the exported graph. Mirroring + ``QuantLinearNbit.from_quant_tensor``'s ``.reshape(...)`` makes it fail here instead. + """ + qt = QuantTensor.from_float(torch.randn(8, 32), bits=4, symmetric=True, group_size=0) + assert tuple(qt.scales.shape) == (1, 1) + with pytest.raises(RuntimeError, match="invalid for input of size"): + QuantEmbeddingNbit.from_quant_tensor(qt) + + def test_per_tensor_quant_tensor_fails_the_same_way_for_linear(self): + """Same semantics as the Embedding path — per-tensor QuantTensors are not exportable.""" + qt = QuantTensor.from_float(torch.randn(8, 32), bits=4, symmetric=True, group_size=0) + with pytest.raises(RuntimeError, match="invalid for input of size"): + QuantLinearNbit.from_quant_tensor(qt) + + @pytest.mark.parametrize("group_size", [-1, 16]) + @pytest.mark.parametrize("symmetric", [True, False]) + def test_supported_group_sizes_round_trip(self, group_size: int, symmetric: bool): + weight = torch.randn(8, 32) + qt = QuantTensor.from_float(weight, bits=4, symmetric=symmetric, group_size=group_size) + module = QuantEmbeddingNbit.from_quant_tensor(qt) + + assert module.group_size == (32 if group_size <= 0 else group_size) + assert torch.equal(module.qweight, qt.qweight) + assert torch.equal(module.scales, qt.scales) + if qt.qzeros is None: + assert module.qzeros is None + else: + assert torch.equal(module.qzeros, qt.qzeros) diff --git a/test/common/quant/test_forward_parity.py b/test/common/quant/test_forward_parity.py new file mode 100644 index 0000000000..3133ce1628 --- /dev/null +++ b/test/common/quant/test_forward_parity.py @@ -0,0 +1,265 @@ +# pylint: disable=protected-access,not-callable +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Numerical-parity tests for QuantTensor-backed modules. + +Confirms that ``nn.Linear`` / ``nn.Embedding`` / fused-3D MoE host modules +whose weight has been swapped to a ``QuantTensor`` produce **bit-exact** +outputs vs the same forward run against ``QuantTensor.to_dense()`` (the +canonical eager dequant), and that the ONNX-export wrappers built by +``make_export_compatible_quant`` round-trip through onnxruntime with +matching outputs. +""" + +from __future__ import annotations + +import copy + +import onnx +import onnxruntime as ort +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from olive.common.hf.quant import make_export_compatible_quant +from olive.common.quant.state_dict import install_quant_tensor_param +from olive.common.quant.tensor import QuantTensor + + +def _quantize_inplace(module: nn.Module, pname: str, *, bits: int, group_size: int, symmetric: bool) -> None: + param = module._parameters[pname] + qt = QuantTensor.from_float( + param.data.detach().clone(), + bits=bits, + group_size=group_size, + symmetric=symmetric, + ) + install_quant_tensor_param(module, pname, qt) + + +def _dense_reference(model: nn.Module) -> nn.Module: + """Return a deep copy of ``model`` with every QuantTensor weight materialized.""" + ref = copy.deepcopy(model) + for module in ref.modules(): + weight = module._parameters.get("weight") + if weight is None: + continue + if isinstance(weight.data, QuantTensor): + module._parameters["weight"] = nn.Parameter(weight.data.to_dense(), requires_grad=False) + return ref + + +@pytest.mark.parametrize(("bits", "group_size", "symmetric"), [(4, 32, False), (4, -1, True), (8, 32, False)]) +def test_full_model_forward_parity(bits, group_size, symmetric): + torch.manual_seed(0) + + class Toy(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embed = nn.Embedding(32, 64) + self.fc1 = nn.Linear(64, 128, bias=False) + self.fc2 = nn.Linear(128, 64, bias=True) + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + return self.fc2(F.silu(self.fc1(self.embed(ids)))) + + model = Toy().eval() + ids = torch.randint(0, 32, (2, 16)) + ref_out = model(ids) + + for sub_module in (model.embed, model.fc1, model.fc2): + _quantize_inplace(sub_module, "weight", bits=bits, group_size=group_size, symmetric=symmetric) + + quant_out = model(ids) + dense_out = _dense_reference(model)(ids) + + # vs canonical dequant: must be bit-exact (same kernels, same data) + torch.testing.assert_close(quant_out, dense_out, rtol=0, atol=0) + # vs original fp: just a sanity bound on quant error + assert (quant_out - ref_out).abs().mean().item() < 1.0 + + +@pytest.mark.parametrize(("bits", "group_size", "symmetric"), [(4, 16, False), (8, -1, True)]) +def test_fused_moe_forward_parity(bits, group_size, symmetric): + """Fused 3D-expert forward: index into a QuantTensor 3D weight per expert.""" + torch.manual_seed(0) + num_experts, in_dim, hidden = 4, 32, 64 + + class FusedMoE(nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate_up = nn.Parameter(torch.randn(num_experts, hidden, in_dim)) + self.down = nn.Parameter(torch.randn(num_experts, in_dim, hidden)) + + def forward(self, x: torch.Tensor, expert_ids: torch.Tensor) -> torch.Tensor: + # x: (tokens, in_dim); expert_ids: (tokens,) routing + outs = [] + for t in range(x.shape[0]): + e = int(expert_ids[t].item()) + h = F.linear(x[t : t + 1], self.gate_up[e]) # (1, hidden) + outs.append(F.linear(F.silu(h), self.down[e])) # (1, in_dim) + return torch.cat(outs, dim=0) + + model = FusedMoE().eval() + x = torch.randn(8, in_dim) + expert_ids = torch.tensor([0, 1, 2, 3, 0, 2, 1, 3]) + ref_out = model(x, expert_ids) + + # Quantize the 3D expert tensors directly via QuantTensor.from_float. + for pname in ("gate_up", "down"): + _quantize_inplace(model, pname, bits=bits, group_size=group_size, symmetric=symmetric) + + quant_out = model(x, expert_ids) + + # Reference: same forward against dense-materialized 3D weights. + dense_model = FusedMoE() + with torch.no_grad(): + dense_model.gate_up = nn.Parameter(model.gate_up.data.to_dense(), requires_grad=False) + dense_model.down = nn.Parameter(model.down.data.to_dense(), requires_grad=False) + dense_out = dense_model(x, expert_ids) + + torch.testing.assert_close(quant_out, dense_out, rtol=0, atol=0) + assert (quant_out - ref_out).abs().mean().item() < 5.0 + + +@pytest.mark.parametrize(("bits", "group_size", "symmetric"), [(4, 32, False), (4, 32, True), (8, 32, False)]) +def test_onnx_export_parity(tmp_path, bits, group_size, symmetric): + """Olive-quantized nn.Linear -> ONNX MatMulNBits -> onnxruntime matches eager.""" + torch.manual_seed(0) + in_dim, out_dim = 64, 32 + + class LinearOnly(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = nn.Linear(in_dim, out_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc(x) + + model = LinearOnly().eval() + _quantize_inplace(model.fc, "weight", bits=bits, group_size=group_size, symmetric=symmetric) + + x = torch.randn(1, 4, in_dim) + eager_out = model(x).detach() + + # Build an export-compatible variant (QuantLinearNbit wrapper) and export to ONNX. + export_model = make_export_compatible_quant(copy.deepcopy(model), dynamo=False) + + onnx_path = tmp_path / "model.onnx" + torch.onnx.export( + export_model, + (x,), + onnx_path, + input_names=["input"], + output_names=["output"], + dynamic_axes={"input": {0: "batch", 1: "seq"}, "output": {0: "batch", 1: "seq"}}, + opset_version=21, + dynamo=False, + ) + + onnx.checker.check_model(str(onnx_path)) + sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_out = sess.run(["output"], {"input": x.numpy()})[0] + + torch.testing.assert_close(torch.from_numpy(ort_out), eager_out, rtol=1e-2, atol=1e-2) + + +def _quantize_model_targets(model, *, moe, embeds, lm_head, bits=4, group_size=32, symmetric=False): + """Quantize every ``iter_quant_targets`` selection in place with a QuantTensor.""" + from olive.common.quant.selection import iter_quant_targets + + for module, pname, _ in list( + iter_quant_targets( + model, + quantize_lm_head=lm_head, + quantize_embeds=embeds, + quantize_moe=moe, + ) + ): + _quantize_inplace(module, pname, bits=bits, group_size=group_size, symmetric=symmetric) + + +def test_real_hf_mixtral_moe_round_trip(): + """Round-trip a real (config-only, no weight download) HF MoE architecture. + + Builds a tiny ``MixtralForCausalLM`` from config, quantizes it with ``moe=True`` (the + fused 3D expert weights plus embeddings / lm_head), then saves and reloads the state + dict and verifies every quantized weight dequantizes bit-identically after reload. + + Note: the model's own ``grouped_mm`` expert forward is not exercised here — Olive's + storage-only MoE quantization dispatches ``F.linear`` / ``F.embedding`` and does not + implement fused ``grouped_mm``; ONNX export / execution of the experts is delegated to + ORT GenAI ModelBuilder / Mobius. This test therefore validates the buffer-backed + save/reload path (the actual round-trip contract) on a real architecture. + """ + transformers = pytest.importorskip("transformers") + + from olive.common.quant.state_dict import refresh_quant_tensor_refs + + def build(): + torch.manual_seed(0) + cfg = transformers.MixtralConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + tie_word_embeddings=False, + ) + return transformers.MixtralForCausalLM(cfg).eval() + + model = build() + _quantize_model_targets(model, moe=True, embeds=True, lm_head=True) + + # every fused expert weight must now be a QuantTensor; biases (if any) stay dense. + experts = model.model.layers[0].mlp.experts + assert isinstance(experts._parameters["gate_up_proj"].data, QuantTensor) + assert isinstance(experts._parameters["down_proj"].data, QuantTensor) + + # Snapshot the dequantized value of every quantized parameter. + def dequant_snapshot(m): + snap = {} + for name, sub in m.named_modules(): + for pname, param in sub._parameters.items(): + if param is not None and isinstance(param.data, QuantTensor): + snap[f"{name}.{pname}"] = param.data.to_dense().clone() + return snap + + before = dequant_snapshot(model) + assert any("experts" in k for k in before), "expected at least one quantized expert weight" + + # Save state dict (QuantTensor params dropped; only plain buffers persisted) then reload + # onto a freshly quantized model and re-verify each dequantized weight is bit-identical. + state = model.state_dict() + assert not any(isinstance(v, QuantTensor) for v in state.values()) + + reloaded = build() + _quantize_model_targets(reloaded, moe=True, embeds=True, lm_head=True) + _missing, unexpected = reloaded.load_state_dict(state, strict=False) + assert not unexpected, f"unexpected keys on reload: {unexpected}" + refresh_quant_tensor_refs(reloaded) + + after = dequant_snapshot(reloaded) + assert after.keys() == before.keys() + for key, ref in before.items(): + torch.testing.assert_close(after[key], ref, rtol=0, atol=0) + + +def test_2bit_quant_linear_export_rejected(): + """Minor: exporting a 2-bit QuantTensor via QuantLinearNbit must raise a clear error. + + ``QuantLinearNbit`` (the ONNX-export wrapper) only supports 4-bit and 8-bit packing; + 2-bit is a PyTorch-checkpoint-only feature and must fail fast (not silently misbehave) + when export is attempted. + """ + from olive.common.hf.quant import QuantLinearNbit + + qt = QuantTensor.from_float(torch.randn(16, 32), bits=2, group_size=16, symmetric=True) + with pytest.raises(ValueError, match="2-bit"): + QuantLinearNbit.from_quant_tensor(qt) diff --git a/test/common/quant/test_hf_utils.py b/test/common/quant/test_hf_utils.py index 92302fa5d0..b5799fb4c5 100644 --- a/test/common/quant/test_hf_utils.py +++ b/test/common/quant/test_hf_utils.py @@ -14,14 +14,21 @@ OliveHfQuantizationOverrideConfig, OliveHfQuantizer, replace_matching_submodules, - tie_quant_modules, tie_quant_word_embeddings, ) -from olive.common.quant.nn import QuantEmbedding, QuantLinear +from olive.common.quant.tensor import QuantTensor # pylint: disable=W0212 +def _is_olive_quant(module: nn.Module) -> bool: + """Return True if ``module`` is an nn.Linear/nn.Embedding whose weight is a QuantTensor.""" + if not isinstance(module, (nn.Linear, nn.Embedding)): + return False + weight = module._parameters.get("weight") + return weight is not None and isinstance(weight.data, QuantTensor) + + class TestOliveHfQuantizationMethod: def test_enum_value(self): """Test that the enum has the expected value.""" @@ -209,13 +216,13 @@ def test_process_model_before_weight_loading(self): # Process model quantizer._process_model_before_weight_loading(model) - # Check that linear layers are replaced with QuantLinear - assert isinstance(model.linear1, QuantLinear) - assert isinstance(model.linear2, QuantLinear) + # Check that linear layers are replaced with QuantTensor weights + assert _is_olive_quant(model.linear1) + assert _is_olive_quant(model.linear2) # Check that embeddings and lm_head are NOT quantized by default - assert not isinstance(model.embed, QuantEmbedding) - assert not isinstance(model.lm_head, QuantLinear) + assert not _is_olive_quant(model.embed) + assert not _is_olive_quant(model.lm_head) def test_process_model_with_embeds_quantization(self): """Test that embeddings are quantized when embeds=True.""" @@ -229,7 +236,7 @@ def test_process_model_with_embeds_quantization(self): quantizer._process_model_before_weight_loading(model) # Check that embeddings are quantized - assert isinstance(model.embed, QuantEmbedding) + assert _is_olive_quant(model.embed) def test_process_model_with_lm_head_quantization(self): """Test that lm_head is quantized when lm_head=True.""" @@ -243,7 +250,7 @@ def test_process_model_with_lm_head_quantization(self): quantizer._process_model_before_weight_loading(model) # Check that lm_head is quantized - assert isinstance(model.lm_head, QuantLinear) + assert _is_olive_quant(model.lm_head) def test_process_model_with_modules_to_not_convert(self): """Test that specified modules are not converted.""" @@ -263,8 +270,9 @@ def test_process_model_with_modules_to_not_convert(self): # Check that linear1 is NOT quantized assert isinstance(model.linear1, nn.Linear) + assert not _is_olive_quant(model.linear1) # But linear2 should be quantized - assert isinstance(model.linear2, QuantLinear) + assert _is_olive_quant(model.linear2) def test_is_serializable(self): """Test that quantizer is serializable.""" @@ -278,6 +286,92 @@ def test_is_trainable(self): quantizer = OliveHfQuantizer(config) assert quantizer.is_trainable is False + def test_checkpoint_keys_default_none(self): + """Test that ``_checkpoint_keys`` defaults to None when no checkpoint_files given.""" + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16) + quantizer = OliveHfQuantizer(config) + assert quantizer._checkpoint_keys is None + + model_config = SimpleConfig() + model = SimpleModel(model_config) + quantizer._process_model_before_weight_loading(model) + assert quantizer._checkpoint_keys is None + + def test_checkpoint_keys_populated_from_safetensors_files(self, tmp_path): + """``_process_model_before_weight_loading`` reads real key names from safetensors shards.""" + from safetensors.torch import save_file + + shard = tmp_path / "model.safetensors" + save_file({"linear1.weight_qweight": torch.zeros(4, 4, dtype=torch.uint8)}, str(shard)) + + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16) + quantizer = OliveHfQuantizer(config) + + model_config = SimpleConfig() + model = SimpleModel(model_config) + quantizer._process_model_before_weight_loading(model, checkpoint_files=[str(shard)]) + + assert quantizer._checkpoint_keys == {"linear1.weight_qweight"} + + def test_checkpoint_keys_none_for_non_safetensors_files(self): + """Legacy ``.bin`` checkpoint shards fall back to ``None`` (no cheap header-only read).""" + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16) + quantizer = OliveHfQuantizer(config) + + model_config = SimpleConfig() + model = SimpleModel(model_config) + quantizer._process_model_before_weight_loading(model, checkpoint_files=["pytorch_model.bin"]) + + assert quantizer._checkpoint_keys is None + + def test_process_model_after_weight_loading_uses_checkpoint_keys(self): + """Test that ``_process_model_after_weight_loading`` uses ``_checkpoint_keys``. + + It should clear ``is_placeholder`` by exact key membership instead of the + buffer-identity heuristic. + """ + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16) + quantizer = OliveHfQuantizer(config) + + model_config = SimpleConfig() + model = SimpleModel(model_config) + quantizer._process_model_before_weight_loading(model) + assert model.linear1.weight.is_placeholder is True + assert model.linear2.weight.is_placeholder is True + + # Simulate an in-place ``.copy_()`` loader: identity of the buffer object never + # changes, so the fallback heuristic alone would miss this load. + for linear in (model.linear1, model.linear2): + linear.weight_qweight.copy_(torch.randint(0, 255, linear.weight_qweight.shape, dtype=torch.uint8)) + + quantizer._checkpoint_keys = { + "linear1.weight_qweight", + "linear1.weight_scales", + "linear2.weight_qweight", + "linear2.weight_scales", + } + quantizer._process_model_after_weight_loading(model) + + assert model.linear1.weight.is_placeholder is False + assert model.linear2.weight.is_placeholder is False + + def test_process_model_after_weight_loading_raises_on_missing_checkpoint_key(self): + """Fail closed: a known manifest that omits a quantized parameter must raise. + + Previously the model was returned with zero-filled placeholder weights and no error. + """ + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16) + quantizer = OliveHfQuantizer(config) + + model_config = SimpleConfig() + model = SimpleModel(model_config) + quantizer._process_model_before_weight_loading(model) + + # linear2's keys are absent from the checkpoint manifest. + quantizer._checkpoint_keys = {"linear1.weight_qweight", "linear1.weight_scales"} + with pytest.raises(RuntimeError, match=r"missing weights for: linear2\.weight"): + quantizer._process_model_after_weight_loading(model) + class TestReplaceMatchingSubmodules: def test_replace_all_linear_layers(self): @@ -357,49 +451,6 @@ def transform(m, name): assert isinstance(result.sub[2], nn.Identity) -class TestTieQuantModules: - def test_tie_quant_linear_modules(self): - """Test tying two QuantLinear modules.""" - # Create two QuantLinear modules - qlinear1 = QuantLinear(32, 10, bits=4, symmetric=True, group_size=16) - qlinear2 = QuantLinear(32, 10, bits=4, symmetric=True, group_size=16) - - # Set some values in qlinear1 - qlinear1.qweight.fill_(1) - qlinear1.scales.fill_(0.5) - - # Initially different - assert not torch.all(qlinear1.qweight == qlinear2.qweight) - - # Tie them - tie_quant_modules(qlinear1, qlinear2) - - # Now they should share buffers - assert qlinear1.qweight is qlinear2.qweight - assert qlinear1.scales is qlinear2.scales - - # Modifications to one should affect the other - qlinear1.qweight.fill_(2) - assert torch.all(qlinear2.qweight == 2) - - def test_tie_quant_embedding_modules(self): - """Test tying two QuantEmbedding modules.""" - # Create two QuantEmbedding modules - qembed1 = QuantEmbedding(100, 64, bits=4, symmetric=True, group_size=16) - qembed2 = QuantEmbedding(100, 64, bits=4, symmetric=True, group_size=16) - - # Set some values in qembed1 - qembed1.qweight.fill_(1) - qembed1.scales.fill_(0.5) - - # Tie them - tie_quant_modules(qembed1, qembed2) - - # Now they should share buffers - assert qembed1.qweight is qembed2.qweight - assert qembed1.scales is qembed2.scales - - class TestTieQuantWordEmbeddings: def test_tie_word_embeddings(self): """Test tying word embeddings in a model.""" @@ -419,13 +470,219 @@ def test_tie_word_embeddings(self): # Process model to quantize embeddings quantizer._process_model_before_weight_loading(model) - # Set different values - model.embed.qweight.fill_(1) - model.lm_head.qweight.fill_(2) + # Set different values on the aliased buffers + model.embed.weight_qweight.fill_(1) + model.lm_head.weight_qweight.fill_(2) # Tie embeddings tie_quant_word_embeddings(model) - # Now they should share buffers - assert model.embed.qweight is model.lm_head.qweight - assert model.embed.scales is model.lm_head.scales + # Now they should share the underlying buffers and the Parameter + assert model.embed.weight_qweight is model.lm_head.weight_qweight + assert model.embed.weight_scales is model.lm_head.weight_scales + assert model.embed._parameters["weight"] is model.lm_head._parameters["weight"] + + +# -- MoE quantization fixtures and tests -------------------------------------- + + +class MoEConfig(PretrainedConfig): + model_type = "moe_simple" + + def __init__(self, hidden_size=64, vocab_size=64, num_experts=4, num_layers=1, **kwargs): + super().__init__(**kwargs) + self.hidden_size = hidden_size + self.vocab_size = vocab_size + self.num_experts = num_experts + self.num_hidden_layers = num_layers + # required by ModelWrapper / LayerWrapper + self.num_attention_heads = 4 + self.num_key_value_heads = 4 + self.head_dim = hidden_size // 4 + self.intermediate_size = hidden_size + + +class _MoEExpert(nn.Module): + """Per-expert sub-module (ModuleList style — like Mixtral / PhiMoE).""" + + def __init__(self, hidden_size: int): + super().__init__() + self.w1 = nn.Linear(hidden_size, hidden_size, bias=False) + self.w2 = nn.Linear(hidden_size, hidden_size, bias=False) + + +class _MoELayer(nn.Module): + def __init__(self, hidden_size: int, num_experts: int): + super().__init__() + self.mlp = nn.Module() + self.mlp.gate = nn.Linear(hidden_size, num_experts, bias=False) + self.mlp.experts = nn.ModuleList([_MoEExpert(hidden_size) for _ in range(num_experts)]) + + +class MoESimpleModel(PreTrainedModel): + config_class = MoEConfig + + def __init__(self, config): + super().__init__(config) + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.model.layers = nn.ModuleList( + [_MoELayer(config.hidden_size, config.num_experts) for _ in range(config.num_hidden_layers)] + ) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def get_input_embeddings(self): + return self.model.embed_tokens + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + +class TestOliveHfQuantizerMoE: + """MoE-specific behavior of ``OliveHfQuantizer``. + + The previous implementation silently quantized every per-expert + ``nn.Linear`` in ``ModuleList(Expert)`` blocks (Mixtral / PhiMoE / + Qwen2/3-MoE). With the new ``moe`` category flag (default ``False``), + every ``nn.Module`` under each experts subtree is added to the skip + set — fixing the silent quantization. + """ + + def test_module_list_experts_skipped_by_default(self): + """Regression: ModuleList(Expert) linears must stay as nn.Linear when moe=False.""" + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16, moe=False) + quantizer = OliveHfQuantizer(config) + model = MoESimpleModel(MoEConfig()) + + quantizer._process_model_before_weight_loading(model) + + for expert in model.model.layers[0].mlp.experts: + assert isinstance(expert.w1, nn.Linear) + assert not _is_olive_quant(expert.w1) + assert isinstance(expert.w2, nn.Linear) + assert not _is_olive_quant(expert.w2) + # The router (gate) is under the mlp but not under .experts. It is always kept in + # full precision: quantizing it perturbs the routing decisions themselves, and it is + # excluded by resolved-module identity (``LayerWrapper.get_router()``) so that bare + # ``nn.Linear`` routers -- e.g. Jamba's -- cannot slip into the ordinary 2D walk. + assert not _is_olive_quant(model.model.layers[0].mlp.gate) + assert isinstance(model.model.layers[0].mlp.gate, nn.Linear) + + def test_module_list_experts_quantized_when_moe_true(self): + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16, moe=True) + quantizer = OliveHfQuantizer(config) + model = MoESimpleModel(MoEConfig()) + + quantizer._process_model_before_weight_loading(model) + + for expert in model.model.layers[0].mlp.experts: + assert _is_olive_quant(expert.w1) + assert _is_olive_quant(expert.w2) + + def test_moe_default_is_false(self): + """``moe`` defaults to False — opt-in, matching lm_head / embeds.""" + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16) + assert config.moe is False + + def test_regex_skip_pattern(self): + """``re:`` prefix opts into regex fullmatch for modules_to_not_convert.""" + config = OliveHfQuantizationConfig( + bits=4, + symmetric=True, + group_size=16, + moe=True, + # Quantize experts but keep the router in fp. + modules_to_not_convert=["re:.*\\.mlp\\.gate"], + ) + quantizer = OliveHfQuantizer(config) + model = MoESimpleModel(MoEConfig()) + + quantizer._process_model_before_weight_loading(model) + + # gate is skipped via regex + assert isinstance(model.model.layers[0].mlp.gate, nn.Linear) + assert not _is_olive_quant(model.model.layers[0].mlp.gate) + # experts are quantized + assert _is_olive_quant(model.model.layers[0].mlp.experts[0].w1) + + +class TestRegexOverrides: + def test_get_qlinear_init_args_with_regex_override(self): + overrides = { + "re:.*\\.mlp\\.experts\\..*\\.w1": {"bits": 8, "group_size": 32}, + "model.lm_head": {"bits": 4}, + } + config = OliveHfQuantizationConfig( + bits=4, + symmetric=True, + group_size=128, + overrides=overrides, + ) + init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1") + assert init_args == {"bits": 8, "symmetric": True, "group_size": 32} + + def test_first_matching_override_wins_in_insertion_order(self): + # Two overlapping overrides match the same target. Per the finalized precedence rule + # (design item 6), the FIRST key in insertion order wins — not the longer/literal one. + overrides = { + "re:.*\\.w1": {"bits": 8}, + "model.layers.0.mlp.experts.0.w1": {"bits": 6}, + } + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=128, overrides=overrides) + init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1") + # The regex is first in insertion order, so it wins even though the literal is "more specific". + assert init_args["bits"] == 8 + + def test_reordering_overlapping_overrides_flips_the_winner(self): + # Same two overlapping overrides, literal placed first -> literal wins. + overrides = { + "model.layers.0.mlp.experts.0.w1": {"bits": 6}, + "re:.*\\.w1": {"bits": 8}, + } + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=128, overrides=overrides) + init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1") + assert init_args["bits"] == 6 + + def test_override_precedence_survives_to_dict_roundtrip(self): + # Regression test for C1: to_dict() must NOT reorder overrides -- match_override + # resolves overlapping keys by first-match-wins in *insertion* order, so re-sorting + # on serialization would silently flip the winner on save -> reload. + overrides = { + "re:.*\\.w1": {"bits": 8}, + "model.layers.0.mlp.experts.0.w1": {"bits": 6}, + } + config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=128, overrides=overrides) + target = "model.layers.0.mlp.experts.0.w1" + + # Before round-trip: the regex (first inserted) wins. + assert config.get_qlinear_init_args(target)["bits"] == 8 + + reloaded = OliveHfQuantizationConfig(**config.to_dict()) + + # After round-trip: the winner must be unchanged. + assert list(reloaded.overrides.keys()) == list(config.overrides.keys()) + assert reloaded.get_qlinear_init_args(target)["bits"] == 8 + + def test_eager_regex_validation_at_construction(self): + # C2: unsafe `re:` patterns must be rejected eagerly at config construction, not + # only lazily on first match (which could let a bad pattern that never matches slip + # through undetected for an entire run). + with pytest.raises(ValueError, match="nested"): + OliveHfQuantizationConfig( + bits=4, + symmetric=True, + group_size=128, + overrides={"re:(a{1,2})+$": {"bits": 8}}, + ) + + def test_eager_regex_validation_for_modules_to_not_convert(self): + with pytest.raises(ValueError, match="nested"): + OliveHfQuantizationConfig( + bits=4, + symmetric=True, + group_size=128, + modules_to_not_convert=["re:(a+)+$"], + ) diff --git a/test/common/quant/test_nn.py b/test/common/quant/test_nn.py deleted file mode 100644 index c496dd21da..0000000000 --- a/test/common/quant/test_nn.py +++ /dev/null @@ -1,571 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -import math - -import pytest -import torch -import torch.nn as nn - -from olive.common.quant.nn import QuantEmbedding, QuantLinear -from olive.common.quant.utils import WeightQuantizer - - -class TestQuantModule: - @pytest.mark.parametrize("bits", [2, 4, 8]) - @pytest.mark.parametrize("symmetric", [True, False]) - @pytest.mark.parametrize("group_size", [-1, 16, 32]) - def test_initialization(self, bits, symmetric, group_size): - """Test QuantModule initialization with various parameters.""" - rows, cols = 64, 128 - - # Since QuantModule is abstract, we'll use QuantLinear - qmodule = QuantLinear( - in_features=cols, - out_features=rows, - bits=bits, - symmetric=symmetric, - group_size=group_size, - ) - - assert qmodule.rows == rows - assert qmodule.cols == cols - assert qmodule.quantizer.bits == bits - assert qmodule.quantizer.symmetric == symmetric - assert qmodule.quantizer.group_size == group_size - - def test_invalid_bits(self): - """Test that invalid bits raise ValueError.""" - with pytest.raises(ValueError, match="Only 2-bit, 4-bit and 8-bit quantization supported"): - QuantLinear(10, 20, bits=16, symmetric=True, group_size=-1) - - def test_invalid_group_size(self): - """Test that invalid group size raises ValueError.""" - with pytest.raises(ValueError, match="group_size must be >= 16 and power of 2"): - QuantLinear(10, 20, bits=4, symmetric=True, group_size=15) - - with pytest.raises(ValueError, match="group_size must be >= 16 and power of 2"): - QuantLinear(10, 20, bits=4, symmetric=True, group_size=24) - - def test_invalid_in_features_for_group_size(self): - """Test that in_features must be divisible by group_size.""" - # in_features=100 is not divisible by group_size=32 - with pytest.raises(ValueError, match=r"cols .* must be divisible by group_size"): - QuantLinear(in_features=100, out_features=20, bits=4, symmetric=True, group_size=32) - - # in_features=50 is not divisible by group_size=16 - with pytest.raises(ValueError, match=r"cols .* must be divisible by group_size"): - QuantLinear(in_features=50, out_features=20, bits=4, symmetric=True, group_size=16) - - # This should work: in_features=64 is divisible by group_size=32 - qlinear = QuantLinear(in_features=64, out_features=20, bits=4, symmetric=True, group_size=32) - assert qlinear.cols == 64 - - def test_buffer_shapes(self): - """Test that buffers have correct shapes.""" - rows, cols = 64, 128 - bits = 4 - packing_factor = 8 // bits - - qmodule = QuantLinear( - in_features=cols, - out_features=rows, - bits=bits, - symmetric=True, - group_size=32, - ) - - # Check qweight shape - assert qmodule.qweight.shape == (rows, math.ceil(cols / packing_factor)) - - # Check scales shape - quantizer = WeightQuantizer(bits=bits, symmetric=True, group_size=32, signed=False) - expected_scale_shape = quantizer.get_qparam_shape((rows, cols)) - assert qmodule.scales.shape == expected_scale_shape - - def test_symmetric_no_qzeros(self): - """Test that symmetric quantization has no qzeros.""" - qmodule = QuantLinear(16, 20, bits=4, symmetric=True, group_size=16) - assert qmodule.qzeros is None - - def test_asymmetric_has_qzeros(self): - """Test that asymmetric quantization has qzeros.""" - qmodule = QuantLinear(16, 20, bits=4, symmetric=False, group_size=16) - assert qmodule.qzeros is not None - - packing_factor = 8 // 4 - quantizer = WeightQuantizer(bits=4, symmetric=False, group_size=16, signed=False) - scale_shape = quantizer.get_qparam_shape((20, 16)) - expected_qzeros_shape = (scale_shape[0], math.ceil(scale_shape[1] / packing_factor)) - assert qmodule.qzeros.shape == expected_qzeros_shape - - -class TestQuantLinear: - def test_basic_initialization(self): - """Test basic QuantLinear initialization.""" - qlinear = QuantLinear(in_features=128, out_features=256, bits=4, symmetric=True, group_size=32) - assert qlinear.cols == 128 - assert qlinear.rows == 256 - assert qlinear.bias is not None - - def test_initialization_without_bias(self): - """Test QuantLinear initialization without bias.""" - qlinear = QuantLinear( - in_features=128, - out_features=256, - bits=4, - symmetric=True, - group_size=32, - bias=False, - ) - assert qlinear.bias is None - - def test_from_module_basic(self): - """Test creating QuantLinear from nn.Linear.""" - linear = nn.Linear(128, 256) - linear.weight.data.normal_(0, 0.02) - - qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32) - - assert qlinear.cols == 128 - assert qlinear.rows == 256 - assert qlinear.bias is not None - - # Check that weights are quantized - assert qlinear.qweight.dtype == torch.uint8 - assert qlinear.scales.dtype == linear.weight.dtype - - def test_from_module_without_bias(self): - """Test creating QuantLinear from nn.Linear without bias.""" - linear = nn.Linear(128, 256, bias=False) - linear.weight.data.normal_(0, 0.02) - - qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32) - - assert qlinear.bias is None - - def test_from_module_preserves_bias(self): - """Test that bias is preserved when converting.""" - linear = nn.Linear(128, 256) - linear.weight.data.normal_(0, 0.02) - linear.bias.data.fill_(0.5) - - qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32) - - assert torch.all(qlinear.bias == 0.5) - - @pytest.mark.parametrize("bits", [2, 4, 8]) - @pytest.mark.parametrize("symmetric", [True, False]) - def test_from_module_quantization_accuracy(self, bits, symmetric): - """Test that quantization/dequantization is reasonably accurate.""" - linear = nn.Linear(64, 32, bias=False) - linear.weight.data.normal_(0, 0.02) - - qlinear = QuantLinear.from_module(linear, bits=bits, symmetric=symmetric, group_size=16) - - # Dequantize and compare - dequantized = qlinear.unpack_and_dequantize(qlinear.qweight, qlinear.scales, qlinear.qzeros) - - # Check shapes match - assert dequantized.shape == linear.weight.shape - - # Check that dequantized weights are close to original (within quantization error) - # The tolerance depends on the bit width and data distribution - max_diff = torch.max(torch.abs(dequantized - linear.weight)) - # For 4-bit, we expect larger errors than 8-bit - tolerance = 0.1 if bits == 4 else 0.05 - assert max_diff < tolerance, f"Max difference {max_diff} exceeds tolerance {tolerance}" - - def test_forward_shape(self): - """Test that forward pass produces correct output shape.""" - linear = nn.Linear(128, 256) - linear.weight.data.normal_(0, 0.02) - - qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32) - - x = torch.randn(16, 128) - output = qlinear(x) - - assert output.shape == (16, 256) - - def test_forward_batch_shape(self): - """Test forward pass with batched input.""" - linear = nn.Linear(64, 128) - linear.weight.data.normal_(0, 0.02) - - qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=16) - - x = torch.randn(8, 32, 64) - output = qlinear(x) - - assert output.shape == (8, 32, 128) - - def test_forward_invalid_shape(self): - """Test that forward raises error for invalid input shape.""" - qlinear = QuantLinear(in_features=128, out_features=256, bits=4, symmetric=True, group_size=32) - - x = torch.randn(16, 64) # Wrong input features - - with pytest.raises(AssertionError, match=r"Input shape .* does not match in_features"): - qlinear(x) - - def test_from_tensors_with_precomputed_params(self): - """Test creating QuantLinear with precomputed scales and zero points.""" - weight = torch.randn(256, 128) - - # Compute quantization parameters - quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False) - scales, zero_points = quantizer.find_qparams(weight) - - # Create QuantLinear with precomputed params - qlinear = QuantLinear.from_tensors( - in_features=128, - out_features=256, - weight=weight, - bits=4, - symmetric=True, - group_size=32, - scales=scales, - zero_points=zero_points, - bias=False, - ) - - assert qlinear.scales.shape == scales.shape - assert torch.all(qlinear.scales == scales) - - def test_from_tensors_quantized_weights(self): - """Test creating QuantLinear from already quantized weights.""" - weight = torch.randn(256, 128) - - # Quantize weights - quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False) - scales, zero_points = quantizer.find_qparams(weight) - qweight = quantizer.quantize(weight, scales, zero_points) - - # Create QuantLinear from quantized weights - qlinear = QuantLinear.from_tensors( - in_features=128, - out_features=256, - weight=qweight, - bits=4, - symmetric=True, - group_size=32, - scales=scales, - zero_points=zero_points, - quantized=True, - bias=False, - ) - - assert qlinear.cols == 128 - assert qlinear.rows == 256 - - def test_from_tensors_missing_params_for_quantized(self): - """Test that error is raised when params are missing for quantized weights.""" - qweight = torch.randint(0, 16, (256, 128)) - - with pytest.raises(ValueError, match="scales and/or zero_points missing"): - QuantLinear.from_tensors( - in_features=128, - out_features=256, - weight=qweight, - bits=4, - symmetric=True, - group_size=32, - quantized=True, - bias=False, - ) - - def test_from_tensors_invalid_symmetric_zero_points(self): - """Test that error is raised for invalid zero points in symmetric quantization.""" - weight = torch.randn(256, 128) - - quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False) - scales, zero_points = quantizer.find_qparams(weight) - - # Modify zero points to be invalid for symmetric quantization - zero_points.fill_(0) - - with pytest.raises(ValueError, match="Zero points must be equal to midq for symmetric quantization"): - QuantLinear.from_tensors( - in_features=128, - out_features=256, - weight=weight, - bits=4, - symmetric=True, - group_size=32, - scales=scales, - zero_points=zero_points, - bias=False, - ) - - def test_extra_repr(self): - """Test string representation of QuantLinear.""" - qlinear = QuantLinear( - in_features=128, - out_features=256, - bits=4, - symmetric=True, - group_size=32, - bias=True, - ) - - repr_str = qlinear.extra_repr() - assert "in_features=128" in repr_str - assert "out_features=256" in repr_str - assert "bits=4" in repr_str - assert "symmetric=True" in repr_str - assert "group_size=32" in repr_str - assert "bias=True" in repr_str - - -class TestQuantEmbedding: - def test_basic_initialization(self): - """Test basic QuantEmbedding initialization.""" - qembed = QuantEmbedding( - num_embeddings=1000, - embedding_dim=256, - bits=4, - symmetric=True, - group_size=32, - ) - assert qembed.rows == 1000 - assert qembed.cols == 256 - assert qembed.padding_idx is None - - def test_initialization_with_padding_idx(self): - """Test QuantEmbedding initialization with padding_idx.""" - qembed = QuantEmbedding( - num_embeddings=1000, - embedding_dim=256, - bits=4, - symmetric=True, - group_size=32, - padding_idx=0, - ) - assert qembed.padding_idx == 0 - - def test_from_module_basic(self): - """Test creating QuantEmbedding from nn.Embedding.""" - embedding = nn.Embedding(1000, 256) - embedding.weight.data.normal_(0, 0.02) - - qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32) - - assert qembed.rows == 1000 - assert qembed.cols == 256 - assert qembed.padding_idx is None - - def test_from_module_with_padding_idx(self): - """Test creating QuantEmbedding with padding_idx.""" - embedding = nn.Embedding(1000, 256, padding_idx=0) - embedding.weight.data.normal_(0, 0.02) - - qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32) - - assert qembed.padding_idx == 0 - - @pytest.mark.parametrize("bits", [2, 4, 8]) - @pytest.mark.parametrize("symmetric", [True, False]) - def test_from_module_quantization_accuracy(self, bits, symmetric): - """Test that quantization/dequantization is reasonably accurate.""" - embedding = nn.Embedding(100, 64) - embedding.weight.data.normal_(0, 0.02) - - qembed = QuantEmbedding.from_module(embedding, bits=bits, symmetric=symmetric, group_size=16) - - # Dequantize and compare - dequantized = qembed.unpack_and_dequantize(qembed.qweight, qembed.scales, qembed.qzeros) - - # Check shapes match - assert dequantized.shape == embedding.weight.shape - - # Check that dequantized weights are close to original - max_diff = torch.max(torch.abs(dequantized - embedding.weight)) - tolerance = 0.1 if bits == 4 else 0.05 - assert max_diff < tolerance, f"Max difference {max_diff} exceeds tolerance {tolerance}" - - def test_forward_shape(self): - """Test that forward pass produces correct output shape.""" - embedding = nn.Embedding(1000, 256) - embedding.weight.data.normal_(0, 0.02) - - qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32) - - x = torch.randint(0, 1000, (16,)) - output = qembed(x) - - assert output.shape == (16, 256) - - def test_forward_2d_input(self): - """Test forward pass with 2D input.""" - embedding = nn.Embedding(1000, 256) - embedding.weight.data.normal_(0, 0.02) - - qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32) - - x = torch.randint(0, 1000, (8, 32)) - output = qembed(x) - - assert output.shape == (8, 32, 256) - - def test_forward_3d_input(self): - """Test forward pass with 3D input.""" - embedding = nn.Embedding(1000, 256) - embedding.weight.data.normal_(0, 0.02) - - qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32) - - x = torch.randint(0, 1000, (4, 8, 32)) - output = qembed(x) - - assert output.shape == (4, 8, 32, 256) - - def test_from_tensors_with_precomputed_params(self): - """Test creating QuantEmbedding with precomputed scales and zero points.""" - weight = torch.randn(1000, 256) - - # Compute quantization parameters - quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False) - scales, zero_points = quantizer.find_qparams(weight) - - # Create QuantEmbedding with precomputed params - qembed = QuantEmbedding.from_tensors( - num_embeddings=1000, - embedding_dim=256, - weight=weight, - bits=4, - symmetric=True, - group_size=32, - scales=scales, - zero_points=zero_points, - ) - - assert qembed.scales.shape == scales.shape - assert torch.all(qembed.scales == scales) - - def test_from_tensors_quantized_weights(self): - """Test creating QuantEmbedding from already quantized weights.""" - weight = torch.randn(1000, 256) - - # Quantize weights - quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False) - scales, zero_points = quantizer.find_qparams(weight) - qweight = quantizer.quantize(weight, scales, zero_points) - - # Create QuantEmbedding from quantized weights - qembed = QuantEmbedding.from_tensors( - num_embeddings=1000, - embedding_dim=256, - weight=qweight, - bits=4, - symmetric=True, - group_size=32, - scales=scales, - zero_points=zero_points, - quantized=True, - ) - - assert qembed.rows == 1000 - assert qembed.cols == 256 - - def test_extra_repr(self): - """Test string representation of QuantEmbedding.""" - qembed = QuantEmbedding( - num_embeddings=1000, - embedding_dim=256, - bits=4, - symmetric=True, - group_size=32, - padding_idx=0, - ) - - repr_str = qembed.extra_repr() - assert "1000" in repr_str - assert "256" in repr_str - assert "bits=4" in repr_str - assert "symmetric=True" in repr_str - assert "group_size=32" in repr_str - assert "padding_idx=0" in repr_str - - -class TestQuantModuleIntegration: - def test_quantlinear_forward_backward_compatibility(self): - """Test that QuantLinear produces similar results to nn.Linear.""" - # Create a linear layer - linear = nn.Linear(128, 64, bias=True) - linear.weight.data.normal_(0, 0.02) - linear.bias.data.zero_() - - # Create quantized version - qlinear = QuantLinear.from_module(linear, bits=8, symmetric=True, group_size=32) - - # Test with same input - x = torch.randn(16, 128) - - # Forward pass - linear_out = linear(x) - qlinear_out = qlinear(x) - - # Outputs should be close (not exact due to quantization) - # For 8-bit, we expect decent accuracy - max_diff = torch.max(torch.abs(linear_out - qlinear_out)) - relative_error = max_diff / torch.max(torch.abs(linear_out)) - assert relative_error < 0.1, f"Relative error {relative_error} too large" - - def test_quantembedding_forward_backward_compatibility(self): - """Test that QuantEmbedding produces similar results to nn.Embedding.""" - # Create an embedding layer - embedding = nn.Embedding(100, 64) - embedding.weight.data.normal_(0, 0.02) - - # Create quantized version - qembed = QuantEmbedding.from_module(embedding, bits=8, symmetric=True, group_size=16) - - # Test with same input - x = torch.randint(0, 100, (16,)) - - # Forward pass - embed_out = embedding(x) - qembed_out = qembed(x) - - # Outputs should be close (not exact due to quantization) - max_diff = torch.max(torch.abs(embed_out - qembed_out)) - relative_error = max_diff / torch.max(torch.abs(embed_out)) - assert relative_error < 0.1, f"Relative error {relative_error} too large" - - def test_device_placement(self): - """Test that QuantLinear respects device placement.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - device = torch.device("cuda") - qlinear = QuantLinear( - in_features=128, - out_features=256, - bits=4, - symmetric=True, - group_size=32, - device=device, - ) - - assert qlinear.qweight.device.type == device.type - assert qlinear.scales.device.type == device.type - if qlinear.bias is not None: - assert qlinear.bias.device.type == device.type - - def test_dtype_consistency(self): - """Test that QuantLinear maintains dtype consistency.""" - dtype = torch.float16 - qlinear = QuantLinear( - in_features=128, - out_features=256, - bits=4, - symmetric=True, - group_size=32, - dtype=dtype, - ) - - assert qlinear.scales.dtype == dtype - if qlinear.bias is not None: - assert qlinear.bias.dtype == dtype diff --git a/test/common/quant/test_patterns.py b/test/common/quant/test_patterns.py new file mode 100644 index 0000000000..5ecc4e59b0 --- /dev/null +++ b/test/common/quant/test_patterns.py @@ -0,0 +1,146 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import pytest + +from olive.common.quant.patterns import _assert_regex_safe, is_regex_pattern, match_override, match_skip + + +class TestIsRegexPattern: + def test_returns_true_for_re_prefix(self): + assert is_regex_pattern("re:foo") + + def test_returns_false_for_plain_string(self): + assert not is_regex_pattern("model.embed_tokens") + + def test_returns_false_for_non_string(self): + assert not is_regex_pattern(None) # type: ignore[arg-type] + + +class TestMatchOverride: + def test_returns_none_when_no_patterns(self): + assert match_override("foo", []) is None + assert match_override("foo", None) is None + + def test_literal_equality(self): + assert match_override("model.embed_tokens", ["model.embed_tokens"]) == "model.embed_tokens" + assert match_override("model.embed", ["model.embed_tokens"]) is None + + def test_regex_fullmatch(self): + assert match_override("layers.0.experts.gate_proj", ["re:layers\\.\\d+\\.experts\\..*"]) == ( + "re:layers\\.\\d+\\.experts\\..*" + ) + # not a fullmatch + assert match_override("prefix_layers.0.experts", ["re:layers\\.\\d+\\.experts"]) is None + + def test_first_matching_pattern_wins_in_insertion_order(self): + # both match — the first in config (insertion) order wins, regardless of length + patterns = ["re:.*\\.experts\\..*", "re:layers\\.\\d+\\.experts\\.gate_proj"] + match = match_override("layers.0.experts.gate_proj", patterns) + assert match == "re:.*\\.experts\\..*" + # reversing the order flips the winner + assert match_override("layers.0.experts.gate_proj", list(reversed(patterns))) == ( + "re:layers\\.\\d+\\.experts\\.gate_proj" + ) + + def test_tie_break_is_deterministic(self): + # same-length patterns: first (and here only) match wins + patterns = ["re:foo.bar", "re:foo.baz"] + # only the first matches + assert match_override("foo.bar", patterns) == "re:foo.bar" + + def test_first_match_wins_even_when_later_pattern_is_more_specific(self): + # A broad ``re:.*`` placed first wins over a later literal, per insertion order. + patterns = ["re:.*", "model.embed_tokens"] + assert match_override("model.embed_tokens", patterns) == "re:.*" + # But a literal placed first wins over a later broad regex. + assert match_override("model.embed_tokens", ["model.embed_tokens", "re:.*"]) == "model.embed_tokens" + + +class TestMatchSkip: + def test_substring_match_for_plain_string(self): + # substring preserves HF semantics + assert match_skip("model.layers.0.experts.gate_proj", ["experts"]) + assert match_skip("model.embed_tokens", ["embed_tokens"]) + assert not match_skip("model.layers.0.attn.q_proj", ["experts"]) + + def test_regex_fullmatch_for_re_prefix(self): + assert match_skip("model.layers.0.experts.router.gate", ["re:.*\\.router\\.gate"]) + assert not match_skip("model.layers.0.experts.router.gate.bias", ["re:.*\\.router\\.gate"]) + + def test_empty_pattern_does_not_match(self): + assert not match_skip("foo", [""]) + + def test_none_or_empty_patterns(self): + assert not match_skip("foo", None) + assert not match_skip("foo", []) + + @pytest.mark.parametrize( + ("name", "patterns", "expected"), + [ + ("model.layers.0.experts.0.w1", ["experts"], True), + ("model.layers.0.attn.q_proj", ["experts"], False), + ("router.gate", ["re:router\\.gate"], True), + ("router_gate", ["re:router\\.gate"], False), + ], + ) + def test_parametrized(self, name, patterns, expected): + assert match_skip(name, patterns) is expected + + +class TestRegexSafetyBound: + """Design item 8: ``re:`` patterns must have a real adversarial-input safety bound.""" + + @pytest.mark.parametrize( + "adversarial", + [ + "(a+)+$", + "(a*)*$", + "(a+)*$", + "(a*)+$", + "(a|a)+$", + "(a|aa)+$", + "([a-z]+)+$", + "(x+x+)+y", + "(.*)*$", + "(a+){2,}", + "(a{1,2})+$", + "(a{1,2}){2,}", + "(a?)+$", + "((a|aa))+$", # round-2 bypass: alternation nested inside an inner group + "(a(b|c))+$", # alternation nested one level deep, still unsafe + ], + ) + def test_rejects_nested_unbounded_quantifiers(self, adversarial): + with pytest.raises(ValueError, match="nested"): + _assert_regex_safe(adversarial) + + def test_rejects_overlong_pattern(self): + with pytest.raises(ValueError, match="too long"): + _assert_regex_safe("a" * 500) + + @pytest.mark.parametrize( + "safe", + [ + "layers\\.\\d+\\.experts\\.gate_proj", + ".*\\.experts\\..*", + "model\\.layers\\.\\d+\\.mlp\\.(gate|up|down)_proj", + "(abc)+def", # quantified group but body has no unbounded quantifier + "a{2,4}b+", + "[a-z]+\\.[0-9]+", + "(a(bc))+$", # nested group, no alternation/repetition anywhere in the body + "(ab(cd)ef)+$", # nested group with plain literals only, still no alternation + ], + ) + def test_accepts_safe_patterns(self, safe): + # Should not raise, and should compile / match through the public API. + _assert_regex_safe(safe) + + def test_adversarial_pattern_rejected_through_match_apis(self): + # The safety bound is enforced when the pattern is actually used, not only when + # validated directly. + with pytest.raises(ValueError, match="nested"): + match_override("aaaaaaaaaaaaaaaaaaaa!", ["re:(a+)+$"]) + with pytest.raises(ValueError, match="nested"): + match_skip("aaaaaaaaaaaaaaaaaaaa!", ["re:(a+)+$"]) diff --git a/test/common/quant/test_selection.py b/test/common/quant/test_selection.py new file mode 100644 index 0000000000..4c5814c0c5 --- /dev/null +++ b/test/common/quant/test_selection.py @@ -0,0 +1,767 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=protected-access +"""Tests for ``olive.common.quant.selection.iter_quant_targets``.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from olive.common.quant.selection import iter_quant_targets + + +class _Toy(nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = nn.Embedding(16, 8) + self.linear = nn.Linear(8, 8, bias=False) + self.lm_head = nn.Linear(8, 16, bias=False) + + def get_input_embeddings(self): + return self.embed_tokens + + def get_output_embeddings(self): + return self.lm_head + + +def _names(targets): + return sorted(full_name for _, _, full_name in targets) + + +def test_default_skips_lm_head_and_embeds(): + m = _Toy() + targets = list(iter_quant_targets(m, quantize_lm_head=False, quantize_embeds=False, quantize_moe=False)) + assert _names(targets) == ["linear"] + + +def test_include_lm_head_and_embeds(): + m = _Toy() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(targets) == ["embed_tokens", "linear", "lm_head"] + + +class _MultiEmbed(nn.Module): + """A model with more than one nn.Embedding, like a legacy BERT/GPT-2 style model.""" + + def __init__(self): + super().__init__() + self.embed_tokens = nn.Embedding(16, 8) + self.position_embeddings = nn.Embedding(32, 8) + self.linear = nn.Linear(8, 8, bias=False) + + def get_input_embeddings(self): + return self.embed_tokens + + +def test_embeds_true_targets_only_input_embeddings_when_resolvable(): + """D: quantize_embeds=True should target ONLY get_input_embeddings(), not every nn.Embedding.""" + m = _MultiEmbed() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(targets) == ["embed_tokens", "linear"] + + +def test_embeds_false_still_skips_all_embeddings(): + """D: quantize_embeds=False must still skip ALL nn.Embedding modules (loophole prevention).""" + m = _MultiEmbed() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=False, quantize_moe=False)) + assert _names(targets) == ["linear"] + + +class _NoInputEmbedsAccessor(nn.Module): + """Synthetic fixture without get_input_embeddings -- fallback path.""" + + def __init__(self): + super().__init__() + self.embed_tokens = nn.Embedding(16, 8) + self.position_embeddings = nn.Embedding(32, 8) + self.linear = nn.Linear(8, 8, bias=False) + + +def test_embeds_true_falls_back_to_all_embeddings_without_accessor(): + """D fallback: when get_input_embeddings is unavailable, retain the broad behavior.""" + m = _NoInputEmbedsAccessor() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(targets) == ["embed_tokens", "linear", "position_embeddings"] + + +def test_skip_patterns_filter_by_name(): + m = _Toy() + targets = list( + iter_quant_targets( + m, + quantize_lm_head=True, + quantize_embeds=True, + quantize_moe=False, + skip_patterns=["re:.*_head"], + ) + ) + assert _names(targets) == ["embed_tokens", "linear"] + + +def test_extra_skip_modules_skip_by_identity(): + m = _Toy() + targets = list( + iter_quant_targets( + m, + quantize_lm_head=True, + quantize_embeds=False, + quantize_moe=False, + extra_skip_modules={m.linear}, + ) + ) + assert _names(targets) == ["lm_head"] + + +def test_already_quantized_param_is_skipped(): + from olive.common.quant.tensor import QuantTensor + + m = _Toy() + qt = QuantTensor.from_packed( + qweight=torch.zeros((8, 4), dtype=torch.uint8), + scales=torch.zeros((8, 1), dtype=torch.float32), + qzeros=None, + bits=4, + group_size=8, + symmetric=True, + shape=(8, 8), + dtype=torch.float32, + ) + m.linear.weight = nn.Parameter(qt, requires_grad=False) + + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=False, quantize_moe=False)) + assert _names(targets) == ["lm_head"] + + +class _ExpertList(nn.Module): + def __init__(self): + super().__init__() + self.experts = nn.ModuleList([nn.Linear(8, 8, bias=False) for _ in range(2)]) + + +def test_moe_disabled_skips_submodules_under_experts(monkeypatch): + """When ``quantize_moe=False``, every nn.Module under an experts subtree is skipped.""" + from olive.common.hf import wrapper as wrapper_mod + + class FakeLayerWrapper: + def __init__(self, experts, name): + self._experts = experts + self._name = name + + def get_experts(self, return_name=True): + return (self._experts, self._name) if return_name else self._experts + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [FakeLayerWrapper(self.model.experts, "experts")] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + m = _ExpertList() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(targets) == [] + + +def test_moe_enabled_yields_3d_fused_params(monkeypatch): + from olive.common.hf import wrapper as wrapper_mod + + class FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False) + self.down_proj = nn.Parameter(torch.zeros(4, 16, 8), requires_grad=False) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.experts = FusedExperts() + + class FakeLayerWrapper: + def __init__(self, experts, name): + self._experts = experts + self._name = name + + def get_experts(self, return_name=True): + return (self._experts, self._name) if return_name else self._experts + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [FakeLayerWrapper(self.model.experts, "experts")] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + m = _Model() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + fused = sorted((full_name, tuple(module._parameters[pname].shape)) for module, pname, full_name in targets) + assert fused == [ + ("experts.down_proj", (4, 16, 8)), + ("experts.gate_up_proj", (4, 8, 16)), + ] + + +def _install_fake_wrapper(monkeypatch, experts_by_layer): + """Patch ModelWrapper.from_model to expose ``experts_by_layer`` (list of (module, name)).""" + from olive.common.hf import wrapper as wrapper_mod + + class FakeLayerWrapper: + def __init__(self, experts, name): + self._experts = experts + self._name = name + + def get_experts(self, return_name=True): + return (self._experts, self._name) if return_name else self._experts + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [FakeLayerWrapper(e, n) for e, n in experts_by_layer] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + +def test_moe_enabled_yields_only_3d_weights_and_skips_2d_bias(monkeypatch): + """Regression (gpt-oss gap): 2D bias params on a fused experts module must NOT be quantized.""" + + class FusedExpertsWithBias(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False) + self.gate_up_proj_bias = nn.Parameter(torch.zeros(4, 8), requires_grad=False) # 2D bias + self.down_proj = nn.Parameter(torch.zeros(4, 16, 8), requires_grad=False) + self.down_proj_bias = nn.Parameter(torch.zeros(4, 16), requires_grad=False) # 2D bias + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.experts = FusedExpertsWithBias() + + m = _Model() + _install_fake_wrapper(monkeypatch, [(m.experts, "experts")]) + + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + assert _names(targets) == ["experts.down_proj", "experts.gate_up_proj"] + # every yielded param is 3D (no 2D bias slipped in) + assert all(module._parameters[pname].dim() == 3 for module, pname, _ in targets) + + +def test_shared_expert_gate_excluded_from_quantization(monkeypatch): + """Verify ``SHARED_EXPERT_GATE``-mapped modules are excluded from quantization. + + Covers qwen2_moe/qwen3_5_moe/qwen3_next-style single-row sigmoid gates, which must be + excluded the same way routers are, regardless of ``quantize_moe``. + """ + from olive.common.hf import wrapper as wrapper_mod + + class FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.experts = FusedExperts() + self.shared_expert_gate = nn.Linear(8, 1, bias=False) + self.other_linear = nn.Linear(8, 8, bias=False) + + class FakeLayerWrapper: + def __init__(self, experts, gate): + self._experts = experts + self._gate = gate + + def get_experts(self, return_name=True): + return (self._experts, "experts") if return_name else self._experts + + def get_shared_expert_gate(self, return_name=True): + return (self._gate, "shared_expert_gate") if return_name else self._gate + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [FakeLayerWrapper(self.model.experts, self.model.shared_expert_gate)] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + m = _Model() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + assert _names(targets) == ["experts.gate_up_proj", "other_linear"] + + +def test_mamba_linear_attn_excluded_from_quantization(monkeypatch): + """Verify ``MAMBA``-mapped modules are excluded from the generic 2D quantization walk. + + Includes qwen3_5_moe's ``linear_attn`` GatedDeltaNet block, which must be excluded + unconditionally. + """ + from olive.common.hf import wrapper as wrapper_mod + + class FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False) + + class LinearAttn(nn.Module): + def __init__(self): + super().__init__() + self.in_proj_qkv = nn.Linear(8, 24, bias=False) + self.out_proj = nn.Linear(8, 8, bias=False) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.experts = FusedExperts() + self.linear_attn = LinearAttn() + self.other_linear = nn.Linear(8, 8, bias=False) + + class FakeLayerWrapper: + def __init__(self, experts, mamba): + self._experts = experts + self._mamba = mamba + + def get_experts(self, return_name=True): + return (self._experts, "experts") if return_name else self._experts + + def get_mamba(self, return_name=True): + return (self._mamba, "linear_attn") if return_name else self._mamba + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [FakeLayerWrapper(self.model.experts, self.model.linear_attn)] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + m = _Model() + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + assert _names(targets) == ["experts.gate_up_proj", "other_linear"] + + +def test_modulelist_experts_moe_flag_controls_selection(monkeypatch): + """Regression (ModuleList bug): per-expert Linears are quantized iff ``moe=True``.""" + m = _ExpertList() + _install_fake_wrapper(monkeypatch, [(m.experts, "experts")]) + + off = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(off) == [] + + on = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + assert _names(on) == ["experts.0", "experts.1"] + + +def test_fail_closed_when_moe_arch_but_experts_not_discovered(monkeypatch): + """``moe=False`` must fail closed for an MoE arch whose experts can't be resolved.""" + + class _MoEConfig: + num_local_experts = 8 + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.config = _MoEConfig() + self.some_linear = nn.Linear(8, 8, bias=False) + + m = _Model() + # Wrapper resolves no experts (unrecognized architecture). + _install_fake_wrapper(monkeypatch, []) + + import pytest + + with pytest.raises(ValueError, match="Mixture-of-Experts"): + list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + + +def test_fail_closed_for_dbrx_shaped_nested_config(monkeypatch): + """R2-4: DBRX-shaped nested config (``config.ffn_config.moe_num_experts``) must be detected. + + Previously ``_config_indicates_moe`` only checked a hardcoded top-level attribute tuple, + so this nested-config MoE architecture silently slipped through the fail-closed guard. + """ + + class _FfnConfig: + moe_num_experts = 8 + + class _DbrxConfig: + ffn_config = _FfnConfig() + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.config = _DbrxConfig() + self.some_linear = nn.Linear(8, 8, bias=False) + + m = _Model() + _install_fake_wrapper(monkeypatch, []) + + import pytest + + with pytest.raises(ValueError, match="Mixture-of-Experts"): + list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + + +def test_fail_closed_for_llama4_shaped_nested_config(monkeypatch): + """R2-4: Llama4-shaped nested config (``config.text_config.num_local_experts``) is detected.""" + + class _TextConfig: + num_local_experts = 16 + + class _Llama4Config: + text_config = _TextConfig() + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.config = _Llama4Config() + self.some_linear = nn.Linear(8, 8, bias=False) + + m = _Model() + _install_fake_wrapper(monkeypatch, []) + + import pytest + + with pytest.raises(ValueError, match="Mixture-of-Experts"): + list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + + +def test_dense_config_does_not_trigger_moe_guard(monkeypatch): + """A plain dense config (only ``num_hidden_layers``) must not trigger the MoE fail-closed guard.""" + + class _DenseConfig: + num_hidden_layers = 12 + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.config = _DenseConfig() + self.some_linear = nn.Linear(8, 8, bias=False) + + m = _Model() + _install_fake_wrapper(monkeypatch, []) + + # Should not raise -- falls through to the plain 2D walk. + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(targets) == ["some_linear"] + + +def test_zero_experts_does_not_trigger_moe_guard(monkeypatch): + """``num_experts = 0`` (falsy/absent MoE) must not trigger the fail-closed guard.""" + + class _Config: + num_experts = 0 + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.config = _Config() + self.some_linear = nn.Linear(8, 8, bias=False) + + m = _Model() + _install_fake_wrapper(monkeypatch, []) + + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(targets) == ["some_linear"] + + +def test_magicmock_config_does_not_trigger_moe_guard(monkeypatch): + """A ``MagicMock`` config must not be spuriously treated as MoE. + + Guards against `Mock` objects being truthy for every ``getattr`` -- the generic + sub-config sweep must reject non-``int`` values (including further ``Mock`` objects). + """ + from unittest.mock import MagicMock + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.config = MagicMock() + self.some_linear = nn.Linear(8, 8, bias=False) + + m = _Model() + _install_fake_wrapper(monkeypatch, []) + + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert _names(targets) == ["some_linear"] + + +def test_fail_closed_per_layer_when_some_layers_have_router_but_no_experts(monkeypatch): + """M4: fail closed per-layer, not only when ALL layers fail to resolve experts. + + Layer 0 resolves experts fine; layer 1 has a resolvable router/gate but its experts + subtree fails to resolve -- this must raise before any target is yielded, even though + layer 0 succeeded. A dense layer with *no* router (layer 2) is legitimately expert-free + (e.g. DeepSeek's ``first_k_dense_replace``) and must NOT trip the guard. + """ + from olive.common.hf import wrapper as wrapper_mod + + class FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False) + + class _Router(nn.Module): + pass + + class FakeLayerWrapper: + def __init__(self, experts, router): + self._experts = experts + self._router = router + + def get_experts(self, return_name=True): + return (self._experts, "experts") if return_name else self._experts + + def get_router(self, return_name=True): + return (self._router, "gate") if return_name else self._router + + resolved_experts = FusedExperts() + layer0 = FakeLayerWrapper(resolved_experts, _Router()) # router + experts resolve fine + layer1 = FakeLayerWrapper(None, _Router()) # router resolves, experts do NOT -> should raise + layer2 = FakeLayerWrapper(None, None) # no router at all -> legitimately dense, exempt + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [layer0, layer1, layer2] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8, bias=False) + + m = _Model() + + import pytest + + with pytest.raises(ValueError, match="router"): + list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + + +def test_no_fail_closed_when_dense_layers_lack_router(monkeypatch): + """A layer with no router at all (dense layer) must not trip the per-layer guard.""" + from olive.common.hf import wrapper as wrapper_mod + + class FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False) + + class _Router(nn.Module): + pass + + class FakeLayerWrapper: + def __init__(self, experts, router): + self._experts = experts + self._router = router + + def get_experts(self, return_name=True): + return (self._experts, "experts") if return_name else self._experts + + def get_router(self, return_name=True): + return (self._router, "gate") if return_name else self._router + + layer0_experts = FusedExperts() + layer0 = FakeLayerWrapper(layer0_experts, _Router()) # MoE layer, resolves fine + layer1 = FakeLayerWrapper(None, None) # dense layer, no router -> exempt + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [layer0, layer1] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8, bias=False) + self.experts = layer0_experts + + m = _Model() + # Should not raise. + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + assert "experts.gate_up_proj" in _names(targets) + + +def test_no_fail_closed_for_dense_model_with_incidental_gate_submodule(monkeypatch): + """Regression: an incidental ``mlp.gate`` submodule on a dense model must not fail closed. + + ``get_router()`` matches purely on attribute name (default lookup key ``gate``), so + without this guard, any dense model with an unrelated ``mlp.gate`` submodule and no MoE + config signal would incorrectly trip the "partially supported MoE architecture" + fail-closed refusal even when ``quantize_moe=False``. + """ + from olive.common.hf import wrapper as wrapper_mod + + class _Gate(nn.Module): + """An unrelated dense-MLP submodule that happens to be named ``gate``.""" + + class FakeLayerWrapper: + def __init__(self, router): + self._router = router + + def get_experts(self, return_name=True): + return (None, None) if return_name else None + + def get_router(self, return_name=True): + return (self._router, "gate") if return_name else self._router + + layer0 = FakeLayerWrapper(_Gate()) + + class FakeWrapper: + def __init__(self, model): + self.model = model + + def get_layer_wrappers(self): + return [layer0] + + monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m))) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8, bias=False) + + def config(self): + return None + + m = _Model() + m.config = type("Config", (), {"model_type": "dense_model_with_gate"})() + + # Should not raise, and quantize_moe=False should behave like an ordinary dense walk. + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False)) + assert "linear" in _names(targets) + + +def test_gptq_then_rtn_moe_composition_skips_already_quantized(monkeypatch): + """Regression for `Gptq`-then-`Rtn(moe=True, embeds=True)` composition. + + Emulates GPTQ having quantized only the `nn.Linear` layers first (they become + ``QuantTensor``-backed), then runs the RTN selection with ``moe=True`` / ``embeds=True``: + the already-quantized Linears must be skipped (kept as their GPTQ tensors) while the MoE + experts and embeddings are newly selected — no conflict, no double quantization. + """ + from olive.common.quant.tensor import QuantTensor + + class FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False) + self.down_proj = nn.Parameter(torch.zeros(4, 16, 8), requires_grad=False) + + class _MoEModel(nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = nn.Embedding(16, 8) + self.q_proj = nn.Linear(8, 8, bias=False) + self.o_proj = nn.Linear(8, 8, bias=False) + self.experts = FusedExperts() + + def get_input_embeddings(self): + return self.embed_tokens + + m = _MoEModel() + _install_fake_wrapper(monkeypatch, [(m.experts, "experts")]) + + # Emulate GPTQ: quantize only the nn.Linear weights (8-bit here so we can tell them apart). + for linear in (m.q_proj, m.o_proj): + qt = QuantTensor.from_float(linear.weight.data.clone(), bits=8, group_size=8, symmetric=True) + linear.weight = nn.Parameter(qt, requires_grad=False) + + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + + # RTN must only pick up what GPTQ left un-quantized: the embeddings and the MoE experts. + assert _names(targets) == ["embed_tokens", "experts.down_proj", "experts.gate_up_proj"] + + # The GPTQ-quantized Linears are untouched (still 8-bit QuantTensor). + assert isinstance(m.q_proj.weight.data, QuantTensor) + assert m.q_proj.weight.data.bits == 8 + assert isinstance(m.o_proj.weight.data, QuantTensor) + assert m.o_proj.weight.data.bits == 8 + + +class _VisionTower(nn.Module): + def __init__(self): + super().__init__() + self.patch_embed = nn.Linear(8, 8, bias=False) + self.blocks = nn.ModuleList([nn.Linear(8, 8, bias=False)]) + self.merger = nn.Linear(8, 8, bias=False) + + +class _VLModel(nn.Module): + """Composite VL model: decoder under ``model.language_model``, tower under ``model.visual``.""" + + def __init__(self, config): + super().__init__() + self.config = config + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.embed_tokens = nn.Embedding(16, 8) + self.model.language_model.linear = nn.Linear(8, 8, bias=False) + self.model.visual = _VisionTower() + self.lm_head = nn.Linear(8, 16, bias=False) + + def get_input_embeddings(self): + return self.model.language_model.embed_tokens + + def get_output_embeddings(self): + return self.lm_head + + +class _FakeConfig: + def __init__(self, vision_config=None): + self.vision_config = vision_config + + +def test_vision_tower_excluded_for_composite_vl_model(): + """Regression: a VL model's vision tower must not be a PyTorch-side quantization target. + + The vision encoder is quantized separately (int8) downstream; including it here would + double-quantize it. + """ + m = _VLModel(_FakeConfig(vision_config=_FakeConfig())) + targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True)) + assert _names(targets) == [ + "lm_head", + "model.language_model.embed_tokens", + "model.language_model.linear", + ] + assert not any(name.startswith("model.visual") for _, _, name in targets) + + +def test_vision_named_modules_kept_when_config_has_no_vision_config(): + """Text-only models are unaffected: nothing is excluded by name alone.""" + m = _VLModel(_FakeConfig(vision_config=None)) + targets = list(iter_quant_targets(m, quantize_lm_head=False, quantize_embeds=False, quantize_moe=False)) + assert "model.visual.merger" in _names(targets) + + +def test_quantize_vision_true_includes_vision_tower(): + """``quantize_vision=True`` opts back into quantizing the vision tower. + + Some callers quantize a VL model end-to-end in a single PyTorch-side pass, with no + separate downstream (e.g. ONNX) vision-quantization step -- ``quantize_vision=True`` must + let them include the vision tower instead of silently leaving it at full precision. + """ + m = _VLModel(_FakeConfig(vision_config=_FakeConfig())) + targets = list( + iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True, quantize_vision=True) + ) + assert "model.visual.merger" in _names(targets) diff --git a/test/common/quant/test_state_dict.py b/test/common/quant/test_state_dict.py new file mode 100644 index 0000000000..6a10874ba0 --- /dev/null +++ b/test/common/quant/test_state_dict.py @@ -0,0 +1,322 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=protected-access +"""Unit tests for ``olive.common.quant.state_dict``. + +The focus is :func:`refresh_quant_tensor_refs`'s handling of **tied/aliased** hosting +modules (``lm_head`` tied to ``model.embed_tokens`` install the *same* ``QuantTensor`` +object on two modules) and its fail-closed behavior for incomplete checkpoints. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from olive.common.quant.state_dict import install_quant_tensor_param, refresh_quant_tensor_refs +from olive.common.quant.tensor import QuantTensor + +OUT_FEATURES = 8 +IN_FEATURES = 32 +BITS = 4 +GROUP_SIZE = 16 +N_GROUPS = IN_FEATURES // GROUP_SIZE +PACKED_IN = IN_FEATURES * BITS // 8 + + +PACKED_GROUPS = N_GROUPS * BITS // 8 + + +def _placeholder_qt(symmetric: bool = True) -> QuantTensor: + return QuantTensor.from_packed( + qweight=torch.zeros(OUT_FEATURES, PACKED_IN, dtype=torch.uint8), + scales=torch.zeros(OUT_FEATURES, N_GROUPS, dtype=torch.float32), + qzeros=None if symmetric else torch.zeros(OUT_FEATURES, PACKED_GROUPS, dtype=torch.uint8), + bits=BITS, + group_size=GROUP_SIZE, + symmetric=symmetric, + shape=(OUT_FEATURES, IN_FEATURES), + dtype=torch.float32, + is_placeholder=True, + ) + + +def _single_quant_module(symmetric: bool = True) -> nn.Module: + """Build a one-layer model whose ``linear1.weight`` is an unloaded quantized placeholder.""" + model = nn.Sequential() + model.add_module("linear1", nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False)) + install_quant_tensor_param(model.linear1, "weight", _placeholder_qt(symmetric=symmetric)) + return model + + +class _TiedModel(nn.Module): + """Two modules hosting the *same* ``QuantTensor`` object, like tied word embeddings. + + ``src_first`` controls child registration order, which is exactly ``named_modules()`` + iteration order — the thing the old last-write-wins implementation was accidentally + sensitive to. + """ + + def __init__(self, src_first: bool = True): + super().__init__() + src = nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False) + dst = nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False) + if src_first: + self.src = src + self.dst = dst + else: + self.dst = dst + self.src = src + + install_quant_tensor_param(self.src, "weight", _placeholder_qt()) + shared = self.src._parameters["weight"] + + # Mimic ``tie_quant_word_embeddings``: alias ``dst``'s buffer dict entries to the + # buffer *objects* ``src`` holds right now (a one-time snapshot, not a live link) + # and share the very same parameter object. + for name in ("weight_qweight", "weight_scales"): + dst._non_persistent_buffers_set.add(name) + dst._buffers[name] = self.src._buffers[name] + dst._parameters["weight"] = shared + + @property + def shared_param(self) -> QuantTensor: + return self.src._parameters["weight"] + + +def _simulate_checkpoint_load(model: _TiedModel) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Replace ``src``'s buffer objects with freshly-loaded real data (HF ``setattr`` loader). + + ``dst``'s buffer dict keeps pointing at the *old* placeholder objects, reproducing the + exact state the loader leaves behind for tied quantized weights. + """ + reference = torch.randn(OUT_FEATURES, IN_FEATURES) + real = QuantTensor.from_float(reference, bits=BITS, symmetric=True, group_size=GROUP_SIZE) + model.src._buffers["weight_qweight"] = real.qweight + model.src._buffers["weight_scales"] = real.scales + return real.qweight, real.scales, real.to_dense() + + +@pytest.mark.parametrize("src_first", [True, False]) +def test_refresh_binds_shared_quant_tensor_to_freshly_loaded_buffers(src_first: bool): + """The shared QuantTensor must bind to the loaded buffers, whatever the module order. + + Regression: ``refresh_quant_tensor_refs`` used to rebind once per *hosting module*, so + for tied weights the last-visited module won — and the alias module holds stale + placeholder buffers, silently zeroing the weight. + """ + model = _TiedModel(src_first=src_first) + fresh_qweight, fresh_scales, expected = _simulate_checkpoint_load(model) + + refresh_quant_tensor_refs(model, checkpoint_keys={"src.weight_qweight", "src.weight_scales"}) + + shared = model.shared_param + assert shared.qweight is fresh_qweight + assert shared.scales is fresh_scales + assert shared.is_placeholder is False + torch.testing.assert_close(shared.to_dense(), expected) + + # Aliased hosting module must agree at the ``_buffers`` dict level too, so + # ``state_dict()`` / save and live forward computation cannot diverge. + assert model.dst._buffers["weight_qweight"] is fresh_qweight + assert model.dst._buffers["weight_scales"] is fresh_scales + assert model.dst._parameters["weight"] is shared + + +@pytest.mark.parametrize("src_first", [True, False]) +def test_refresh_is_order_independent_with_identity_heuristic(src_first: bool): + """Same fix, but through the ``checkpoint_keys=None`` identity-heuristic path.""" + model = _TiedModel(src_first=src_first) + fresh_qweight, fresh_scales, expected = _simulate_checkpoint_load(model) + + refresh_quant_tensor_refs(model, checkpoint_keys=None) + + shared = model.shared_param + assert shared.qweight is fresh_qweight + assert shared.scales is fresh_scales + assert shared.is_placeholder is False + torch.testing.assert_close(shared.to_dense(), expected) + assert model.dst._buffers["weight_qweight"] is fresh_qweight + + +def test_refresh_raises_when_checkpoint_manifest_omits_a_parameter(): + """Fail closed: a known manifest missing a quantized parameter must raise, not zero it.""" + model = nn.Sequential() + model.add_module("linear1", nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False)) + model.add_module("linear2", nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False)) + install_quant_tensor_param(model.linear1, "weight", _placeholder_qt()) + install_quant_tensor_param(model.linear2, "weight", _placeholder_qt()) + + with pytest.raises(RuntimeError, match=r"missing weights for: linear2\.weight"): + refresh_quant_tensor_refs(model, checkpoint_keys={"linear1.weight_qweight", "linear1.weight_scales"}) + + +def test_refresh_does_not_raise_for_tied_alias_absent_from_manifest(): + """Only the tie *source*'s keys are persisted; the alias must not be reported missing.""" + model = _TiedModel() + _simulate_checkpoint_load(model) + + # ``dst.weight_qweight`` is intentionally absent from the manifest (non-persistent). + refresh_quant_tensor_refs(model, checkpoint_keys={"src.weight_qweight", "src.weight_scales"}) + + assert model.shared_param.is_placeholder is False + + +def test_refresh_does_not_raise_when_loader_remapped_the_checkpoint_key(): + """``checkpoint_keys`` holds *raw on-disk* names, which HF may remap while loading. + + ``save_pretrained(save_original_format=True)`` writes MoE experts as legacy per-expert + keys that the loader fuses back into ``experts.gate_up_proj_qweight``; the fused name is + therefore absent from the raw manifest even though real data *was* loaded. Fail-closed + must require both signals (manifest miss **and** untouched buffers) to agree. + """ + model = nn.Sequential() + model.add_module("linear1", nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False)) + install_quant_tensor_param(model.linear1, "weight", _placeholder_qt()) + + real = QuantTensor.from_float( + torch.randn(OUT_FEATURES, IN_FEATURES), bits=BITS, symmetric=True, group_size=GROUP_SIZE + ) + model.linear1._buffers["weight_qweight"] = real.qweight + model.linear1._buffers["weight_scales"] = real.scales + + refresh_quant_tensor_refs(model, checkpoint_keys={"some.legacy.key_qweight"}) + + shared = model.linear1._parameters["weight"] + assert shared.qweight is real.qweight + assert shared.is_placeholder is False + + +def test_refresh_with_unknown_checkpoint_keys_stays_permissive(): + """Regression: ``checkpoint_keys=None`` keeps the permissive identity heuristic. + + Nothing was loaded, so the parameter stays a placeholder — but we must NOT raise, + because the checkpoint format/manifest is unknown and "not loaded" cannot be + distinguished from "loaded in place". + """ + model = nn.Sequential() + model.add_module("linear1", nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False)) + install_quant_tensor_param(model.linear1, "weight", _placeholder_qt()) + + refresh_quant_tensor_refs(model, checkpoint_keys=None) + + assert model.linear1._parameters["weight"].is_placeholder is True + + +def test_refresh_ignores_modules_without_quant_tensors(): + model = nn.Sequential(nn.Linear(4, 4), nn.ReLU()) + refresh_quant_tensor_refs(model, checkpoint_keys=set()) + + +# ---------------------------------------------------------------------- +# Partial loads must not count as "loaded" (fail-closed false-negative gap). +# ---------------------------------------------------------------------- + + +def test_refresh_raises_when_manifest_has_scales_but_not_qweight(): + """A manifest carrying only *some* of a parameter's buffers is an incomplete checkpoint. + + Regression: the fail-closed check used to look at ``_qweight`` membership alone + (and, worse, treated *any* single signal as proof of a load), so a truncated checkpoint + that contains ``weight_scales`` but not ``weight_qweight`` was reported as loaded — the + placeholder flag was cleared over an all-zero ``qweight``. + """ + model = _single_quant_module() + qt = model.linear1._parameters["weight"] + + # The loader wrote the one key it found; ``weight_qweight`` stays the zero placeholder. + model.linear1._buffers["weight_scales"] = torch.randn(OUT_FEATURES, N_GROUPS) + + with pytest.raises(RuntimeError, match=r"missing weights for: linear1\.weight"): + refresh_quant_tensor_refs(model, checkpoint_keys={"linear1.weight_scales"}) + + assert qt.is_placeholder is True + assert int(qt.qweight.sum()) == 0 + + +def test_refresh_raises_when_manifest_omits_qzeros_for_asymmetric_parameter(): + """``_qzeros`` is mandatory when the QuantTensor is asymmetric.""" + model = _single_quant_module(symmetric=False) + qt = model.linear1._parameters["weight"] + + with pytest.raises(RuntimeError, match=r"missing weights for: linear1\.weight"): + refresh_quant_tensor_refs(model, checkpoint_keys={"linear1.weight_qweight", "linear1.weight_scales"}) + + assert qt.is_placeholder is True + + +def test_refresh_treats_symmetric_parameter_without_qzeros_as_fully_loaded(): + """Symmetric quantization has no ``_qzeros`` buffer — don't demand a key for it.""" + model = _single_quant_module() + + refresh_quant_tensor_refs(model, checkpoint_keys={"linear1.weight_qweight", "linear1.weight_scales"}) + + assert model.linear1._parameters["weight"].is_placeholder is False + + +def test_refresh_treats_asymmetric_parameter_with_all_keys_as_fully_loaded(): + """The asymmetric counterpart: all three keys present means fully loaded.""" + model = _single_quant_module(symmetric=False) + + refresh_quant_tensor_refs( + model, + checkpoint_keys={"linear1.weight_qweight", "linear1.weight_scales", "linear1.weight_qzeros"}, + ) + + assert model.linear1._parameters["weight"].is_placeholder is False + + +def test_refresh_identity_heuristic_ignores_partial_buffer_swap(): + """``checkpoint_keys=None``: a partial buffer swap must not clear ``is_placeholder``. + + The identity heuristic used to OR across ``qweight``/``scales``/``qzeros``, so swapping + any single buffer object marked the whole parameter loaded. It now requires *all* + mandatory buffers to have been replaced. The permissive contract of this code path is + unchanged (nothing is raised when the manifest is unknown) — the parameter simply stays + a placeholder. + """ + model = _single_quant_module() + qt = model.linear1._parameters["weight"] + placeholder_qweight = qt.qweight + + model.linear1._buffers["weight_scales"] = torch.randn(OUT_FEATURES, N_GROUPS) + + refresh_quant_tensor_refs(model, checkpoint_keys=None) + + assert qt.is_placeholder is True + assert qt.qweight is placeholder_qweight + assert int(qt.qweight.sum()) == 0 + + +def test_refresh_identity_heuristic_ignores_missing_qzeros_swap(): + """Asymmetric variant: ``qweight``/``scales`` swapped but ``qzeros`` left stale.""" + model = _single_quant_module(symmetric=False) + qt = model.linear1._parameters["weight"] + placeholder_qzeros = qt.qzeros + + model.linear1._buffers["weight_qweight"] = torch.randint(0, 255, (OUT_FEATURES, PACKED_IN), dtype=torch.uint8) + model.linear1._buffers["weight_scales"] = torch.randn(OUT_FEATURES, N_GROUPS) + + refresh_quant_tensor_refs(model, checkpoint_keys=None) + + assert qt.is_placeholder is True + assert qt.qzeros is placeholder_qzeros + + +def test_refresh_identity_heuristic_accepts_full_buffer_swap(): + """Sanity check that the AND-ed identity heuristic still recognises a real full load.""" + model = _single_quant_module(symmetric=False) + qt = model.linear1._parameters["weight"] + + model.linear1._buffers["weight_qweight"] = torch.randint(0, 255, (OUT_FEATURES, PACKED_IN), dtype=torch.uint8) + model.linear1._buffers["weight_scales"] = torch.randn(OUT_FEATURES, N_GROUPS) + model.linear1._buffers["weight_qzeros"] = torch.randint(0, 255, (OUT_FEATURES, PACKED_GROUPS), dtype=torch.uint8) + + refresh_quant_tensor_refs(model, checkpoint_keys=None) + + assert qt.is_placeholder is False + assert qt.qweight is model.linear1._buffers["weight_qweight"] + assert qt.qzeros is model.linear1._buffers["weight_qzeros"] diff --git a/test/common/quant/test_tensor.py b/test/common/quant/test_tensor.py new file mode 100644 index 0000000000..91e699ddf8 --- /dev/null +++ b/test/common/quant/test_tensor.py @@ -0,0 +1,637 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=redefined-outer-name,not-callable +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from olive.common.quant.tensor import QuantTensor + + +@pytest.fixture +def w2d(): + torch.manual_seed(0) + return torch.randn(64, 128, dtype=torch.float32) + + +@pytest.fixture +def w3d(): + torch.manual_seed(0) + return torch.randn(4, 32, 128, dtype=torch.float32) + + +class TestQuantTensor2D: + def test_shape_dtype_device_preserved(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + assert qt.shape == w2d.shape + assert qt.dtype == w2d.dtype + assert qt.device == w2d.device + assert qt.requires_grad is False + + def test_inner_buffer_layout(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + # 4-bit packed → in_features / 2 + assert qt.qweight.shape == (64, 64) + assert qt.qweight.dtype == torch.uint8 + # groupwise scales: (out, num_groups) + assert qt.scales.shape == (64, 128 // 32) + # symmetric → no zero_points + assert qt.qzeros is None + + def test_asymmetric_has_qzeros(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=False, group_size=32) + assert qt.qzeros is not None + assert qt.qzeros.dtype == torch.uint8 + + def test_to_dense_round_trip(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=False, group_size=32) + dense = qt.to_dense() + assert dense.shape == w2d.shape + # Round trip should be close (4-bit groupwise is reasonably accurate) + assert (dense - w2d).abs().mean().item() < 0.1 + + def test_dispatches_through_f_linear(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + x = torch.randn(2, 128) + out_quant = F.linear(x, qt) + out_dense = F.linear(x, qt.to_dense()) + assert torch.allclose(out_quant, out_dense, atol=1e-5) + + def test_nn_parameter_preserves_subclass(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + p = nn.Parameter(qt, requires_grad=False) + assert isinstance(p, QuantTensor) + assert isinstance(p.data, QuantTensor) + + def test_nn_linear_forward_with_quant_tensor_weight(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + layer = nn.Linear(128, 64, bias=False) + layer.weight = nn.Parameter(qt, requires_grad=False) + x = torch.randn(2, 128) + out_layer = layer(x) + out_ref = F.linear(x, qt.to_dense()) + assert torch.allclose(out_layer, out_ref, atol=1e-5) + + def test_model_to_dtype_propagates(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=False, group_size=32) + + class M(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(128, 64, bias=False) + + m = M() + m.lin.weight = nn.Parameter(qt, requires_grad=False) + m = m.to(torch.float16) + # dtype follows the wrapper subclass; scales are floating-point + assert m.lin.weight.dtype == torch.float16 + assert m.lin.weight.scales.dtype == torch.float16 + # qweight is uint8 — non-floating-point, kept as-is + assert m.lin.weight.qweight.dtype == torch.uint8 + + def test_nn_embedding_forward(self, w2d): + # 64 embeddings of dim 128 + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + emb = nn.Embedding(64, 128) + emb.weight = nn.Parameter(qt, requires_grad=False) + ids = torch.tensor([0, 5, 60]) + out = emb(ids) + out_ref = F.embedding(ids, qt.to_dense()) + assert torch.allclose(out, out_ref, atol=1e-5) + + +class TestQuantTensor3D: + def test_3d_shape(self, w3d): + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + assert qt.shape == w3d.shape + assert qt.qweight.shape == (4, 32, 64) + assert qt.scales.shape == (4, 32, 4) + assert qt.qzeros is not None + assert qt.qzeros.shape == (4, 32, 2) + + def test_3d_round_trip(self, w3d): + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + dense = qt.to_dense() + assert (dense - w3d).abs().mean().item() < 0.1 + + def test_slice_returns_2d_quant_tensor(self, w3d): + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + sliced = qt[2] + assert isinstance(sliced, QuantTensor) + assert sliced.shape == w3d[2].shape + # F.linear over the slice + x = torch.randn(2, 128) + out = F.linear(x, sliced) + out_ref = F.linear(x, qt.to_dense()[2]) + assert torch.allclose(out, out_ref, atol=1e-5) + + +class TestQuantTensorEqual: + """``torch.equal`` must compare packed buffers, never dequantize. + + ``transformers>=5``'s ``PreTrainedModel.tie_weights`` calls ``torch.equal`` on the two + tied word-embedding parameters inside ``from_pretrained``'s ``_finalize_model_loading``, + i.e. *before* the quantizer's ``_process_model_after_weight_loading`` hook runs. At that + point a placeholder ``QuantTensor``'s inner buffers can still be on ``meta``, and the + generic dequantizing fallback hard-fails with + ``NotImplementedError: aten::equal ... with Meta tensors``. + """ + + @staticmethod + def _meta_qt() -> QuantTensor: + return QuantTensor.from_packed( + qweight=torch.zeros(8, 16, dtype=torch.uint8, device="meta"), + scales=torch.zeros(8, 2, dtype=torch.float32, device="meta"), + qzeros=None, + bits=4, + group_size=16, + symmetric=True, + shape=(8, 32), + dtype=torch.float32, + is_placeholder=True, + ) + + def test_tied_meta_quant_tensors_compare_equal_without_crashing(self): + qt = self._meta_qt() + assert torch.equal(qt, qt) is True + + def test_distinct_meta_quant_tensors_are_not_equal(self): + assert torch.equal(self._meta_qt(), self._meta_qt()) is False + + def test_equal_compares_packed_buffers(self, w2d): + a = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + b = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + c = QuantTensor.from_float(w2d + 1.0, bits=4, symmetric=True, group_size=32) + assert torch.equal(a, b) is True + assert torch.equal(a, c) is False + + def test_equal_is_false_for_mismatched_quant_metadata(self, w2d): + a = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + b = QuantTensor.from_float(w2d, bits=8, symmetric=True, group_size=32) + assert torch.equal(a, b) is False + + def test_equal_against_dense_tensor_dequantizes(self, w2d): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + assert torch.equal(qt, qt.to_dense()) is True + assert torch.equal(qt, torch.zeros_like(w2d)) is False + + +class TestQuantTensorOnnxExportGuards: + def test_linear_raises_when_in_onnx_export(self, w2d, monkeypatch): + qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32) + x = torch.randn(1, 128) + # Simulate being inside ONNX export + monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True) + with pytest.raises(RuntimeError, match="QuantTensor cannot be traced"): + F.linear(x, qt) + + +class TestQuantTensor3DExpertRouting: + def test_tensor_index_routing_preserves_quantized_storage(self, w3d): + """Advanced/tensor-index expert selection (how real MoE routes) stays quantized.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + expert_ids = torch.tensor([0, 2, 2, 1]) + selected = qt[expert_ids] + assert isinstance(selected, QuantTensor) + assert selected.shape == (4, *w3d.shape[1:]) + # Values match a dense gather. + ref = qt.to_dense()[expert_ids] + assert torch.allclose(selected.to_dense(), ref, atol=1e-6) + + def test_list_index_routing_preserves_quantized_storage(self, w3d): + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + selected = qt[[0, 3]] + assert isinstance(selected, QuantTensor) + assert selected.shape == (2, *w3d.shape[1:]) + + def test_tuple_leading_only_index_preserves_quantized_storage(self, w3d): + """M2: tuple-form leading-dim indexing ``w[expert_ids, :, :]`` must stay quantized.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + expert_ids = torch.tensor([0, 2]) + selected = qt[expert_ids, :, :] + assert isinstance(selected, QuantTensor) + assert selected.shape == (2, *w3d.shape[1:]) + ref = qt.to_dense()[expert_ids, :, :] + assert torch.allclose(selected.to_dense(), ref, atol=1e-6) + + def test_rank2_tensor_index_raises_instead_of_producing_unusable_quant_tensor(self, w3d): + """#2598 item 4: rank>=2 integer-tensor indices (e.g. (tokens, k) top-k routing) raise. + + They must raise immediately instead of silently producing a >3D ``QuantTensor`` that + can never be dequantized (``to_dense()`` refuses rank > 3) or re-indexed + (``__getitem__`` also refuses rank > 3) -- i.e. a dead-end object. Callers needing a + multi-dim batch of expert ids should flatten to 1-D first + (``weight[expert_ids.flatten()]``) and reshape the *dense output* back afterward. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + idx = torch.tensor([[0, 1], [2, 3]]) + with pytest.raises(RuntimeError, match="Unsupported indexing pattern"): + _ = qt[idx] + + def test_flattened_rank2_index_workaround_stays_quantized_and_dequantizes(self, w3d): + """The documented workaround for a rank>=2 expert-id batch: flatten to 1-D first.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + idx = torch.tensor([[0, 1], [2, 3]]) + selected = qt[idx.flatten()] + assert isinstance(selected, QuantTensor) + assert selected.dim() == 3 + dense = selected.to_dense() + # Reshape the dense *output* back to the original (tokens, k, ...) batch shape. + reshaped = dense.reshape(*idx.shape, *w3d.shape[1:]) + ref = qt.to_dense()[idx] + assert torch.allclose(reshaped, ref, atol=1e-6) + + def test_unsupported_3d_indexing_raises_instead_of_dequantizing(self, w3d): + """Multi-axis / advanced indexing that isn't leading-dim-only must raise, not OOM-dequant.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + with pytest.raises(RuntimeError, match="Unsupported indexing pattern"): + _ = qt[:, 0, :] + + def test_rank2_bool_mask_raises_instead_of_misclassifying(self, w3d): + """Round-2 regression: a rank>1 boolean mask must raise, not silently misclassify shape. + + Previously any boolean tensor was treated as "leading-dim only", producing a + QuantTensor whose ``.shape`` metadata disagreed with ``.to_dense()``'s actual shape. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + mask = torch.ones((4, 8), dtype=torch.bool) + with pytest.raises(RuntimeError, match="Unsupported indexing pattern"): + _ = qt[mask] + + def test_1d_bool_mask_preserves_quantized_storage_and_shape(self, w3d): + """A 1-D boolean mask over the leading (expert) dim is a safe, quantized-preserving index.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + mask = torch.tensor([True, False, True, False]) + selected = qt[mask] + assert isinstance(selected, QuantTensor) + assert selected.shape == (2, *w3d.shape[1:]) + assert selected.to_dense().shape == selected.shape + ref = qt.to_dense()[mask] + assert torch.allclose(selected.to_dense(), ref, atol=1e-6) + + def test_1d_uint8_mask_matching_length_preserves_quantized_storage_and_shape(self, w3d): + """#2598 item 1: a correctly-shaped 1-D uint8 mask is torch's legacy boolean mask form. + + It must be treated identically to a real bool mask (selects, doesn't gather) so the + resulting shape/values match a dense uint8-indexed selection. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + mask = torch.tensor([1, 0, 1, 0], dtype=torch.uint8) + selected = qt[mask] + assert isinstance(selected, QuantTensor) + assert selected.shape == (2, *w3d.shape[1:]) + assert selected.to_dense().shape == selected.shape + ref = qt.to_dense()[mask] + assert torch.allclose(selected.to_dense(), ref, atol=1e-6) + + def test_1d_uint8_mask_wrong_length_raises_instead_of_misclassifying(self, w3d): + """#2598 item 1: a length-mismatched uint8 mask must raise, not silently misclassify.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + mask = torch.tensor([1, 0], dtype=torch.uint8) + with pytest.raises(RuntimeError, match="Unsupported indexing pattern"): + _ = qt[mask] + + def test_0d_uint8_scalar_raises_instead_of_misclassifying(self, w3d): + """#2598 item 1: a 0-D uint8 scalar (e.g. ``expert_idx[0]`` in real MoE routing code). + + Previously fell through to the "any non-float/complex tensor is a safe gather index" + branch and silently produced a shape-inserting result instead of a scalar selection. + Now must raise instead of silently producing a wrong shape. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + idx = torch.tensor(1, dtype=torch.uint8) + with pytest.raises(RuntimeError, match="Unsupported indexing pattern"): + _ = qt[idx] + + def test_0d_int64_scalar_still_selects_a_single_expert(self, w3d): + """Regression guard: the uint8 fix must not affect ordinary 0-D integer scalar indexing.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + idx = torch.tensor(1, dtype=torch.int64) + selected = qt[idx] + assert isinstance(selected, QuantTensor) + assert selected.shape == w3d.shape[1:] + + @pytest.mark.parametrize( + "idx", + [ + 0, + slice(0, 2), + [0, 2, 3], + torch.tensor([0, 2]), + torch.tensor([True, False, True, False]), + ], + ) + def test_getitem_shape_matches_dense_for_every_accepted_index_form(self, w3d, idx): + """Property test: for every accepted index form, ``.shape`` metadata must agree with dense. + + This is the invariant both this bug and any future ``_getitem`` extension must preserve. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + result = qt[idx] + dense_result = qt.to_dense()[idx] + assert tuple(result.shape) == tuple(dense_result.shape) + if isinstance(result, QuantTensor) and result.dim() in (2, 3): + assert torch.allclose(result.to_dense(), dense_result, atol=1e-6) + + def test_unsupported_3d_op_raises_during_onnx_export(self, w3d, monkeypatch): + """Central guard: any 3D QuantTensor op reaching the dense fallback under export must raise. + + ``torch.index_select`` is an op that is not individually special-cased, so it exercises the + central ``_maybe_dense`` rejection rather than an op-specific check. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True) + with pytest.raises(RuntimeError, match=r"ModelBuilder|Mobius|MoE"): + torch.index_select(qt, 1, torch.tensor([0, 1])) + + def test_unsupported_3d_op_raises_in_eager_mode(self, w3d): + """M1: the same central guard must also raise in plain eager mode (no ONNX export). + + Previously an unregistered op on a 3D QuantTensor fell through to the generic + ``__torch_dispatch__`` fallback and fully dequantized the (potentially huge) expert + tensor in eager mode -- an OOM risk this guard must refuse instead. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + with pytest.raises(RuntimeError, match=r"OOM|dequant|refused"): + torch.index_select(qt, 1, torch.tensor([0, 1])) + + @pytest.mark.parametrize( + "op", + [ + lambda qt: qt.transpose(-2, -1), + lambda qt: qt.reshape(4, -1), + lambda qt: qt.view(4, -1), + lambda qt: qt.permute(0, 2, 1), + lambda qt: qt.flatten(), + lambda qt: torch.transpose(qt, -2, -1), + ], + ) + def test_movement_ops_raise_instead_of_silently_misbehaving(self, w3d, op): + """Shape-movement / view ops must raise a clear error rather than produce a malformed tensor.""" + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + with pytest.raises(RuntimeError, match="storage-only"): + op(qt) + + def test_movement_op_under_onnx_export_raises_moe_message(self, w3d, monkeypatch): + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True) + with pytest.raises(RuntimeError, match=r"ModelBuilder|Mobius|MoE"): + qt.transpose(-2, -1) + + +class TestQuantTensorPlaceholderInit: + def test_inplace_initializers_are_noop_on_placeholder(self, w3d): + """Regression: HF's ``_initialize_weights`` calls in-place initializers on placeholders. + + HF's ``PreTrainedModel._initialize_weights`` calls in-place initializers (e.g. + ``nn.init.normal_``) on freshly-installed placeholder QuantTensor params before the + checkpoint's real buffers are loaded. These must be safe (harmless) no-ops rather + than raising (the M1 eager-3D guard would otherwise break real HF model loading for + MoE / fused-3D targets) or silently dequantizing just to throw the result away. + """ + qt = QuantTensor.from_packed( + qweight=torch.zeros(4, 32, 64, dtype=torch.uint8), + scales=torch.zeros(4, 32, 4, dtype=torch.float32), + qzeros=torch.zeros(4, 32, 2, dtype=torch.uint8), + bits=4, + group_size=32, + symmetric=False, + shape=(4, 32, 128), + dtype=torch.float32, + is_placeholder=True, + ) + assert qt.is_placeholder is True + # Must not raise. + torch.nn.init.normal_(qt, mean=0.0, std=0.02) + torch.nn.init.zeros_(qt) + assert isinstance(qt, QuantTensor) + + def test_inplace_initializers_raise_on_real_quant_tensor(self, w3d): + """Round-2 regression: in-place init on a *real* (non-placeholder) QuantTensor must raise. + + Silently no-oping here would let ``torch.nn.init.zeros_(real_qt)`` "succeed" while the + dequantized values stay unchanged -- a silent data-integrity bug. + """ + qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + assert qt.is_placeholder is False + with pytest.raises(RuntimeError, match="In-place initializer"): + torch.nn.init.zeros_(qt) + with pytest.raises(RuntimeError, match="In-place initializer"): + torch.nn.init.normal_(qt, mean=0.0, std=0.02) + + def test_is_placeholder_survives_nn_parameter_detach_and_to(self, w3d): + """The placeholder flag must survive ``nn.Parameter(qt)``'s ``.detach()`` and ``.to(...)``. + + ``torch.nn.Parameter(qt, requires_grad=False)`` for a tensor subclass returns + ``qt.detach()``, which constructs a *new* QuantTensor via ``_apply_fn_to_data``. If the + flag were lost here, HF's ``_initialize_weights`` would raise instead of no-oping on the + placeholder it installs before checkpoint loading. + """ + qt = QuantTensor.from_packed( + qweight=torch.zeros(4, 32, 64, dtype=torch.uint8), + scales=torch.zeros(4, 32, 4, dtype=torch.float32), + qzeros=torch.zeros(4, 32, 2, dtype=torch.uint8), + bits=4, + group_size=32, + symmetric=False, + shape=(4, 32, 128), + dtype=torch.float32, + is_placeholder=True, + ) + p = torch.nn.Parameter(qt, requires_grad=False) + assert isinstance(p, QuantTensor) + assert p.is_placeholder is True + p2 = p.to(torch.float16) + assert isinstance(p2, QuantTensor) + assert p2.is_placeholder is True + assert p[0].is_placeholder is True + + def test_refresh_quant_tensor_refs_clears_placeholder_flag(self): + """``refresh_quant_tensor_refs`` binds real checkpoint buffers and clears the flag.""" + from olive.common.quant.state_dict import refresh_quant_tensor_refs + + qt = QuantTensor.from_packed( + qweight=torch.zeros(64, 64, dtype=torch.uint8), + scales=torch.zeros(64, 4, dtype=torch.float32), + qzeros=None, + bits=4, + group_size=32, + symmetric=True, + shape=(64, 128), + dtype=torch.float32, + is_placeholder=True, + ) + layer = nn.Linear(128, 64, bias=False) + layer.weight = nn.Parameter(qt, requires_grad=False) + assert layer.weight.is_placeholder is True + + # Simulate HF's checkpoint buffer load: real ``_qweight`` / ``_scales`` buffers + # land in ``_buffers`` under the naming convention ``buffer_names`` expects. + layer.register_buffer("weight_qweight", torch.randint(0, 255, (64, 64), dtype=torch.uint8)) + layer.register_buffer("weight_scales", torch.randn(64, 4, dtype=torch.float32)) + + refresh_quant_tensor_refs(layer) + assert layer.weight.is_placeholder is False + + def test_refresh_quant_tensor_refs_keeps_placeholder_when_key_missing(self): + """Round-3 regression (#2598 item 3): a missing checkpoint key must not clear the flag. + + ``refresh_quant_tensor_refs`` is called once for the *whole model* after HF's loader + finishes, with no per-parameter "was this key actually in the checkpoint" signal. A + parameter whose key was missing from the checkpoint keeps the exact placeholder buffer + objects installed by ``install_quant_tensor_param`` -- only a real + ``load_state_dict(..., assign=True)`` replaces the buffer objects. Unconditionally + clearing ``is_placeholder`` here would make a later in-place initializer (which HF may + still call for missing-key parameters) raise instead of safely no-oping. + """ + from olive.common.quant.state_dict import refresh_quant_tensor_refs + + qt = QuantTensor.from_packed( + qweight=torch.zeros(64, 64, dtype=torch.uint8), + scales=torch.zeros(64, 4, dtype=torch.float32), + qzeros=None, + bits=4, + group_size=32, + symmetric=True, + shape=(64, 128), + dtype=torch.float32, + is_placeholder=True, + ) + layer = nn.Linear(128, 64, bias=False) + # ``nn.Parameter(qt)`` for a tensor subclass returns ``qt.detach()``, which produces + # *new* (storage-aliased) inner tensor objects -- so read them back off + # ``layer.weight`` (like ``install_quant_tensor_param`` does), not off ``qt`` itself. + layer.weight = nn.Parameter(qt, requires_grad=False) + # Register the placeholder buffers aliasing the exact same tensor objects the + # installed QuantTensor parameter already holds -- i.e. simulate "checkpoint load + # ran, but this parameter's key was missing so its buffers were never reassigned". + layer.register_buffer("weight_qweight", layer.weight.qweight) + layer.register_buffer("weight_scales", layer.weight.scales) + + refresh_quant_tensor_refs(layer) + assert layer.weight.is_placeholder is True + # Must still be a safe no-op, not a raise. + torch.nn.init.zeros_(layer.weight) + + def test_refresh_quant_tensor_refs_checkpoint_keys_clears_flag_despite_unchanged_buffer_identity(self): + """``checkpoint_keys`` is authoritative even when the loader mutated buffers in place. + + The buffer-identity heuristic (the ``checkpoint_keys=None`` fallback) only detects a + real load when the loader *replaces* the buffer object (e.g. via ``setattr``); a + loader that instead does an in-place ``.copy_()`` into the existing buffer object would + leave identity unchanged and be missed. Passing the checkpoint's own key manifest sidesteps + that entirely by checking exact key membership instead of object identity. + """ + from olive.common.quant.state_dict import refresh_quant_tensor_refs + + qt = QuantTensor.from_packed( + qweight=torch.zeros(64, 64, dtype=torch.uint8), + scales=torch.zeros(64, 4, dtype=torch.float32), + qzeros=None, + bits=4, + group_size=32, + symmetric=True, + shape=(64, 128), + dtype=torch.float32, + is_placeholder=True, + ) + layer = nn.Linear(128, 64, bias=False) + layer.weight = nn.Parameter(qt, requires_grad=False) + # Simulate an in-place ``.copy_()`` loader: real data is written into the *same* + # buffer objects rather than replacing them, so identity never changes. + layer.weight.qweight.copy_(torch.randint(0, 255, (64, 64), dtype=torch.uint8)) + layer.weight.scales.copy_(torch.randn(64, 4, dtype=torch.float32)) + layer.register_buffer("weight_qweight", layer.weight.qweight) + layer.register_buffer("weight_scales", layer.weight.scales) + + # Without checkpoint_keys, the identity heuristic is fooled (buffers were mutated, + # not replaced) and incorrectly keeps the placeholder flag set. + refresh_quant_tensor_refs(layer) + assert layer.weight.is_placeholder is True + + # With the checkpoint's own key manifest, membership is checked directly and + # correctly clears the flag regardless of how the loader wrote the data. + refresh_quant_tensor_refs(layer, checkpoint_keys={"weight_qweight", "weight_scales"}) + assert layer.weight.is_placeholder is False + + def test_refresh_quant_tensor_refs_raises_when_checkpoint_key_absent(self): + """Fail closed when a known manifest omits a quantized parameter. + + A missing key must raise instead of silently leaving the model with zero-filled + placeholder weights. + """ + from olive.common.quant.state_dict import refresh_quant_tensor_refs + + qt = QuantTensor.from_packed( + qweight=torch.zeros(64, 64, dtype=torch.uint8), + scales=torch.zeros(64, 4, dtype=torch.float32), + qzeros=None, + bits=4, + group_size=32, + symmetric=True, + shape=(64, 128), + dtype=torch.float32, + is_placeholder=True, + ) + layer = nn.Linear(128, 64, bias=False) + layer.weight = nn.Parameter(qt, requires_grad=False) + layer.register_buffer("weight_qweight", layer.weight.qweight) + layer.register_buffer("weight_scales", layer.weight.scales) + + # Checkpoint manifest doesn't mention this parameter's keys at all. + with pytest.raises(RuntimeError, match="missing weights for: weight"): + refresh_quant_tensor_refs(layer, checkpoint_keys={"some_other_param_qweight", "some_other_param_scales"}) + assert layer.weight.is_placeholder is True + + def test_copy_into_placeholder_clears_flag(self, w3d): + """Round-3 regression (#2598 item 2): ``copy_`` must propagate ``is_placeholder``. + + Copying real (non-placeholder) data into a placeholder ``QuantTensor`` makes it real + too. If ``is_placeholder`` were left ``True``, a later in-place initializer (e.g. a + module re-init call) could still silently no-op and discard the just-copied real data. + """ + placeholder = QuantTensor.from_packed( + qweight=torch.zeros(4, 32, 64, dtype=torch.uint8), + scales=torch.zeros(4, 32, 4, dtype=torch.float32), + qzeros=torch.zeros(4, 32, 2, dtype=torch.uint8), + bits=4, + group_size=32, + symmetric=False, + shape=(4, 32, 128), + dtype=torch.float32, + is_placeholder=True, + ) + real = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + assert placeholder.is_placeholder is True + assert real.is_placeholder is False + + placeholder.copy_(real) + assert placeholder.is_placeholder is False + # Data integrity: an in-place initializer must now raise (real data present), + # not silently no-op. + with pytest.raises(RuntimeError, match="In-place initializer"): + torch.nn.init.zeros_(placeholder) + + def test_copy_from_placeholder_keeps_flag(self, w3d): + """Copying a placeholder's (throwaway) data into another QuantTensor keeps it a placeholder.""" + src_placeholder = QuantTensor.from_packed( + qweight=torch.zeros(4, 32, 64, dtype=torch.uint8), + scales=torch.zeros(4, 32, 4, dtype=torch.float32), + qzeros=torch.zeros(4, 32, 2, dtype=torch.uint8), + bits=4, + group_size=32, + symmetric=False, + shape=(4, 32, 128), + dtype=torch.float32, + is_placeholder=True, + ) + dst_real = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32) + assert dst_real.is_placeholder is False + + dst_real.copy_(src_placeholder) + assert dst_real.is_placeholder is True + # Must not raise now that it's (again) a placeholder. + torch.nn.init.zeros_(dst_real) diff --git a/test/common/quant/test_utils.py b/test/common/quant/test_utils.py index 802eab746d..8e9ea928b1 100644 --- a/test/common/quant/test_utils.py +++ b/test/common/quant/test_utils.py @@ -63,7 +63,7 @@ def test_initialization(self, bits, symmetric, group_size, signed): def test_invalid_bits(self): """Test that invalid bits raise AssertionError.""" - with pytest.raises(AssertionError, match="Only 4-bit and 8-bit quantization supported"): + with pytest.raises(AssertionError, match="Only 2-bit, 4-bit and 8-bit quantization supported"): WeightQuantizer(bits=16, symmetric=True, group_size=0) def test_midq_calculation(self): @@ -98,7 +98,7 @@ def test_get_num_groups_invalid_divisibility(self): quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32) shape = (64, 100) # 100 is not divisible by 32 - with pytest.raises(AssertionError, match=r"in_features .* must be divisible by group_size"): + with pytest.raises(AssertionError, match=r"last dim .* must be divisible by group_size"): quantizer.get_num_groups(shape) @pytest.mark.parametrize("group_size", [0, 16, -1]) @@ -399,3 +399,51 @@ def test_pack_unpack_device_consistency(self, device): unpacked = unpack_from_uint8(packed, bits=4, shape=tensor.shape) assert unpacked.device.type == device.type assert torch.all(unpacked == tensor.to(torch.int32)) + + +class TestNDimensional: + """N-D quantization tests. + + Verify that the quantizer / pack helpers operate identically on N-D tensors, + always quantizing along the last dim, without an explicit leading-dim loop. + """ + + @pytest.mark.parametrize("bits", [2, 4, 8]) + @pytest.mark.parametrize("group_size", [-1, 16, 32]) + def test_quantizer_3d_matches_2d_per_slice(self, bits, group_size): + """A 3D quantize must match independently quantizing each ``[i]`` slice.""" + torch.manual_seed(0) + weight = torch.randn(3, 8, 64) + quantizer = WeightQuantizer(bits=bits, symmetric=False, group_size=group_size) + + scales, zero_points = quantizer.find_qparams(weight) + q = quantizer.quantize(weight, scales, zero_points) + dq = quantizer.dequantize(q, scales, zero_points) + + # quantization parameters carry shape (num_experts, out, num_groups) + num_groups = 1 if group_size == -1 else 64 // group_size + assert scales.shape == (3, 8, num_groups) + assert zero_points.shape == (3, 8, num_groups) + assert q.shape == weight.shape + assert dq.shape == weight.shape + + for i in range(weight.shape[0]): + s_i, zp_i = quantizer.find_qparams(weight[i]) + assert torch.equal(scales[i], s_i) + assert torch.equal(zero_points[i], zp_i) + assert torch.equal(q[i], quantizer.quantize(weight[i], s_i, zp_i)) + + @pytest.mark.parametrize("bits", [2, 4, 8]) + def test_pack_unpack_3d_round_trip(self, bits): + """``pack_to_uint8`` / ``unpack_from_uint8`` round-trip on 3D inputs.""" + torch.manual_seed(0) + shape = (3, 8, 64) + tensor = torch.randint(0, 2**bits, shape, dtype=torch.uint8) + + packed = pack_to_uint8(tensor, bits) + packing_factor = 8 // bits + assert packed.shape == (3, 8, 64 // packing_factor) + + unpacked = unpack_from_uint8(packed, bits, shape) + assert unpacked.shape == shape + assert torch.all(unpacked == tensor.to(torch.int32)) diff --git a/test/common/test_hardlink_copy.py b/test/common/test_hardlink_copy.py index a0395e90a3..a15e26ebe1 100644 --- a/test/common/test_hardlink_copy.py +++ b/test/common/test_hardlink_copy.py @@ -1,6 +1,7 @@ import random import string import sys +from unittest.mock import patch import pytest @@ -67,6 +68,18 @@ def test_copy_file_to_file(create_dir, tmp_path): assert list(dst_dirpath.glob("**/*")) == [dst_filepath] +@patch("olive.common.utils.os.link", side_effect=NotImplementedError) +def test_copy_file_falls_back_when_hardlink_options_are_unsupported(mock_link, create_dir, tmp_path): + src_filepath = create_dir / "file1.ext1" + dst_filepath = tmp_path / "file1.ext1" + + hardlink_copy_file(src_filepath, dst_filepath) + + assert mock_link.call_count == 2 + assert dst_filepath.read_bytes() == src_filepath.read_bytes() + assert not dst_filepath.samefile(src_filepath) + + def test_copy_file_to_dir_overwrites(create_dir, tmp_path): # setup src_dirpath = create_dir diff --git a/test/common/test_hf_wrapper.py b/test/common/test_hf_wrapper.py index 667d645ce0..e25b7f7485 100644 --- a/test/common/test_hf_wrapper.py +++ b/test/common/test_hf_wrapper.py @@ -4,6 +4,7 @@ # -------------------------------------------------------------------------- import pytest from torch import nn +from transformers import PretrainedConfig from olive.common.hf.wrapper import ModelWrapper from test.utils import get_tiny_phi3, make_local_tiny_llama @@ -137,3 +138,169 @@ def test_hf_wrapper_lfm2(): # LFM2 must have both layer types assert has_attn_layer, "Expected at least one attention layer" assert has_conv_layer, "Expected at least one conv layer" + + +def test_hf_wrapper_composite_vl_config(): + """Composite VL configs (e.g. Qwen3.5/3.6-MoE VL) nest decoder attributes under ``text_config``. + + Verifies the ``text_config.*`` alias fallbacks in ``defaults.yaml`` resolve them, and the + ``qwen3_5_moe`` entries in ``LAYERS``/``EMBEDDINGS``/``PRE_HEAD_LAYERNORM``/ + ``ROTARY_EMBEDDING`` point at the VL decoder's actual nested module path + (``model.language_model.*``), without needing to download/load any real model or weights. + + ``text_config``/``vision_config`` must be actual (attribute-accessible) nested + ``PretrainedConfig`` objects here -- as real composite VL config classes (e.g. + ``Qwen3_5MoeConfig``) construct them -- not plain dicts, since ``resolve_alias``'s nested + path lookup uses ``getattr``, not dict indexing. + """ + text_config = PretrainedConfig() + text_config.hidden_size = 2048 + text_config.num_hidden_layers = 40 + text_config.num_attention_heads = 16 + text_config.num_key_value_heads = 2 + text_config.head_dim = 256 + text_config.num_experts = 256 + + vision_config = PretrainedConfig() + vision_config.hidden_size = 1152 + + composite_config = PretrainedConfig() + composite_config.model_type = "qwen3_5_moe" + composite_config.text_config = text_config + composite_config.vision_config = vision_config + + model_wrapper = ModelWrapper(composite_config) + + assert model_wrapper.model_type == "qwen3_5_moe" + assert model_wrapper.hidden_size == 2048 + assert model_wrapper.num_attention_heads == 16 + assert model_wrapper.num_key_value_heads == 2 + assert model_wrapper.head_dim == 256 + assert model_wrapper.num_hidden_layers == 40 + + assert model_wrapper.LAYERS[model_wrapper.model_type] == "model.language_model.layers" + assert model_wrapper.EMBEDDINGS[model_wrapper.model_type] == ["model.language_model.embed_tokens"] + assert model_wrapper.PRE_HEAD_LAYERNORM[model_wrapper.model_type] == "model.language_model.norm" + assert model_wrapper.ROTARY_EMBEDDING[model_wrapper.model_type] == "model.language_model.rotary_emb" + + +def test_hf_wrapper_text_only_config_unaffected_by_vl_aliases(): + """Verify the flat text-only Qwen3.5/3.6 config resolves exactly as before. + + ``model_type == "qwen3_5_moe_text"`` already has flat attributes matching the "default" + aliases, and has no ``qwen3_5_moe`` entry in the ``LAYERS``/etc. mappings (a different + model_type), so the VL-specific additions must not change its resolution. + """ + text_only_config = { + "model_type": "qwen3_5_moe_text", + "hidden_size": 2048, + "num_hidden_layers": 40, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "num_experts": 256, + } + + model_wrapper = ModelWrapper(text_only_config) + + assert model_wrapper.model_type == "qwen3_5_moe_text" + assert model_wrapper.hidden_size == 2048 + assert model_wrapper.num_hidden_layers == 40 + # No "qwen3_5_moe_text" entry in LAYERS/EMBEDDINGS/etc. -- falls back to "default". + assert model_wrapper.LAYERS.get(model_wrapper.model_type, model_wrapper.LAYERS["default"]) == "model.layers" + + +def test_hf_wrapper_flat_text_only_qwen3_5_moe_config(): + """Regression: flat text-only checkpoints also report ``model_type == "qwen3_5_moe"``. + + Only the composite VL checkpoint nests its decoder under ``model.language_model``; a + text-only ``Qwen3_5MoeForCausalLM`` checkpoint keeps the flat ``model.layers`` layout + while still reporting ``model_type == "qwen3_5_moe"``. The two are disambiguated by the + presence of a ``vision_config`` (same rule mobius uses), so the flat config must resolve + the "default" module paths instead of crashing on the VL-only paths. + """ + + class _Layer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.mlp = nn.Module() + + class _FlatQwenMoE(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.model = nn.Module() + self.model.embed_tokens = nn.Embedding(16, 8) + self.model.layers = nn.ModuleList([_Layer(), _Layer()]) + self.model.norm = nn.LayerNorm(8) + self.lm_head = nn.Linear(8, 16, bias=False) + + flat_config = PretrainedConfig() + flat_config.model_type = "qwen3_5_moe" + flat_config.hidden_size = 2048 + flat_config.num_hidden_layers = 2 + flat_config.num_attention_heads = 16 + flat_config.num_key_value_heads = 2 + flat_config.head_dim = 256 + flat_config.num_experts = 256 + + model_wrapper = ModelWrapper(flat_config) + + # normalized to the text-only model_type, so the flat "default" paths are used + assert model_wrapper.model_type == "qwen3_5_moe_text" + assert model_wrapper.hidden_size == 2048 + assert model_wrapper.head_dim == 256 + + model = _FlatQwenMoE(flat_config) + model_wrapper.set_model(model) + + assert model_wrapper.get_layers(False) is model.model.layers + assert model_wrapper.get_embeds(False)[0] is model.model.embed_tokens + assert model_wrapper.get_pre_head_layernorm(False) is model.model.norm + assert model_wrapper.get_lm_head(False) is model.lm_head + + +def test_hf_wrapper_vl_qwen3_5_moe_config_uses_nested_paths(): + """The composite VL config (has ``vision_config``) keeps the nested decoder paths.""" + + class _Layer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.mlp = nn.Module() + + class _VLQwenMoE(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.embed_tokens = nn.Embedding(16, 8) + self.model.language_model.layers = nn.ModuleList([_Layer()]) + self.model.language_model.norm = nn.LayerNorm(8) + self.model.visual = nn.Module() + self.model.visual.proj = nn.Linear(8, 8, bias=False) + self.lm_head = nn.Linear(8, 16, bias=False) + + text_config = PretrainedConfig() + text_config.hidden_size = 2048 + text_config.num_hidden_layers = 1 + text_config.num_attention_heads = 16 + text_config.num_key_value_heads = 2 + text_config.head_dim = 256 + + vl_config = PretrainedConfig() + vl_config.model_type = "qwen3_5_moe" + vl_config.text_config = text_config + vl_config.vision_config = PretrainedConfig() + + model_wrapper = ModelWrapper(vl_config) + assert model_wrapper.model_type == "qwen3_5_moe" + + model = _VLQwenMoE(vl_config) + model_wrapper.set_model(model) + + assert model_wrapper.get_layers(False) is model.model.language_model.layers + assert model_wrapper.get_embeds(False)[0] is model.model.language_model.embed_tokens + assert model_wrapper.get_pre_head_layernorm(False) is model.model.language_model.norm diff --git a/test/common/test_mobius_utils.py b/test/common/test_mobius_utils.py index e5dd7951b2..872832be92 100644 --- a/test/common/test_mobius_utils.py +++ b/test/common/test_mobius_utils.py @@ -71,5 +71,5 @@ def test_inspect_components_coerces_mobius_objects(monkeypatch): def test_inspect_components_raises_importerror_when_mobius_missing(monkeypatch): monkeypatch.setitem(sys.modules, "mobius", None) - with pytest.raises(ImportError, match="mobius-ai is required"): + with pytest.raises(ImportError, match="mobius-onnx is required"): inspect_components("fake/llava") diff --git a/test/common/test_onnx_io.py b/test/common/test_onnx_io.py new file mode 100644 index 0000000000..72ba468e91 --- /dev/null +++ b/test/common/test_onnx_io.py @@ -0,0 +1,30 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import json + +from olive.common.onnx_io import get_genai_decoder_config + +DECODER_CONFIG = {"head_size": 64, "num_key_value_heads": 2} + + +def _save_genai_config(model_dir): + model_dir.mkdir(parents=True, exist_ok=True) + (model_dir / "genai_config.json").write_text(json.dumps({"model": {"decoder": DECODER_CONFIG}})) + + +def test_get_genai_decoder_config_returns_decoder_section_for_model_file(tmp_path): + _save_genai_config(tmp_path) + + assert get_genai_decoder_config(tmp_path / "model.onnx") == DECODER_CONFIG + + +def test_get_genai_decoder_config_returns_decoder_section_for_model_dir(tmp_path): + _save_genai_config(tmp_path) + + assert get_genai_decoder_config(tmp_path) == DECODER_CONFIG + + +def test_get_genai_decoder_config_returns_none_when_config_is_missing(tmp_path): + assert get_genai_decoder_config(tmp_path / "model.onnx") is None diff --git a/test/conftest.py b/test/conftest.py index db97c685af..fb70f9b4c5 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -2,13 +2,14 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +import os import shutil from unittest.mock import patch import pytest from packaging import version -from olive.telemetry.telemetry import Telemetry +import olive.telemetry.telemetry as telemetry_module from test.utils import create_onnx_model_file, delete_onnx_model_files @@ -44,4 +45,11 @@ def maybe_patch_inc(): @pytest.fixture(scope="session", autouse=True) def disable_telemetry(): - Telemetry().disable_telemetry() + # Apply the full opt-out before singleton construction so tests create no + # telemetry identity, store, uploader, heartbeat thread, or network traffic. + with patch.dict(os.environ, {"ORT_DISABLE_TELEMETRY": "1"}): + telemetry = telemetry_module.Telemetry() + try: + yield + finally: + telemetry.shutdown() diff --git a/test/data_container/test_text_generation.py b/test/data_container/test_text_generation.py new file mode 100644 index 0000000000..810ca0633f --- /dev/null +++ b/test/data_container/test_text_generation.py @@ -0,0 +1,315 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import logging + +import pytest +import torch +from datasets import Dataset + +from olive.data.component.text_generation import TextGenStrategy, text_gen_pre_process + + +@pytest.fixture(name="propagate_olive_logs", autouse=True) +def propagate_olive_logs_fixture(): + """Olive disables propagation on its root logger, so caplog cannot see the records without this. + + This mirrors the inline propagation toggling already used elsewhere in the suite, e.g. in + test/engine/test_engine.py, test/common/test_copy_dir.py and test/hardware/test_accelerator.py. + """ + logger = logging.getLogger("olive") + original = logger.propagate + logger.propagate = True + yield + logger.propagate = original + + +class FakeEncoding(dict): + """Dict that also supports attribute access, like transformers' BatchEncoding.""" + + def __getattr__(self, name): + try: + return self[name] + except KeyError as exc: + raise AttributeError(name) from exc + + +class FakeTokenizer: + """Minimal word level tokenizer so the tests don't need to download a real tokenizer.""" + + def __init__(self, pad_token_id: int = 0): + self.padding_side = "right" + self.pad_token_id = pad_token_id + + def encode(self, text, add_special_tokens=False): + # only the number of tokens matters for these tests, the token ids themselves are arbitrary + return [len(word) for word in text.split()] + + def __call__( + self, + text, + add_special_tokens=False, + truncation=False, + max_length=None, + padding=False, + return_tensors=None, + **kwargs, + ): + texts = [text] if isinstance(text, str) else text + input_ids = [self.encode(single_text) for single_text in texts] + if truncation and max_length is not None: + input_ids = [ids[:max_length] for ids in input_ids] + attention_mask = [[1] * len(ids) for ids in input_ids] + if padding == "max_length" and max_length is not None: + attention_mask = [mask + [0] * (max_length - len(mask)) for mask in attention_mask] + input_ids = [ids + [self.pad_token_id] * (max_length - len(ids)) for ids in input_ids] + if return_tensors == "pt": + return FakeEncoding(input_ids=torch.tensor(input_ids), attention_mask=torch.tensor(attention_mask)) + if isinstance(text, str): + return FakeEncoding(input_ids=input_ids[0], attention_mask=attention_mask[0]) + return FakeEncoding(input_ids=input_ids, attention_mask=attention_mask) + + +def make_dataset(num_rows: int, words_per_row: int) -> Dataset: + return Dataset.from_dict({"text": [" ".join(["word"] * words_per_row) for _ in range(num_rows)]}) + + +def get_kwargs(max_samples, max_seq_len: int = 8, **kwargs) -> dict: + return { + "strategy": TextGenStrategy.JOIN, + "add_special_tokens": False, + "max_samples": max_samples, + "max_seq_len": max_seq_len, + "joiner": "", + **kwargs, + } + + +@pytest.mark.parametrize( + ("strategy", "expected_step", "expected_required", "expected_samples"), + [ + # JOIN: step is max_seq_len, so 8 + (10 - 1) * 8 = 80 tokens are needed, 16 tokens give 2 samples + (TextGenStrategy.JOIN, 8, 80, 2), + # JOIN_SLIDING_WINDOW: step is stride, so 8 + (10 - 1) * 4 = 44 tokens are needed, 16 tokens give 3 samples + (TextGenStrategy.JOIN_SLIDING_WINDOW, 4, 44, 3), + ], +) +def test_text_gen_pre_process_warns_when_join_corpus_exhausted_before_max_samples( + strategy, expected_step, expected_required, expected_samples, caplog +): + # setup + # 4 rows x 4 tokens = 16 tokens + dataset = make_dataset(num_rows=4, words_per_row=4) + kwargs = get_kwargs(max_samples=10, strategy=strategy, stride=4) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == expected_samples + assert len(caplog.records) == 1 + message = caplog.records[0].getMessage() + assert f"Only {expected_samples} samples were generated" in message + assert "max_samples=10" in message + assert f"step={expected_step}" in message + assert f"= {expected_required} tokens" in message + + +def test_text_gen_pre_process_warns_with_stride_remedy_when_strategy_is_sliding_window(caplog): + # setup + dataset = make_dataset(num_rows=4, words_per_row=4) + kwargs = get_kwargs(max_samples=10, strategy=TextGenStrategy.JOIN_SLIDING_WINDOW, stride=4) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert "lower stride" in caplog.records[0].getMessage() + + +def test_text_gen_pre_process_delivers_max_samples_when_corpus_is_an_exact_fit(caplog): + # setup + # 4 rows x 8 tokens = 32 tokens, exactly max_samples x max_seq_len = 4 x 8 + dataset = make_dataset(num_rows=4, words_per_row=8) + kwargs = get_kwargs(max_samples=4) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 4 + assert all(len(sample[0]["input_ids"]) == 8 for sample in result) + assert not caplog.records + + +def test_text_gen_pre_process_warns_when_random_strategy_cannot_fill_max_samples(caplog): + # setup + # each row is shorter than max_seq_len and there is only one row, so no sample can be built + dataset = make_dataset(num_rows=1, words_per_row=2) + kwargs = get_kwargs(max_samples=3, strategy=TextGenStrategy.JOIN_RANDOM, random_seed=0, random_retries=5) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 0 + assert len(caplog.records) == 1 + message = caplog.records[0].getMessage() + assert "max_samples=3" in message + assert "random_retries=5" in message + # max_samples must not be presented as part of the token requirement for the random strategy + assert "max_samples does not increase that requirement" in message + assert "lower max_samples" not in message + + +@pytest.mark.parametrize( + "strategy", + [TextGenStrategy.JOIN, TextGenStrategy.JOIN_SLIDING_WINDOW], +) +def test_text_gen_pre_process_does_not_warn_when_max_samples_delivered(strategy, caplog): + # setup + # 20 rows x 8 tokens = 160 tokens, more than enough for 4 samples with either step + dataset = make_dataset(num_rows=20, words_per_row=8) + kwargs = get_kwargs(max_samples=4, strategy=strategy, stride=4) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 4 + assert not caplog.records + + +def test_text_gen_pre_process_does_not_warn_when_max_samples_is_none(caplog): + # setup + dataset = make_dataset(num_rows=4, words_per_row=4) + kwargs = get_kwargs(max_samples=None) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert not caplog.records + + +def test_text_gen_pre_process_warns_when_line_by_line_has_too_few_rows(caplog): + # setup + # only 3 usable rows for 5 requested samples + dataset = make_dataset(num_rows=3, words_per_row=8) + kwargs = get_kwargs(max_samples=5, strategy=TextGenStrategy.LINE_BY_LINE) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 3 + assert len(caplog.records) == 1 + message = caplog.records[0].getMessage() + assert "Only 3 samples were generated" in message + assert "3 non-empty rows" in message + assert "lower max_samples" in message + + +def test_text_gen_pre_process_warns_when_line_by_line_drops_short_rows(caplog): + # setup + # every row is shorter than max_seq_len and drop_short_sequences filters them all out + dataset = make_dataset(num_rows=5, words_per_row=2) + kwargs = get_kwargs( + max_samples=5, + strategy=TextGenStrategy.LINE_BY_LINE, + pad_to_max_len=False, + drop_short_sequences=True, + ) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 0 + assert "drop_short_sequences=True" in caplog.records[0].getMessage() + + +def test_text_gen_pre_process_does_not_warn_when_line_by_line_delivers_max_samples(caplog): + # setup + dataset = make_dataset(num_rows=5, words_per_row=8) + kwargs = get_kwargs(max_samples=5, strategy=TextGenStrategy.LINE_BY_LINE) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 5 + assert not caplog.records + + +def test_text_gen_pre_process_warns_when_line_by_line_random_finds_no_sample(caplog): + # setup + # random_retries=0 means no row is ever accepted, so every sample is skipped + dataset = make_dataset(num_rows=5, words_per_row=2) + kwargs = get_kwargs( + max_samples=4, + strategy=TextGenStrategy.LINE_BY_LINE_RANDOM, + random_seed=0, + random_retries=0, + pad_to_max_len=False, + drop_short_sequences=True, + ) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 0 + assert len(caplog.records) == 1 + message = caplog.records[0].getMessage() + assert "max_samples=4" in message + assert "random_retries=0" in message + + +def test_text_gen_pre_process_drops_short_rows_when_line_by_line_random_exhausts_retries(caplog): + # setup + # every row is shorter than max_seq_len, so all random_retries attempts are rejected. The loop variable + # still holds the last rejected row, which used to be emitted because a BatchEncoding is always truthy. + dataset = make_dataset(num_rows=5, words_per_row=2) + kwargs = get_kwargs( + max_samples=4, + strategy=TextGenStrategy.LINE_BY_LINE_RANDOM, + random_seed=0, + random_retries=3, + pad_to_max_len=False, + drop_short_sequences=True, + ) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 0 + assert len(caplog.records) == 1 + assert "drop_short_sequences=True" in caplog.records[0].getMessage() + + +def test_text_gen_pre_process_does_not_warn_when_line_by_line_random_delivers_max_samples(caplog): + # setup + dataset = make_dataset(num_rows=5, words_per_row=8) + kwargs = get_kwargs(max_samples=4, strategy=TextGenStrategy.LINE_BY_LINE_RANDOM, random_seed=0) + + # execute + with caplog.at_level(logging.WARNING, logger="olive.data.component.text_generation"): + result = text_gen_pre_process(dataset, FakeTokenizer(), kwargs) + + # assert + assert len(result) == 4 + assert not caplog.records diff --git a/test/model/test_hf_model.py b/test/model/test_hf_model.py index 019ddefe2d..087d62d3ad 100644 --- a/test/model/test_hf_model.py +++ b/test/model/test_hf_model.py @@ -3,6 +3,8 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import json +import logging +from contextlib import contextmanager from pathlib import Path from unittest.mock import ANY, patch @@ -100,6 +102,161 @@ def test_save_pretrained_metadata(self, local, tmp_path): } +def test_save_metadata_saves_processor_for_vl_model(tmp_path): + """save_metadata should save preprocessor_config.json for multimodal models. + + Uses a mocked AutoProcessor since there's no lightweight real VL checkpoint fixture + available; verifies the processor.save_pretrained output is captured and the + tokenizer-like AutoProcessor return value (text-only models) is correctly skipped. + """ + olive_model = HfModelHandler( + model_path="katuni4ka/tiny-random-phi3", + task="text-generation-with-past", + load_kwargs={"revision": "585361abfee667f3c63f8b2dc4ad58405c4e34e2"}, + ) + + processor_config_path = tmp_path / "preprocessor_config.json" + + class FakeProcessor: + def save_pretrained(self, output_dir, **kwargs): + path = Path(output_dir) / "preprocessor_config.json" + path.write_text("{}") + return [str(path)] + + with patch("transformers.AutoProcessor.from_pretrained", return_value=FakeProcessor()): + saved_filepaths = olive_model.save_metadata(tmp_path) + + assert str(processor_config_path) in saved_filepaths + assert processor_config_path.exists() + + +def test_save_metadata_processor_fills_in_missing_files_without_overwriting_existing(tmp_path): + """A processor with some files already saved must still be filled in without clobbering existing ones. + + save_metadata should still look for processor files even if one (e.g. preprocessor_config.json) + already exists, since a processor can emit several files and an earlier step may have saved only + some of them -- but it must never overwrite a file that's already there. + """ + olive_model = HfModelHandler( + model_path="katuni4ka/tiny-random-phi3", + task="text-generation-with-past", + load_kwargs={"revision": "585361abfee667f3c63f8b2dc4ad58405c4e34e2"}, + ) + + tmp_path.mkdir(parents=True, exist_ok=True) + (tmp_path / "preprocessor_config.json").write_text('{"existing": true}') + + class FakeProcessor: + def save_pretrained(self, output_dir, **kwargs): + out = Path(output_dir) + (out / "preprocessor_config.json").write_text('{"existing": false}') + (out / "chat_template.json").write_text("{}") + return [str(out / "preprocessor_config.json"), str(out / "chat_template.json")] + + with patch("transformers.AutoProcessor.from_pretrained", return_value=FakeProcessor()) as mock_from_pretrained: + saved_filepaths = olive_model.save_metadata(tmp_path) + + mock_from_pretrained.assert_called_once() + # the pre-existing preprocessor_config.json is untouched, but the missing chat_template.json is filled in + assert (tmp_path / "preprocessor_config.json").read_text() == '{"existing": true}' + assert (tmp_path / "chat_template.json").exists() + assert str(tmp_path / "chat_template.json") in saved_filepaths + assert str(tmp_path / "preprocessor_config.json") not in saved_filepaths + + +@contextmanager +def capture_warnings(logger_name: str): + """Collect warning-or-worse records from ``logger_name`` (Olive loggers don't propagate).""" + records: list[str] = [] + + class _Handler(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + logger = logging.getLogger(logger_name) + handler = _Handler(level=logging.WARNING) + previous_level = logger.level + logger.addHandler(handler) + logger.setLevel(logging.WARNING) + try: + yield records + finally: + logger.removeHandler(handler) + logger.setLevel(previous_level) + + +def test_save_metadata_warns_on_unexpected_processor_failure(tmp_path): + """An unexpected AutoProcessor failure must be visible, not silently swallowed at debug level.""" + olive_model = HfModelHandler( + model_path="katuni4ka/tiny-random-phi3", + task="text-generation-with-past", + load_kwargs={"revision": "585361abfee667f3c63f8b2dc4ad58405c4e34e2"}, + ) + + with ( + patch("transformers.AutoProcessor.from_pretrained", side_effect=RuntimeError("boom")), + capture_warnings("olive.model.handler.mixin.hf") as warnings, + ): + saved_filepaths = olive_model.save_metadata(tmp_path) + + assert any("boom" in message for message in warnings) + assert not (tmp_path / "preprocessor_config.json").exists() + # the rest of the metadata is still saved + assert str(tmp_path / "config.json") in saved_filepaths + + +def test_save_metadata_skips_tokenizer_return_silently(tmp_path): + """Text-only models: AutoProcessor returns the tokenizer -- expected, so no warning.""" + olive_model = HfModelHandler( + model_path="katuni4ka/tiny-random-phi3", + task="text-generation-with-past", + load_kwargs={"revision": "585361abfee667f3c63f8b2dc4ad58405c4e34e2"}, + ) + + tokenizer = transformers.AutoTokenizer.from_pretrained( + "katuni4ka/tiny-random-phi3", revision="585361abfee667f3c63f8b2dc4ad58405c4e34e2" + ) + with ( + patch("transformers.AutoProcessor.from_pretrained", return_value=tokenizer), + capture_warnings("olive.model.handler.mixin.hf") as warnings, + ): + olive_model.save_metadata(tmp_path) + + assert not warnings + assert not (tmp_path / "preprocessor_config.json").exists() + + +def test_save_metadata_processor_does_not_overwrite_custom_tokenizer(tmp_path): + """A processor save must not clobber a tokenizer an earlier step customized and saved.""" + olive_model = HfModelHandler( + model_path="katuni4ka/tiny-random-phi3", + task="text-generation-with-past", + load_kwargs={"revision": "585361abfee667f3c63f8b2dc4ad58405c4e34e2"}, + ) + + custom_tokenizer_config = '{"custom": true}' + (tmp_path / "tokenizer_config.json").write_text(custom_tokenizer_config) + + class FakeProcessor: + """Mimics ProcessorMixin.save_pretrained, which also re-saves the tokenizer.""" + + def save_pretrained(self, output_dir, **kwargs): + out = Path(output_dir) + (out / "preprocessor_config.json").write_text('{"image_processor": true}') + (out / "tokenizer_config.json").write_text('{"custom": false}') + (out / "tokenizer.json").write_text("{}") + return [str(out / "preprocessor_config.json")] + + with patch("transformers.AutoProcessor.from_pretrained", return_value=FakeProcessor()): + saved_filepaths = olive_model.save_metadata(tmp_path) + + # the pre-existing tokenizer config is untouched, the processor config is new + assert (tmp_path / "tokenizer_config.json").read_text() == custom_tokenizer_config + assert (tmp_path / "preprocessor_config.json").read_text() == '{"image_processor": true}' + assert str(tmp_path / "preprocessor_config.json") in saved_filepaths + assert str(tmp_path / "tokenizer_config.json") not in saved_filepaths + + @pytest.mark.parametrize("trust_remote_code", [True, False]) def test_save_metadata_with_module_files(trust_remote_code, tmp_path): load_kwargs = {"trust_remote_code": trust_remote_code, "revision": "585361abfee667f3c63f8b2dc4ad58405c4e34e2"} @@ -110,12 +267,13 @@ def test_save_metadata_with_module_files(trust_remote_code, tmp_path): saved_filepaths = olive_model.save_metadata(tmp_path) assert all(Path(fp).exists() for fp in saved_filepaths) + config = transformers.AutoConfig.from_pretrained(tmp_path, **load_kwargs) + assert config.__class__.__name__ == "Phi3Config" if trust_remote_code: - expected_class_name = f"transformers_modules.{tmp_path.name}.configuration_phi3.Phi3Config" + assert config.__module__.startswith(f"transformers_modules.{tmp_path.name}.") + assert config.__module__.endswith(".configuration_phi3") else: - expected_class_name = "transformers.models.phi3.configuration_phi3.Phi3Config" - config = transformers.AutoConfig.from_pretrained(tmp_path, **load_kwargs) - assert f"{config.__module__}.{config.__class__.__name__}" == expected_class_name + assert config.__module__ == "transformers.models.phi3.configuration_phi3" assert isinstance( transformers.AutoTokenizer.from_pretrained(tmp_path, **load_kwargs), transformers.PreTrainedTokenizerBase, diff --git a/test/passes/onnx/test_common.py b/test/passes/onnx/test_common.py index 82f1e3f4cd..9498cd9c39 100644 --- a/test/passes/onnx/test_common.py +++ b/test/passes/onnx/test_common.py @@ -40,9 +40,13 @@ def test_model_proto_to_olive_model(external_data_config, tmp_path): @pytest.mark.parametrize("has_external_data", [True, False]) def test_resave_model(has_external_data, tmp_path): # setup + from transformers.cache_utils import DynamicLayer + + original_lazy_initialization = DynamicLayer.lazy_initialization input_model = create_pass_from_dict( OnnxConversion, {"save_as_external_data": has_external_data, "use_dynamo_exporter": True}, disable_search=True ).run(get_hf_model(), str(tmp_path / "input")) + assert DynamicLayer.lazy_initialization is original_lazy_initialization # execute resave_path = tmp_path / "resave" / "resave.onnx" diff --git a/test/passes/onnx/test_conversion.py b/test/passes/onnx/test_conversion.py index 0b1a48680c..001c251c27 100644 --- a/test/passes/onnx/test_conversion.py +++ b/test/passes/onnx/test_conversion.py @@ -17,7 +17,11 @@ from olive.model import PyTorchModelHandler from olive.model.config import IoConfig from olive.passes.olive_pass import create_pass_from_dict -from olive.passes.onnx.conversion import OnnxConversion, OnnxOpVersionConversion +from olive.passes.onnx.conversion import ( + OnnxConversion, + OnnxOpVersionConversion, + _patch_dynamic_layer_for_export, +) from olive.passes.pytorch.autogptq import GptqQuantizer from olive.passes.pytorch.rtn import Rtn from test.utils import ( @@ -35,6 +39,82 @@ def _torch_is_older_than(version_str: str) -> bool: return torch_version < version.parse(version_str).release +def test_dynamic_layer_export_patch_preserves_key_and_value_shapes(): + from transformers.cache_utils import DynamicLayer + + original_lazy_initialization = DynamicLayer.lazy_initialization + with _patch_dynamic_layer_for_export(): + layer = DynamicLayer() + keys, values = layer.update(torch.ones(1, 2, 1, 16), torch.ones(1, 2, 1, 8)) + assert keys.shape == (1, 2, 1, 16) + assert values.shape == (1, 2, 1, 8) + + assert DynamicLayer.lazy_initialization is original_lazy_initialization + + +def test_dynamic_layer_export_patch_restores_method_on_error(): + from transformers.cache_utils import DynamicLayer + + original_lazy_initialization = DynamicLayer.lazy_initialization + # Use a plain try/except (rather than `pytest.raises`, even nested/separated) so CodeQL's + # control-flow analysis can see that the code below is reachable: CodeQL doesn't model + # `pytest.raises.__exit__` as suppressing the exception, so it treats any statement after a + # `with` block whose body unconditionally raises as unreachable, regardless of an outer + # `pytest.raises`. A native try/except is understood correctly. + raised = None + try: + with _patch_dynamic_layer_for_export(): + raise RuntimeError("export failed") # noqa: TRY301 + except RuntimeError as error: + raised = error + assert raised is not None + assert str(raised) == "export failed" + assert DynamicLayer.lazy_initialization is original_lazy_initialization + + +def test_dynamic_layer_export_patch_nested_usage_restores_after_outer_exit(): + from transformers.cache_utils import DynamicLayer + + original_lazy_initialization = DynamicLayer.lazy_initialization + with _patch_dynamic_layer_for_export(): + patched_lazy_initialization = DynamicLayer.lazy_initialization + assert patched_lazy_initialization is not original_lazy_initialization + with _patch_dynamic_layer_for_export(): + assert DynamicLayer.lazy_initialization is patched_lazy_initialization + assert DynamicLayer.lazy_initialization is patched_lazy_initialization + + assert DynamicLayer.lazy_initialization is original_lazy_initialization + + +def test_dynamic_layer_export_patch_non_lifo_overlapping_usage_restores_after_last_exit(): + from transformers.cache_utils import DynamicLayer + + original_lazy_initialization = DynamicLayer.lazy_initialization + context_a = _patch_dynamic_layer_for_export() + context_b = _patch_dynamic_layer_for_export() + context_a_active = context_b_active = False + try: + context_a.__enter__() + context_a_active = True + patched_lazy_initialization = DynamicLayer.lazy_initialization + context_b.__enter__() + context_b_active = True + assert DynamicLayer.lazy_initialization is patched_lazy_initialization + + context_a.__exit__(None, None, None) + context_a_active = False + assert DynamicLayer.lazy_initialization is patched_lazy_initialization + + context_b.__exit__(None, None, None) + context_b_active = False + assert DynamicLayer.lazy_initialization is original_lazy_initialization + finally: + if context_b_active: + context_b.__exit__(None, None, None) + if context_a_active: + context_a.__exit__(None, None, None) + + @pytest.mark.parametrize( ("input_model", "use_dynamo_exporter", "dynamic"), [ diff --git a/test/passes/onnx/test_discrepancy_check.py b/test/passes/onnx/test_discrepancy_check.py index 404bdb3b49..e236a44d21 100644 --- a/test/passes/onnx/test_discrepancy_check.py +++ b/test/passes/onnx/test_discrepancy_check.py @@ -11,8 +11,11 @@ from olive.passes.onnx.discrepancy_check import ( _expand_genai_output_names, + _has_bfloat16, + _infer_shape, _longest_common_token_sequence, _reconcile_genai_speech_output_names, + _run_onnx_session, ) @@ -51,6 +54,19 @@ def test_common_tokens_later_not_counted(self): assert _longest_common_token_sequence([10, 1, 2, 3], [20, 1, 2, 3]) == 0 +def test_infer_shape_resolves_kv_cache_dim_from_known_values(): + inferred = _infer_shape( + ["batch_size", 8, "past_sequence_length", "kv_cache_dim"], + {"kv_cache_dim": 16}, + ) + assert inferred == (1, 8, 0, 16) + + +def test_infer_shape_error_message_handles_mixed_known_symbol_keys(): + with pytest.raises(KeyError, match="Unsupported symbolic dimension 'mystery_dim'"): + _infer_shape(["batch_size", "mystery_dim"], {"kv_cache_dim": 16, 8: 8}) + + def _whisper_genai_config(num_layers=2): """Build a minimal Whisper-style genai_config with cross-attention cache outputs.""" return { @@ -1425,8 +1441,9 @@ def test_run_speech_generation_comparison_surfaces_transformers_on_genai_failure ): results = pass_instance._run_speech_generation_comparison(model, config, MagicMock(), "ref_path") - # Status is not skipped: transformers figures are surfaced despite the GenAI failure. - assert results["status"] == "passed" + # Status is "partial": transformers figures are surfaced despite the GenAI failure, + # but GenAI was unavailable so the comparison is incomplete. + assert results["status"] == "partial" assert results["transformers_first_token"] == 50258 assert results["transformers_ttft_s"] == 0.01 assert "present_key_cross_2" in results["genai_generation_error"] @@ -1804,3 +1821,318 @@ 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 TestRunComponentSubprocess: + """Tests for _run_component_subprocess error handling and result parsing.""" + + def test_returns_result_on_success(self, tmp_path): + """Subprocess writes a valid result JSON → returned as-is.""" + from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck + + script = tmp_path / "worker.py" + script.write_text( + "import json, sys\n" + 'result = {"max_abs_error": 0.001, "elements_above_0_1": 0, ' + '"elements_above_0_01": 2, "total_elements": 100, "output_compared": "logits"}\n' + 'with open(sys.argv[2], "w") as f: json.dump(result, f)\n' + ) + result = OnnxDiscrepancyCheck._run_component_subprocess(script, tmp_path, {"component": "decoder"}) + assert result["max_abs_error"] == 0.001 + assert "error" not in result + + def test_returns_error_on_crash(self, tmp_path): + """Subprocess exits non-zero → error dict with stderr.""" + from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck + + script = tmp_path / "worker.py" + script.write_text('import sys\nsys.stderr.write("RuntimeError: boom\\n")\nsys.exit(1)\n') + result = OnnxDiscrepancyCheck._run_component_subprocess(script, tmp_path, {"component": "encoder"}) + assert "error" in result + assert "boom" in result["error"] + + def test_returns_error_when_no_output(self, tmp_path): + """Subprocess exits 0 but writes no result file → error dict.""" + from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck + + script = tmp_path / "worker.py" + script.write_text("pass\n") + result = OnnxDiscrepancyCheck._run_component_subprocess(script, tmp_path, {"component": "encoder"}) + assert "error" in result + + def test_strips_extension_modules_from_faulthandler_stderr(self, tmp_path): + """Faulthandler output with Extension modules list is truncated.""" + from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck + + script = tmp_path / "worker.py" + # Simulate faulthandler-style output with Extension modules list + script.write_text( + "import sys\n" + 'sys.stderr.write("Fatal Python error: Segmentation fault\\n")\n' + 'sys.stderr.write(" File \\"foo.py\\", line 10 in run\\n")\n' + 'sys.stderr.write("Extension modules: numpy, torch, ort\\n")\n' + "sys.exit(-11)\n" + ) + result = OnnxDiscrepancyCheck._run_component_subprocess(script, tmp_path, {"component": "encoder"}) + assert "error" in result + assert "Fatal Python error" in result["error"] + assert "Extension modules" not in result["error"] + + def test_returns_error_on_invalid_json_output(self, tmp_path): + """Subprocess exits 0 but writes invalid JSON → error dict.""" + from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck + + script = tmp_path / "worker.py" + script.write_text('import sys\nwith open(sys.argv[2], "w") as f: f.write("{invalid")\n') + result = OnnxDiscrepancyCheck._run_component_subprocess(script, tmp_path, {"component": "decoder"}) + assert "error" in result + assert "invalid JSON" in result["error"] + + +class TestComputeSpeechComponentDiscrepancy: + """Tests for _compute_speech_component_discrepancy orchestration.""" + + def _make_pass_instance(self): + from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck + + return OnnxDiscrepancyCheck.__new__(OnnxDiscrepancyCheck) + + def test_skips_non_composite_model(self): + """Non-CompositeModelHandler → empty dict.""" + instance = self._make_pass_instance() + result = instance._compute_speech_component_discrepancy(MagicMock(), "/some/path") + assert result == {} + + def test_skips_when_ref_path_not_directory(self, tmp_path): + """ref_model_path is not an existing directory → empty dict.""" + from olive.model import CompositeModelHandler + + model = CompositeModelHandler.__new__(CompositeModelHandler) + enc = MagicMock() + dec = MagicMock() + model.get_model_components = MagicMock(return_value=[("encoder.onnx", enc), ("decoder.onnx", dec)]) + + instance = self._make_pass_instance() + result = instance._compute_speech_component_discrepancy(model, str(tmp_path / "nonexistent")) + assert result == {} + + def test_aggregates_results_from_both_components(self, tmp_path): + """Both subprocesses succeed → aggregate max_abs_error is the max of both.""" + from olive.model import CompositeModelHandler + + model = CompositeModelHandler.__new__(CompositeModelHandler) + enc = MagicMock(model_path=str(tmp_path / "encoder.onnx")) + dec = MagicMock(model_path=str(tmp_path / "decoder.onnx")) + model.get_model_components = MagicMock(return_value=[("encoder.onnx", enc), ("decoder.onnx", dec)]) + + ref_path = tmp_path / "ref_model" + ref_path.mkdir() + + enc_result = { + "max_abs_error": 0.05, + "elements_above_0_1": 0, + "elements_above_0_01": 10, + "total_elements": 1000, + "output_compared": "hidden_states", + } + dec_result = { + "max_abs_error": 0.002, + "elements_above_0_1": 0, + "elements_above_0_01": 3, + "total_elements": 500, + "output_compared": "logits", + } + + instance = self._make_pass_instance() + with patch.object(type(instance), "_run_component_subprocess", side_effect=[enc_result, dec_result]): + result = instance._compute_speech_component_discrepancy(model, str(ref_path)) + + assert result["max_abs_error"] == 0.05 + assert result["elements_above_0_01"] == 13 + assert result["total_elements"] == 1500 + assert result["components"]["encoder"] == enc_result + assert result["components"]["decoder"] == dec_result + + def test_encoder_failure_does_not_block_decoder(self, tmp_path): + """Encoder returns error → decoder still runs, aggregate uses decoder only.""" + from olive.model import CompositeModelHandler + + model = CompositeModelHandler.__new__(CompositeModelHandler) + enc = MagicMock(model_path=str(tmp_path / "encoder.onnx")) + dec = MagicMock(model_path=str(tmp_path / "decoder.onnx")) + model.get_model_components = MagicMock(return_value=[("encoder.onnx", enc), ("decoder.onnx", dec)]) + + ref_path = tmp_path / "ref_model" + ref_path.mkdir() + + enc_result = {"error": "segfault"} + dec_result = { + "max_abs_error": 0.001, + "elements_above_0_1": 0, + "elements_above_0_01": 0, + "total_elements": 200, + "output_compared": "logits", + } + + instance = self._make_pass_instance() + with patch.object(type(instance), "_run_component_subprocess", side_effect=[enc_result, dec_result]): + result = instance._compute_speech_component_discrepancy(model, str(ref_path)) + + assert result["max_abs_error"] == 0.001 + assert "error" in result["components"]["encoder"] + assert "error" not in result["components"]["decoder"] + + def test_both_fail_returns_components_but_no_aggregate(self, tmp_path): + """Both components error → no aggregate max_abs_error key.""" + from olive.model import CompositeModelHandler + + model = CompositeModelHandler.__new__(CompositeModelHandler) + enc = MagicMock(model_path=str(tmp_path / "encoder.onnx")) + dec = MagicMock(model_path=str(tmp_path / "decoder.onnx")) + model.get_model_components = MagicMock(return_value=[("encoder.onnx", enc), ("decoder.onnx", dec)]) + + ref_path = tmp_path / "ref_model" + ref_path.mkdir() + + instance = self._make_pass_instance() + with patch.object( + type(instance), "_run_component_subprocess", side_effect=[{"error": "enc crash"}, {"error": "dec crash"}] + ): + result = instance._compute_speech_component_discrepancy(model, str(ref_path)) + + assert "max_abs_error" not in result + assert "error" in result["components"]["encoder"] + assert "error" in result["components"]["decoder"] + + def test_component_name_matching_with_extensions(self, tmp_path): + """Component names like 'encoder.onnx' and 'decoder.onnx' are matched correctly.""" + from olive.model import CompositeModelHandler + + model = CompositeModelHandler.__new__(CompositeModelHandler) + enc = MagicMock(model_path=str(tmp_path / "encoder.onnx")) + dec = MagicMock(model_path=str(tmp_path / "decoder.onnx")) + model.get_model_components = MagicMock(return_value=[("encoder.onnx", enc), ("decoder.onnx", dec)]) + + ref_path = tmp_path / "ref_model" + ref_path.mkdir() + + enc_result = { + "max_abs_error": 0.01, + "elements_above_0_1": 0, + "elements_above_0_01": 1, + "total_elements": 100, + "output_compared": "hidden_states", + } + dec_result = { + "max_abs_error": 0.02, + "elements_above_0_1": 0, + "elements_above_0_01": 2, + "total_elements": 200, + "output_compared": "logits", + } + + instance = self._make_pass_instance() + with patch.object(type(instance), "_run_component_subprocess", side_effect=[enc_result, dec_result]): + result = instance._compute_speech_component_discrepancy(model, str(ref_path)) + + assert "encoder" in result["components"] + assert "decoder" in result["components"] + + +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) diff --git a/test/passes/onnx/test_mnb_to_qdq.py b/test/passes/onnx/test_mnb_to_qdq.py index 391130e2af..db8c372878 100644 --- a/test/passes/onnx/test_mnb_to_qdq.py +++ b/test/passes/onnx/test_mnb_to_qdq.py @@ -149,14 +149,14 @@ def test_mnb_to_qdq(create_mnb_model, nodes_to_exclude, add_zero_point, use_sign original_session.disable_fallback() # disable qdq to mnb fusion so we can test the output of the DQ nodes directly disabled_optimizers = ["QDQSelectorActionTransformer"] - if is_symmetric and use_signed_int and not add_zero_point and use_transpose_op: - # there seems to be a bug in ORT graph optimization which changes the int4 DQ to uint8 DQ - with pytest.raises(Exception, match="uint8"): - onnxruntime.InferenceSession(str(qdq_model.model_path), disabled_optimizers=disabled_optimizers) - return - else: + try: qdq_session = onnxruntime.InferenceSession(str(qdq_model.model_path), disabled_optimizers=disabled_optimizers) - qdq_session.disable_fallback() + except onnxruntime.capi.onnxruntime_pybind11_state.Fail as e: + if is_symmetric and use_signed_int and not add_zero_point and use_transpose_op and "uint8" in str(e): + # Older ORT versions incorrectly require uint8 here. Newer versions accept the signed type. + return + raise + qdq_session.disable_fallback() input_data = {"input": np.random.randn(1, 1, in_dim).astype(np.float32)} original_output = original_session.run(None, input_data)[0] diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index a9ccc956b9..92a925a47a 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -29,7 +29,7 @@ def _stub_mobius_module(): """Stub the optional mobius package into sys.modules for the duration of this module. patch("mobius.build") resolves the module via sys.modules, so it works correctly - even in environments where mobius-ai is not installed (e.g. Olive CI). + even in environments where mobius-onnx is not installed (e.g. Olive CI). The stub is only injected when mobius is absent; if the real package is installed, this fixture is a no-op. """ @@ -74,18 +74,26 @@ def _make_pass(ep: str = ExecutionProvider.CPUExecutionProvider) -> MobiusBuilde def _fake_pkg(keys: list[str], _output_dir: Path) -> MagicMock: - """Create a fake ModelPackage that writes dummy .onnx files when .save() is called.""" + """Create a fake ModelPackage that writes dummy .onnx files when .save() is called. - def _save(directory: str, **_kwargs): + Respects the optional ``components`` filter kwarg passed to ``save()``: only writes + files for components for which ``components(name)`` returns True (or all if None). + """ + + def _save(directory: str, components=None, **_kwargs): out = Path(directory) if len(keys) == 1: - # Single-component: saved as /model.onnx - (out / "model.onnx").write_text("dummy") + # Single-component: saved as /model.onnx. + # Apply the components filter consistently with multi-component behaviour. + key = keys[0] + if components is None or components(key): + (out / "model.onnx").write_text("dummy") else: # Multi-component: saved as //model.onnx for k in keys: - (out / k).mkdir(parents=True, exist_ok=True) - (out / k / "model.onnx").write_text("dummy") + if components is None or components(k): + (out / k).mkdir(parents=True, exist_ok=True) + (out / k / "model.onnx").write_text("dummy") pkg = MagicMock() pkg.keys.return_value = keys @@ -95,6 +103,52 @@ def _save(directory: str, **_kwargs): return pkg +def test_components_to_export_skips_genai_config_generation(tmp_path): + """_write_genai_config is not called when components_to_export filters a subset. + + Regression test: mobius has no API to scope ORT GenAI config generation to a + subset of a package, so generating it against the full (unfiltered) pkg would + reference components that were never actually saved to disk (e.g. a "decoder" + filename when only vision_encoder/embedding were exported). See + https://github.com/microsoft/Olive/pull/2456#discussion_r3807156856. + """ + out = tmp_path / "out" + keys = ["decoder", "vision_encoder", "embedding"] + pkg = _fake_pkg(keys, out) + + p = _make_filtered_pass(["vision_encoder", "embedding"]) + + with ( + patch("mobius.build", return_value=pkg), + patch.object(MobiusBuilder, "_write_genai_config") as mock_write_genai_config, + ): + result = p.run(_make_hf_model("org/vlm"), out) + + mock_write_genai_config.assert_not_called() + assert isinstance(result, CompositeModelHandler) + assert result.model_attributes["additional_files"] == [] + + +def test_components_to_export_none_still_generates_genai_config(tmp_path): + """_write_genai_config is still called for a full (unfiltered) export.""" + out = tmp_path / "out" + keys = ["decoder", "vision_encoder", "embedding"] + pkg = _fake_pkg(keys, out) + + p = _make_pass() + + mock_genai_artifacts = {"genai_config": str(out / "genai_config.json")} + with ( + patch("mobius.build", return_value=pkg), + patch.object(MobiusBuilder, "_write_genai_config", return_value=mock_genai_artifacts) as mock_write, + ): + result = p.run(_make_hf_model("org/vlm"), out) + + mock_write.assert_called_once() + assert isinstance(result, CompositeModelHandler) + assert result.model_attributes["additional_files"] == [str(out / "genai_config.json")] + + def _patch_build(pkg: MagicMock): # Patch mobius.build directly — lazy import inside _run_for_config means # patching the module attribute, not the local binding. @@ -368,7 +422,7 @@ def test_none_execution_provider_falls_back_to_default(tmp_path): assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT -@pytest.mark.skipif(not _HAS_REAL_MOBIUS, reason="mobius-ai is not publicly available in CI yet") +@pytest.mark.skipif(not _HAS_REAL_MOBIUS, reason="mobius-onnx is not publicly available in CI yet") def test_write_genai_config_requires_real_mobius(tmp_path): """Integration smoke test for _write_genai_config when real mobius is installed.""" # This test is intentionally lightweight and only verifies the import path. @@ -479,3 +533,167 @@ def test_no_warning_when_trust_remote_code_false(tmp_path): warning_messages = [call.args[0] for call in mock_logger.warning.call_args_list] assert not any("trust_remote_code" in msg for msg in warning_messages) + + +# --------------------------------------------------------------------------- +# components_to_export filter tests +# --------------------------------------------------------------------------- + + +def _make_filtered_pass(components_to_export, precision: str = "fp16") -> MobiusBuilder: + accelerator_spec = AcceleratorSpec( + accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider + ) + return create_pass_from_dict( + MobiusBuilder, + {"precision": precision, "components_to_export": components_to_export}, + disable_search=True, + accelerator_spec=accelerator_spec, + ) + + +def test_components_to_export_filters_subset(tmp_path): + """Only requested components are saved and returned when components_to_export is set.""" + out = tmp_path / "out" + keys = ["decoder", "vision_encoder", "embedding"] + pkg = _fake_pkg(keys, out) + + p = _make_filtered_pass(["vision_encoder", "embedding"]) + + with _patch_build(pkg): + result = p.run(_make_hf_model("org/vlm"), out) + + assert isinstance(result, CompositeModelHandler) + assert result.model_component_names == ["vision_encoder", "embedding"] + + # pkg.save must have been called with a components filter that excludes decoder + save_kwargs = pkg.save.call_args.kwargs + components_filter = save_kwargs.get("components") + assert components_filter is not None + assert components_filter("vision_encoder") is True + assert components_filter("embedding") is True + assert components_filter("decoder") is False + + # Verify skipped component directory is absent from disk + assert (out / "vision_encoder" / "model.onnx").exists(), "vision_encoder should be on disk" + assert (out / "embedding" / "model.onnx").exists(), "embedding should be on disk" + assert not (out / "decoder").exists(), "decoder directory should not exist on disk (was skipped)" + + +def test_components_to_export_preserves_package_order(tmp_path): + """Returned component order follows the package's own key order, not the request order.""" + out = tmp_path / "out" + keys = ["decoder", "vision_encoder", "embedding"] + pkg = _fake_pkg(keys, out) + + # Request components in the reverse of their package order. + p = _make_filtered_pass(["embedding", "vision_encoder"]) + + with _patch_build(pkg): + result = p.run(_make_hf_model("org/vlm"), out) + + assert isinstance(result, CompositeModelHandler) + assert result.model_component_names == ["vision_encoder", "embedding"], ( + "component order must follow the package's own key order (all_keys), not the " + "order components_to_export happened to list them in" + ) + + +def test_components_to_export_none_exports_all(tmp_path): + """All components are exported when components_to_export is None (default).""" + out = tmp_path / "out" + keys = ["decoder", "vision_encoder", "embedding"] + pkg = _fake_pkg(keys, out) + + with _patch_build(pkg): + result = _make_pass().run(_make_hf_model("org/vlm"), out) + + assert isinstance(result, CompositeModelHandler) + assert result.model_component_names == keys + # pkg.save must have been called without a filter (components=None) + save_kwargs = pkg.save.call_args.kwargs + assert save_kwargs.get("components") is None + + +def test_components_to_export_single_component_via_filter(tmp_path): + """Filtering a multi-component model to one component still returns a CompositeModelHandler. + + Unlike an architecturally single-component model (which uses the root layout), a + filtered multi-component model still uses the component sub-directory layout, + so we always return CompositeModelHandler for multi-component packages. + """ + out = tmp_path / "out" + keys = ["decoder", "vision_encoder", "embedding"] + pkg = _fake_pkg(keys, out) + + p = _make_filtered_pass(["decoder"]) + + with _patch_build(pkg): + result = p.run(_make_hf_model("org/vlm"), out) + + # Multi-component model filtered to 1 → still CompositeModelHandler (component sub-dir layout) + assert isinstance(result, CompositeModelHandler) + assert result.model_component_names == ["decoder"] + # The composite sub-directory layout must be preserved for ORT GenAI. + assert result.model_attributes["no_flatten"] is True + assert (out / "decoder" / "model.onnx").exists() + + +def test_components_to_export_unknown_component_raises(tmp_path): + """ValueError when components_to_export names a component not in the package.""" + out = tmp_path / "out" + pkg = _fake_pkg(["decoder", "vision_encoder"], out) + + p = _make_filtered_pass(["nonexistent"]) + + with _patch_build(pkg), pytest.raises(ValueError, match="unknown component"): + p.run(_make_hf_model("org/vlm"), out) + + +def test_components_to_export_empty_list_raises(tmp_path): + """components_to_export=[] must raise ValueError — empty list is always a mistake.""" + out = tmp_path / "out" + pkg = _fake_pkg(["decoder", "vision_encoder"], out) + + p = _make_filtered_pass([]) + + with _patch_build(pkg), pytest.raises(ValueError, match="cannot be empty"): + p.run(_make_hf_model("org/vlm"), out) + + +def test_components_to_export_in_default_config(): + """components_to_export parameter must appear in _default_config with None default.""" + accelerator_spec = AcceleratorSpec( + accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider + ) + config = MobiusBuilder._default_config(accelerator_spec) # pylint: disable=protected-access + assert "components_to_export" in config + assert config["components_to_export"].default_value is None + assert config["components_to_export"].required is False + + +def test_pkg_save_components_filter_applied(tmp_path): + """pkg.save() is always called with the 'components' filter kwarg. + + Only the requested components land on disk. + """ + out = tmp_path / "out" + keys = ["decoder", "vision_encoder", "embedding"] + pkg = _fake_pkg(keys, out) # _fake_pkg sets a __signature__ that includes 'components' + + p = _make_filtered_pass(["vision_encoder", "embedding"]) + + with _patch_build(pkg): + result = p.run(_make_hf_model("org/vlm"), out) + + # Only the requested components should be returned. + assert isinstance(result, CompositeModelHandler) + assert result.model_component_names == ["vision_encoder", "embedding"] + + # pkg.save must have been called WITH the 'components=' kwarg (modern API path). + assert "components" in pkg.save.call_args.kwargs + + # Requested components must be on disk; decoder must not be. + assert (out / "vision_encoder" / "model.onnx").exists() + assert (out / "embedding" / "model.onnx").exists() + assert not (out / "decoder").exists(), "decoder must not be written when filtered out" diff --git a/test/passes/onnx/test_model_builder.py b/test/passes/onnx/test_model_builder.py index 634cc762aa..3de060b106 100644 --- a/test/passes/onnx/test_model_builder.py +++ b/test/passes/onnx/test_model_builder.py @@ -29,13 +29,14 @@ def _create_test_onnx_model(model_path: Path, node_name: str): onnx.save(model, model_path) -def _mock_genai_builder(monkeypatch, create_model_fn): +def _mock_genai_builder(monkeypatch, create_model_fn, check_extra_options_fn=None): builder_module = types.ModuleType("onnxruntime_genai.models.builder") builder_module.create_model = create_model_fn + builder_module.check_extra_options = check_extra_options_fn or (lambda *args, **kwargs: None) models_module = types.ModuleType("onnxruntime_genai.models") models_module.builder = builder_module genai_module = types.ModuleType("onnxruntime_genai") - genai_module.__version__ = "0.8.0" + genai_module.__version__ = "0.15.0" genai_module.models = models_module monkeypatch.setitem(sys.modules, "onnxruntime_genai", genai_module) monkeypatch.setitem(sys.modules, "onnxruntime_genai.models", models_module) @@ -156,6 +157,7 @@ def fake_create_model(*_, **kwargs): fake_builder = types.ModuleType("onnxruntime_genai.models.builder") fake_builder.create_model = MagicMock(side_effect=fake_create_model) + fake_builder.check_extra_options = MagicMock() fake_models = types.ModuleType("onnxruntime_genai.models") fake_models.builder = fake_builder fake_ort_genai = types.ModuleType("onnxruntime_genai") @@ -225,6 +227,7 @@ def fake_create_model(*_, **kwargs): fake_builder = types.ModuleType("onnxruntime_genai.models.builder") fake_builder.create_model = MagicMock(side_effect=fake_create_model) + fake_builder.check_extra_options = MagicMock() fake_models = types.ModuleType("onnxruntime_genai.models") fake_models.builder = fake_builder fake_ort_genai = types.ModuleType("onnxruntime_genai") @@ -326,3 +329,244 @@ def fake_create_model( assert str(output_folder / "encoder.onnx.data") not in additional_files assert str(output_folder / "decoder.onnx.data") not in additional_files assert str(output_folder / "tokenizer.json") in additional_files + + +def test_olive_quantized_model_raises_for_moe(): + """ModelBuilder must reject Olive-quantized MoE checkpoints. + + Errors out cleanly so the user reaches for an alternative builder + or re-runs RTN without ``moe=True``. + """ + from olive.passes.onnx.model_builder import OliveQuantizedModel + + quant_attrs = { + "config": { + "bits": 4, + "group_size": 32, + "symmetric": True, + "embeds": False, + "lm_head": False, + "tie_word_embeddings": False, + "moe": True, + "overrides": {}, + } + } + with pytest.raises(NotImplementedError, match="MoE"): + OliveQuantizedModel( + quant_type="olive", + input_path="/tmp/does_not_matter", + quant_attrs=quant_attrs, + q_size=64, + kv_size=64, + intermediate_size=64, + num_layers=1, + ) + + +def test_olive_quantized_model_migrates_non_moe_keys(tmp_path): + """M7 regression: ``set_tensor``'s non-MoE key migration must be correct. + + It must correctly map Olive's ``_qweight`` / ``_scales`` / ``_qzeros`` naming + onto ``QuantizedTensorModule``'s bare ``qweight`` / ``scales`` / ``qzeros`` attributes, + with correct ``in_features`` / ``out_features`` / block reshape -- previously only the + ``moe=True``-rejection path had coverage for this code. + """ + from olive.passes.onnx.model_builder import OliveQuantizedModel + + # Produce a real Olive-quantized (non-MoE) checkpoint via the actual Rtn pass. + input_model = make_local_tiny_llama(tmp_path / "hf_model", "hf") + quantized_model = create_pass_from_dict( + Rtn, + { + "bits": 4, + "group_size": 16, + "symmetric": False, + "lm_head": True, + "embeds": True, + }, + disable_search=True, + ).run(input_model, tmp_path / "quantized_model") + + loaded = quantized_model.load_model() + qcfg = loaded.config.quantization_config.to_dict() + + quant_attrs = { + "config": { + "bits": qcfg["bits"], + "group_size": qcfg["group_size"], + "symmetric": qcfg["symmetric"], + "embeds": qcfg["embeds"], + "lm_head": qcfg["lm_head"], + "tie_word_embeddings": qcfg["tie_word_embeddings"], + "moe": qcfg["moe"], + "overrides": qcfg.get("overrides") or {}, + } + } + hidden_size = loaded.config.hidden_size + num_heads = loaded.config.num_attention_heads + num_kv_heads = getattr(loaded.config, "num_key_value_heads", num_heads) + head_dim = hidden_size // num_heads + + model = OliveQuantizedModel( + quant_type="olive", + input_path=quantized_model.model_path, + quant_attrs=quant_attrs, + q_size=hidden_size, + kv_size=num_kv_heads * head_dim, + intermediate_size=loaded.config.intermediate_size, + num_layers=loaded.config.num_hidden_layers, + ) + + q_proj = model.layers[0].self_attn.q_proj + assert q_proj.qweight is not None + assert q_proj.scales is not None + assert q_proj.bits == 4 + assert q_proj.in_features == hidden_size + assert q_proj.out_features == hidden_size + # qweight reshaped to (out_features, num_blocks, blob_size) + assert q_proj.qweight.dim() == 3 + assert q_proj.qweight.shape[0] == hidden_size + + down_proj = model.layers[0].mlp.down_proj + assert down_proj.qweight is not None + assert down_proj.bits == 4 + assert down_proj.in_features == loaded.config.intermediate_size + assert down_proj.out_features == hidden_size + + +def test_olive_quantized_model_applies_regex_overrides(tmp_path): + """``re:``-prefixed override keys must be honored by ModelBuilder. + + ``overrides`` keys are documented (``olive.common.quant.patterns``) to support ``re:`` + regex patterns matched with ``re.fullmatch``. ModelBuilder used to look them up with a + plain ``dict.get``, so a regex-keyed override silently fell back to the global + ``bits``/``group_size`` -- which then miscomputes ``in_features`` and reshapes the packed + ``qweight`` incorrectly. + """ + from olive.passes.onnx.model_builder import OliveQuantizedModel + + default_bits, default_group_size = 4, 16 + override_bits, override_group_size = 8, 32 + override_key = r"re:model\.layers\.0\.mlp\.down_proj" + + input_model = make_local_tiny_llama(tmp_path / "hf_model", "hf") + quantized_model = create_pass_from_dict( + Rtn, + { + "bits": default_bits, + "group_size": default_group_size, + "symmetric": False, + "overrides": {override_key: {"bits": override_bits, "group_size": override_group_size}}, + }, + disable_search=True, + ).run(input_model, tmp_path / "quantized_model") + + loaded = quantized_model.load_model() + qcfg = loaded.config.quantization_config.to_dict() + # The regex key must survive serialization, otherwise this test would pass vacuously. + assert override_key in (qcfg.get("overrides") or {}) + assert loaded.config.num_hidden_layers > 1, "need a second layer to check the non-matched case" + + hidden_size = loaded.config.hidden_size + num_heads = loaded.config.num_attention_heads + num_kv_heads = getattr(loaded.config, "num_key_value_heads", num_heads) + model = OliveQuantizedModel( + quant_type="olive", + input_path=quantized_model.model_path, + quant_attrs={ + "config": { + "bits": qcfg["bits"], + "group_size": qcfg["group_size"], + "symmetric": qcfg["symmetric"], + "embeds": qcfg["embeds"], + "lm_head": qcfg["lm_head"], + "tie_word_embeddings": qcfg["tie_word_embeddings"], + "moe": qcfg["moe"], + "overrides": qcfg.get("overrides") or {}, + } + }, + q_size=hidden_size, + kv_size=num_kv_heads * (hidden_size // num_heads), + intermediate_size=loaded.config.intermediate_size, + num_layers=loaded.config.num_hidden_layers, + ) + + # Matched layer -> overridden bits / group_size (and therefore correct in_features). + matched = model.layers[0].mlp.down_proj + assert matched.bits == override_bits + assert matched.group_size == override_group_size + assert matched.in_features == loaded.config.intermediate_size + assert matched.qweight.shape == ( + hidden_size, + loaded.config.intermediate_size // override_group_size, + override_group_size * override_bits // 8, + ) + + # Non-matched layer -> pass-level defaults. + unmatched = model.layers[1].mlp.down_proj + assert unmatched.bits == default_bits + assert unmatched.group_size == default_group_size + assert unmatched.in_features == loaded.config.intermediate_size + assert unmatched.qweight.shape == ( + hidden_size, + loaded.config.intermediate_size // default_group_size, + default_group_size * default_bits // 8, + ) + + +def test_model_builder_prechecks_extra_options(tmp_path, monkeypatch): + def fake_check_extra_options( + model_name, input_path, output_dir, precision, execution_provider, cache_dir, extra_options + ): + assert model_name == "dummy-model" + assert input_path == "dummy-model" + assert output_dir == str(tmp_path / "output_model") + assert precision == "fp32" + assert execution_provider == "cpu" + assert cache_dir + # Values are serialized the way `--extra_options key=value` would produce them. + assert extra_options["exclude_embeds"] == "true" + assert extra_options["use_qdq"] == "false" + assert extra_options["int4_op_types_to_quantize"] == "MatMul/Gather" + assert extra_options["int4_nodes_to_exclude"] == "node_1,node_2" + # An option the model builder does not treat as a list is left alone. + assert extra_options["int4_block_size"] == 32 + extra_options["hf_details"] = { + "extra_kwargs": {}, + "hf_name": model_name, + "hf_config": Mock(), + } + + def fake_create_model( + model_name, input_path, output_dir, precision, execution_provider, cache_dir, filename, **kwargs + ): + assert "hf_details" in kwargs + output_dir = Path(output_dir) + _create_test_onnx_model(output_dir / filename, "test_node") + (output_dir / "genai_config.json").write_text(json.dumps({"search": {}})) + + _mock_genai_builder(monkeypatch, fake_create_model, fake_check_extra_options) + + input_model = Mock(spec=HfModelHandler) + input_model.model_name_or_path = "dummy-model" + input_model.adapter_path = None + input_model.test_model_config = None + input_model.test_model_path = None + input_model.model_attributes = {} + + p = create_pass_from_dict( + ModelBuilder, + { + "precision": "fp32", + "exclude_embeds": True, + "use_qdq": False, + "int4_block_size": 32, + "int4_op_types_to_quantize": ["MatMul", "Gather"], + "int4_nodes_to_exclude": ["node_1", "node_2"], + }, + disable_search=True, + ) + output_model = p.run(input_model, tmp_path / "output_model") + + assert isinstance(output_model, ONNXModelHandler) + assert Path(output_model.model_path).exists() diff --git a/test/passes/pytorch/test_gptq.py b/test/passes/pytorch/test_gptq.py index 0675255aa9..e18677f57e 100644 --- a/test/passes/pytorch/test_gptq.py +++ b/test/passes/pytorch/test_gptq.py @@ -2,13 +2,14 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +# pylint: disable=protected-access from pathlib import Path import pytest import torch from olive.common.quant.hf_utils import OliveHfQuantizationConfig -from olive.common.quant.nn import QuantLinear +from olive.common.quant.tensor import QuantTensor from olive.hardware.accelerator import AcceleratorSpec, Device from olive.model import HfModelHandler from olive.passes.olive_pass import create_pass_from_dict @@ -16,6 +17,29 @@ from test.utils import get_tiny_phi3, make_local_tiny_llama +def _is_quant(module: torch.nn.Module) -> bool: + if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)): + return False + weight = module._parameters.get("weight") + return weight is not None and isinstance(weight.data, QuantTensor) + + +def _bits(module: torch.nn.Module) -> int: + return module.weight.data.bits + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_gptq_rejects_non_finite_moe_fallback_min_k_multiple(value): + p = create_pass_from_dict( + Gptq, + {"moe_fallback_min_k_multiple": value}, + disable_search=True, + accelerator_spec=AcceleratorSpec(accelerator_type=Device.CPU, execution_provider="CPUExecutionProvider"), + ) + + assert not Gptq.validate_config(p.config, p.accelerator_spec) + + # running on CPU takes time so will only run a subset of tests when GPU is not available @pytest.mark.parametrize( ("model_path", "expected_model_type"), @@ -58,9 +82,9 @@ def test_gptq(tmp_path: Path, model_path: str, expected_model_type: str, group_s assert hasattr(loaded_model.config, "quantization_config") assert isinstance(loaded_model.config.quantization_config, OliveHfQuantizationConfig) assert loaded_model.config.quantization_config.group_size == group_size - assert not any(isinstance(m, torch.nn.Linear) for m in loaded_model.model.layers.modules()) - assert isinstance(loaded_model.model.layers[0].self_attn.o_proj, QuantLinear) - assert loaded_model.model.layers[0].self_attn.o_proj.quantizer.bits == 8 - assert loaded_model.model.layers[0].mlp.down_proj.quantizer.bits == 4 + assert not any(isinstance(m, torch.nn.Linear) and not _is_quant(m) for m in loaded_model.model.layers.modules()) + assert _is_quant(loaded_model.model.layers[0].self_attn.o_proj) + assert _bits(loaded_model.model.layers[0].self_attn.o_proj) == 8 + assert _bits(loaded_model.model.layers[0].mlp.down_proj) == 4 assert loaded_model.config.quantization_config.lm_head == lm_head - assert isinstance(loaded_model.lm_head, QuantLinear) == lm_head + assert _is_quant(loaded_model.lm_head) == lm_head diff --git a/test/passes/pytorch/test_gptq_moe.py b/test/passes/pytorch/test_gptq_moe.py new file mode 100644 index 0000000000..e00e4b4a1f --- /dev/null +++ b/test/passes/pytorch/test_gptq_moe.py @@ -0,0 +1,1075 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=protected-access +"""Tests for calibrated (GPTQ) quantization of fused-MoE expert weights. + +Every model here is a randomly-initialised tiny model built from the architecture's own HF +config (no hub checkpoints), following the precedent in ``test_rtn.py``. +""" + +import logging +import threading +from contextlib import contextmanager +from pathlib import Path + +import pytest +import torch + +from olive.common.hf.wrapper import ModelWrapper +from olive.common.quant.selection import iter_quant_targets +from olive.common.quant.tensor import QuantTensor +from olive.common.quant.utils import WeightQuantizer +from olive.model import HfModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.pytorch.gptq import Gptq +from olive.passes.pytorch.moe_calib import ( + OLIVE_MOE_CALIB_IMPLEMENTATION, + MoeCalibrationError, + MoeCalibrationSession, + _register_calib_implementation, + check_moe_gptq_support, +) +from olive.passes.pytorch.quant_utils import QuantInfo + +# Representative K-last architectures covered by the end-to-end tests. +TESTED_MOE_MODEL_TYPES = [ + "deepseek_v3", + "granitemoe", + "jamba", + "mixtral", + "olmoe", + "phimoe", + "qwen2_moe", + "qwen3_moe", +] + +VOCAB_SIZE = 64 +NUM_EXPERTS = 4 + + +def build_tiny_moe_model(model_type: str) -> torch.nn.Module: + """Build a tiny, randomly-initialised model for one tested MoE architecture.""" + # pylint: disable=unexpected-keyword-arg + # HF config ``__init__``s are not statically resolvable by astroid (every kwarg below is + # a real, documented config field); ``test_rtn.py`` suppresses the same false positive. + import transformers as tf + + torch.manual_seed(0) + common = { + "vocab_size": VOCAB_SIZE, + "hidden_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "intermediate_size": 32, + } + builders = { + "qwen2_moe": lambda: tf.Qwen2MoeForCausalLM( + tf.Qwen2MoeConfig( + **common, + moe_intermediate_size=16, + shared_expert_intermediate_size=16, + num_experts=NUM_EXPERTS, + num_experts_per_tok=2, + decoder_sparse_step=1, + ) + ), + "qwen3_moe": lambda: tf.Qwen3MoeForCausalLM( + tf.Qwen3MoeConfig( + **common, + moe_intermediate_size=16, + num_experts=NUM_EXPERTS, + num_experts_per_tok=2, + decoder_sparse_step=1, + head_dim=16, + ) + ), + "phimoe": lambda: tf.PhimoeForCausalLM( + tf.PhimoeConfig(**common, num_local_experts=NUM_EXPERTS, num_experts_per_tok=2) + ), + "mixtral": lambda: tf.MixtralForCausalLM( + tf.MixtralConfig(**common, num_local_experts=NUM_EXPERTS, num_experts_per_tok=2) + ), + "deepseek_v3": lambda: tf.DeepseekV3ForCausalLM( + tf.DeepseekV3Config( + **common, + moe_intermediate_size=16, + n_routed_experts=NUM_EXPERTS, + num_experts_per_tok=2, + first_k_dense_replace=1, + n_group=1, + topk_group=1, + n_shared_experts=1, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=16, + kv_lora_rank=16, + q_lora_rank=None, + ) + ), + "granitemoe": lambda: tf.GraniteMoeForCausalLM( + tf.GraniteMoeConfig(**common, num_local_experts=NUM_EXPERTS, num_experts_per_tok=2) + ), + "olmoe": lambda: tf.OlmoeForCausalLM(tf.OlmoeConfig(**common, num_experts=NUM_EXPERTS, num_experts_per_tok=2)), + "jamba": lambda: tf.JambaForCausalLM( + tf.JambaConfig( + **common, + num_experts=NUM_EXPERTS, + num_experts_per_tok=2, + expert_layer_period=2, + expert_layer_offset=1, + attn_layer_period=2, + attn_layer_offset=1, + mamba_d_state=8, + mamba_d_conv=2, + mamba_dt_rank=16, + ) + ), + } + return builders[model_type]().eval() + + +def save_tiny_moe_model(save_path, model_type: str) -> HfModelHandler: + """Save a tiny MoE model (plus a trivial local tokenizer) and return its handler. + + The saved config pins ``experts_implementation="eager"``: the default ``grouped_mm`` + backend transposes the fused expert weight, which an Olive ``QuantTensor`` cannot do + (quantization is storage-only). Same precedent as ``test_rtn.py``. + """ + import json + + from tokenizers import Tokenizer, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast + + save_path = Path(save_path) + save_path.mkdir(parents=True, exist_ok=True) + build_tiny_moe_model(model_type).save_pretrained(save_path) + + tokenizer = Tokenizer(models.WordLevel({f"t{i}": i for i in range(VOCAB_SIZE)}, unk_token="t0")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="t0", pad_token="t0").save_pretrained(save_path) + + config_path = save_path / "config.json" + config = json.loads(config_path.read_text()) + config["experts_implementation"] = "eager" + config_path.write_text(json.dumps(config, indent=2)) + + return HfModelHandler(model_path=str(save_path)) + + +@pytest.fixture +def _patched_calibration_dataset(monkeypatch): + """Replace the wikitext calibration dataset with deterministic local random token batches.""" + + def fake_calibration_dataset(model, data_config=None, **kwargs): + generator = torch.Generator().manual_seed(0) + return [ + { + "input_ids": torch.randint(0, VOCAB_SIZE, (1, 16), generator=generator), + "attention_mask": torch.ones(1, 16, dtype=torch.long), + } + for _ in range(3) + ] + + # ``monkeypatch.setattr`` with a dotted string target resolves the module internally, + # so this file never needs its own ``import olive.passes.pytorch.quant_utils`` alongside + # the top-level ``from olive.passes.pytorch.quant_utils import QuantInfo``. + monkeypatch.setattr("olive.passes.pytorch.quant_utils.get_calibration_dataset", fake_calibration_dataset) + + +@contextmanager +def capture_logs(logger_name: str): + """Collect records emitted by ``logger_name`` (Olive's loggers don't propagate to caplog).""" + records: list[str] = [] + + class _Handler(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + logger = logging.getLogger(logger_name) + handler = _Handler(level=logging.INFO) + previous_level = logger.level + logger.addHandler(handler) + logger.setLevel(logging.INFO) + try: + yield records + finally: + logger.removeHandler(handler) + logger.setLevel(previous_level) + + +def _attach_quant_info(experts: torch.nn.Module, group_size: int = 16) -> None: + for pname in ("gate_up_proj", "down_proj"): + experts._parameters[pname].quant_info = QuantInfo( + quantizer=WeightQuantizer(bits=4, symmetric=True, group_size=group_size) + ) + + +# --------------------------------------------------------------------------- +# LayerWrapper mappings +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("model_type", TESTED_MOE_MODEL_TYPES) +def test_layer_wrapper_resolves_experts_and_router(model_type: str): + """Every tested architecture must resolve its experts + router submodules. + + Regression for ``granitemoe`` (``layer.block_sparse_moe``) and ``jamba`` + (``layer.feed_forward``), which previously raised ``AttributeError`` in + ``LayerWrapper.__init__`` before any MoE guard could run. + """ + wrapper = ModelWrapper.from_model(build_tiny_moe_model(model_type)) + assert wrapper.model_type == model_type + + moe_layers = 0 + for layer_wrapper in wrapper.get_layer_wrappers(): + experts = layer_wrapper.get_experts(return_name=False) + router = layer_wrapper.get_router(return_name=False) + if experts is None: + # architectures that interleave dense layers (deepseek_v3, jamba) resolve neither + assert router is None + continue + moe_layers += 1 + assert router is not None + assert experts.gate_up_proj.dim() == 3 + assert experts.gate_up_proj.shape[0] == NUM_EXPERTS + + assert moe_layers > 0 + + +@pytest.mark.parametrize("model_type", TESTED_MOE_MODEL_TYPES) +def test_router_is_never_a_quant_target(model_type: str): + """Routers stay full precision, including Jamba's bare ``nn.Linear`` router.""" + model = build_tiny_moe_model(model_type) + wrapper = ModelWrapper.from_model(model) + router_ids = { + id(sub) + for lw in wrapper.get_layer_wrappers() + if (router := lw.get_router(return_name=False)) is not None + for sub in router.modules() + } + assert router_ids, "expected at least one resolvable router" + + for quantize_moe in (False, True): + targets = list( + iter_quant_targets(model, quantize_lm_head=True, quantize_embeds=True, quantize_moe=quantize_moe) + ) + assert not any(id(module) in router_ids for module, _, _ in targets) + # sanity: the walk is not empty (so the assertion above is meaningful) + assert targets + + +def test_jamba_router_linear_would_otherwise_be_selected(): + """Guard the specific Jamba risk: its router *is* an ``nn.Linear`` in the 2D walk.""" + model = build_tiny_moe_model("jamba") + wrapper = ModelWrapper.from_model(model) + routers = [ + router for lw in wrapper.get_layer_wrappers() if (router := lw.get_router(return_name=False)) is not None + ] + assert routers + assert all(isinstance(router, torch.nn.Linear) for router in routers) + + +def test_mamba_block_is_never_a_quant_target(): + """Jamba's Mamba/SSM block stays full precision. + + Covers ``dt_proj``/``in_proj``/``x_proj``/``out_proj``. Regression for a latent selection + gap: ``iter_quant_targets`` had no exclusion for Mamba/SSM submodules, so its bare + ``nn.Linear`` projections (e.g. ``dt_proj``) were swept into the generic 2D walk like any + ordinary MLP weight. This went unnoticed while an older ``transformers`` Mamba forward + implementation happened to call ``self.dt_proj(...)`` directly (so GPTQ's calibration + hook fired and silently mis-quantized it); a later ``transformers`` Mamba forward + refactor stopped calling it directly, leaving ``quant_info.data`` uncollected and + turning the same latent bug into a hard failure. + """ + model = build_tiny_moe_model("jamba") + wrapper = ModelWrapper.from_model(model) + mamba_ids = { + id(sub) + for lw in wrapper.get_layer_wrappers() + if (mamba := lw.get_mamba(return_name=False)) is not None + for sub in mamba.modules() + } + assert mamba_ids, "expected at least one resolvable Mamba block" + + for quantize_moe in (False, True): + targets = list( + iter_quant_targets(model, quantize_lm_head=True, quantize_embeds=True, quantize_moe=quantize_moe) + ) + assert not any(id(module) in mamba_ids for module, _, _ in targets) + # sanity: the walk is not empty (so the assertion above is meaningful) + assert targets + + +# --------------------------------------------------------------------------- +# support gating (fail closed) +# --------------------------------------------------------------------------- + + +class _FakeExperts(torch.nn.Module): + def __init__(self, **flags): + super().__init__() + self.config = object() + for key, value in flags.items(): + setattr(self, key, value) + + +class _ShapeOnly: + """Minimal stand-in for a fused expert parameter: only ``shape`` is needed.""" + + def __init__(self, shape): + self.shape = shape + + +def make_fake_experts(class_name: str = "Qwen3MoeExperts", **flags) -> torch.nn.Module: + """Build a stand-in experts module whose *class name* is ``class_name``.""" + return type(class_name, (_FakeExperts,), {})(**flags) + + +def test_check_support_accepts_model_type_outside_old_allow_list(): + check_moe_gptq_support( + "some_future_moe_architecture", + [make_fake_experts("FutureExperts", is_transposed=False, has_bias=False, has_gate=True)], + ) + + +def test_check_support_rejects_missing_capability(): + """No ``is_transposed`` attribute => transformers too old / class not decorated.""" + with pytest.raises(MoeCalibrationError, match="is_transposed"): + check_moe_gptq_support("qwen3_moe", [make_fake_experts()]) + + +def test_check_support_rejects_transposed_layout(): + with pytest.raises(MoeCalibrationError, match="transposed fused-weight layout"): + check_moe_gptq_support("qwen3_moe", [make_fake_experts(is_transposed=True)]) + + +@pytest.mark.parametrize("bad_value", [None, 0, 1, "False", "", torch.tensor(False)]) +def test_check_support_rejects_non_bool_is_transposed(bad_value): + with pytest.raises(MoeCalibrationError, match="cannot verify"): + check_moe_gptq_support("qwen3_moe", [make_fake_experts(is_transposed=bad_value)]) + + +def test_check_support_rejects_old_transformers(monkeypatch): + monkeypatch.setattr("olive.passes.pytorch.moe_calib._transformers_supports_experts_registry", lambda: False) + with pytest.raises(MoeCalibrationError, match="requires transformers >="): + check_moe_gptq_support("qwen3_moe", [make_fake_experts(is_transposed=False)]) + + +def test_check_support_rejects_module_list_experts(): + experts = torch.nn.ModuleList([torch.nn.Linear(4, 4), torch.nn.Linear(4, 4)]) + with pytest.raises(MoeCalibrationError, match="cannot intercept") as exc_info: + check_moe_gptq_support("classic_moe", [experts]) + assert "moe=False" in str(exc_info.value) + assert "Rtn pass" in str(exc_info.value) + + +def test_check_support_rejects_expert_bias(): + experts = make_fake_experts(is_transposed=False, has_bias=True, has_gate=True) + with pytest.raises(MoeCalibrationError, match="expert biases"): + check_moe_gptq_support("qwen3_moe", [experts]) + + +def test_check_support_rejects_non_gated_experts(): + experts = make_fake_experts(is_transposed=False, has_bias=False, has_gate=False) + with pytest.raises(MoeCalibrationError, match="non-gated experts"): + check_moe_gptq_support("qwen3_moe", [experts]) + + +def test_hessian_memory_preflight_warns_for_large_configs(): + """A DeepSeek-V3-scale config must warn about Hessian memory before calibration starts.""" + experts = make_fake_experts("DeepseekV3Experts", is_transposed=False, has_bias=False, has_gate=True) + experts.num_experts = 256 + # shapes only: (num_experts, out, in) with hidden=7168 and moe_intermediate=2048 + experts.gate_up_proj = _ShapeOnly((256, 4096, 7168)) + experts.down_proj = _ShapeOnly((256, 7168, 2048)) + + with capture_logs("olive.passes.pytorch.moe_calib") as records: + check_moe_gptq_support("deepseek_v3", [experts]) + + messages = "\n".join(records) + assert "Hessians" in messages + # 256 * (7168^2 + 2048^2) * 4 bytes ~= 53.0 GiB + assert "53." in messages + + +def test_hessian_memory_preflight_silent_for_small_configs(): + """Tiny models must not emit the memory warning.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + with capture_logs("olive.passes.pytorch.moe_calib") as records: + assert MoeCalibrationSession.create(wrapper) is not None + assert not any("Hessians" in message for message in records) + + +def test_session_create_returns_none_for_dense_model(): + from transformers import LlamaConfig, LlamaForCausalLM + + config = LlamaConfig( # pylint: disable=unexpected-keyword-arg + vocab_size=32, hidden_size=32, intermediate_size=64, num_hidden_layers=1, num_attention_heads=2 + ) + assert MoeCalibrationSession.create(ModelWrapper.from_model(LlamaForCausalLM(config))) is None + + +# --------------------------------------------------------------------------- +# per-expert recording +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("model_type", TESTED_MOE_MODEL_TYPES) +def test_per_expert_hessians_are_isolated(model_type: str): + """Each expert gets its own (K, K) Hessian built only from the tokens routed to it.""" + model = build_tiny_moe_model(model_type) + wrapper = ModelWrapper.from_model(model) + session = MoeCalibrationSession.create(wrapper) + assert session is not None + + experts = next( + lw.get_experts(return_name=False) + for lw in wrapper.get_layer_wrappers() + if lw.get_experts(return_name=False) is not None + ) + _attach_quant_info(experts) + + session.start() + try: + with session.record([experts]), torch.no_grad(): + model(torch.randint(0, VOCAB_SIZE, (2, 16))) + finally: + session.finish() + + data = experts.gate_up_proj.quant_info.data + assert data["moe"] is True + assert data["tokens_seen"] == 32 + assert len(data["token_counts"]) == NUM_EXPERTS + # top_k=2 => every token is counted for exactly two experts + assert sum(data["token_counts"]) == 32 * 2 + + hidden_size = experts.gate_up_proj.shape[-1] + for expert_idx, entry in data["experts"].items(): + assert entry["H"].shape == (hidden_size, hidden_size) + # the Hessian saw exactly the tokens routed to that expert -- no cross-expert pooling + assert entry["N"] == data["token_counts"][expert_idx] + + down_data = experts.down_proj.quant_info.data + intermediate_size = experts.down_proj.shape[-1] + for entry in down_data["experts"].values(): + assert entry["H"].shape == (intermediate_size, intermediate_size) + + +def test_recording_switch_suppresses_second_pass(): + """The true-sequential re-run must not double-count into the Hessians.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + session = MoeCalibrationSession.create(wrapper) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + _attach_quant_info(experts) + + inputs = torch.randint(0, VOCAB_SIZE, (2, 16)) + session.start() + try: + with session.record([experts]), torch.no_grad(): + model(inputs) + recorded = {e: entry["N"] for e, entry in experts.gate_up_proj.quant_info.data["experts"].items()} + hessians = {e: entry["H"].clone() for e, entry in experts.gate_up_proj.quant_info.data["experts"].items()} + + # second (post-quantization) pass -- recording is off + with torch.no_grad(): + model(inputs) + finally: + session.finish() + + after = {e: entry["N"] for e, entry in experts.gate_up_proj.quant_info.data["experts"].items()} + assert after == recorded + for expert_idx, hessian in hessians.items(): + assert torch.equal(experts.gate_up_proj.quant_info.data["experts"][expert_idx]["H"], hessian) + + +def test_calibration_forward_matches_eager_output(): + """The recording experts forward must return the model's normal output.""" + model = build_tiny_moe_model("qwen3_moe") + model.set_experts_implementation("eager") + inputs = torch.randint(0, VOCAB_SIZE, (2, 16)) + with torch.no_grad(): + expected = model(inputs).logits + + wrapper = ModelWrapper.from_model(model) + session = MoeCalibrationSession.create(wrapper) + session.start() + try: + with torch.no_grad(): + actual = model(inputs).logits + finally: + session.finish() + + torch.testing.assert_close(actual, expected) + assert model.config._experts_implementation == "eager" + + +def test_coverage_report_flags_unseen_experts(): + """Coverage is reported per layer + summarized, and warns (never raises) when thin.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + session = MoeCalibrationSession.create(wrapper) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + _attach_quant_info(experts) + + # bias routing hard toward expert 0 so at least one expert is never routed to + router = wrapper.get_layer_wrappers()[0].get_router(return_name=False) + with torch.no_grad(): + router.weight.zero_() + router.weight[0] = 10.0 + + session.start() + with capture_logs("olive.passes.pytorch.moe_calib") as records: + try: + with session.record([experts]), torch.no_grad(): + model(torch.randint(0, VOCAB_SIZE, (2, 16))) + session.add_coverage("model.layers.0", experts) + finally: + session.finish() + + counts = experts.gate_up_proj.quant_info.data["token_counts"] + assert counts[0] > 0 + assert sum(1 for c in counts if c == 0) >= 1 + messages = "\n".join(records) + assert "MoE coverage [model.layers.0]" in messages + assert "MoE coverage summary" in messages + assert "unseen" in messages + + +def test_token_counts_do_not_depend_on_gate_up_proj_being_quantized(): + """Coverage must stay correct when only ``down_proj`` carries ``quant_info``. + + ``token_counts`` is derived from the recorded Hessian sample counts, so excluding + ``gate_up_proj`` (e.g. via ``modules_to_not_convert``) must not report every expert as + unseen. + """ + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + session = MoeCalibrationSession.create(wrapper) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + experts._parameters["down_proj"].quant_info = QuantInfo( + quantizer=WeightQuantizer(bits=4, symmetric=True, group_size=16) + ) + + session.start() + try: + with session.record([experts]), torch.no_grad(): + model(torch.randint(0, VOCAB_SIZE, (2, 16))) + finally: + session.finish() + + data = experts.down_proj.quant_info.data + assert not hasattr(experts.gate_up_proj, "quant_info") + assert sum(data["token_counts"]) == 32 * 2 + for expert_idx, entry in data["experts"].items(): + assert entry["N"] == data["token_counts"][expert_idx] + + +# --------------------------------------------------------------------------- +# session lifecycle +# --------------------------------------------------------------------------- + + +def test_start_is_not_re_entrant(): + """A second ``start()`` would clobber the saved implementation, so it must be refused.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + original_implementation = model.get_experts_implementation() + session = MoeCalibrationSession.create(wrapper) + + session.start() + try: + with pytest.raises(MoeCalibrationError, match="already active"): + session.start() + finally: + session.finish() + + assert model.get_experts_implementation() == original_implementation + + +def test_second_session_on_the_same_model_is_refused(): + """Two overlapping sessions would restore each other's calibration implementation.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + original_implementation = model.get_experts_implementation() + first = MoeCalibrationSession.create(wrapper) + second = MoeCalibrationSession.create(wrapper) + + first.start() + try: + with pytest.raises(MoeCalibrationError, match="another MoE calibration session"): + second.start() + finally: + first.finish() + + assert model.get_experts_implementation() == original_implementation + + +def test_concurrent_sessions_cannot_both_claim_the_same_model(monkeypatch): + """A session reserves the model before swapping, closing the start() check/set race.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + first = MoeCalibrationSession.create(wrapper) + second = MoeCalibrationSession.create(wrapper) + original_setter = model.set_experts_implementation + setter_entered = threading.Event() + release_setter = threading.Event() + block_lock = threading.Lock() + should_block = True + + def blocking_setter(implementation): + nonlocal should_block + with block_lock: + block_this_call = implementation == OLIVE_MOE_CALIB_IMPLEMENTATION and should_block + if block_this_call: + should_block = False + if block_this_call: + setter_entered.set() + assert release_setter.wait(timeout=10) + return original_setter(implementation) + + monkeypatch.setattr(model, "set_experts_implementation", blocking_setter) + first_errors = [] + + def start_first(): + try: + first.start() + except Exception as exc: + first_errors.append(exc) + + thread = threading.Thread(target=start_first) + thread.start() + assert setter_entered.wait(timeout=10) + try: + with pytest.raises(MoeCalibrationError, match="another MoE calibration session"): + second.start() + finally: + release_setter.set() + thread.join(timeout=10) + if first._active: + first.finish() + + assert not thread.is_alive() + assert not first_errors + + +def test_record_rejects_nested_context_for_same_experts(): + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + session = MoeCalibrationSession.create(wrapper) + experts = session.experts_modules[0] + + session.start() + try: + with ( + session.record([experts]), + pytest.raises(MoeCalibrationError, match=r"nested or concurrent record\(\) contexts"), + session.record([experts]), + ): + pass + finally: + session.finish() + + +def test_record_requires_active_session(): + model = build_tiny_moe_model("qwen3_moe") + session = MoeCalibrationSession.create(ModelWrapper.from_model(model)) + + with ( + pytest.raises(MoeCalibrationError, match=r"record\(\) requires an active session"), + session.record([session.experts_modules[0]]), + ): + pass + + +def test_start_restores_the_model_when_the_swap_is_stale(): + """A failed swap validation must not leave the model in calibration mode.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + original_implementation = model.get_experts_implementation() + session = MoeCalibrationSession.create(wrapper) + + # an experts module whose config never picks up the swap => stale + unswitchable = make_fake_experts(is_transposed=False) + unswitchable.config = type("_Config", (), {"_experts_implementation": "eager"})() + session.experts_modules = [*session.experts_modules, unswitchable] + + with pytest.raises(MoeCalibrationError, match="could not switch the experts implementation"): + session.start() + + assert model.get_experts_implementation() == original_implementation + assert session._active is False + + +def test_start_reports_stale_swap_when_restore_fails(monkeypatch): + """A restore failure must not mask the actionable stale-swap error.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + original_implementation = model.get_experts_implementation() + original_setter = model.set_experts_implementation + session = MoeCalibrationSession.create(wrapper) + + unswitchable = make_fake_experts(is_transposed=False) + unswitchable.config = type("_Config", (), {"_experts_implementation": "eager"})() + session.experts_modules = [*session.experts_modules, unswitchable] + + def setter_with_failed_restore(implementation): + if implementation == original_implementation: + raise RuntimeError("restore failed") + return original_setter(implementation) + + monkeypatch.setattr(model, "set_experts_implementation", setter_with_failed_restore) + try: + with ( + capture_logs("olive.passes.pytorch.moe_calib") as records, + pytest.raises(MoeCalibrationError, match="restoring the original experts implementation also failed"), + ): + session.start() + finally: + original_setter(original_implementation) + + assert any("Failed to restore the original experts implementation" in message for message in records) + assert session._active is False + + +def test_registry_key_collision_is_detected(monkeypatch): + """A foreign function under Olive's registry key must fail closed, not run silently.""" + from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS + + key = OLIVE_MOE_CALIB_IMPLEMENTATION + previous = ALL_EXPERTS_FUNCTIONS.get(key) + + def foreign_experts_forward(*args, **kwargs): + raise AssertionError("this must never be called") + + ALL_EXPERTS_FUNCTIONS.register(key, foreign_experts_forward) + try: + with pytest.raises(MoeCalibrationError, match="already maps"): + _register_calib_implementation() + finally: + if previous is None: + del ALL_EXPERTS_FUNCTIONS[key] + else: + ALL_EXPERTS_FUNCTIONS.register(key, previous) + + +@pytest.mark.usefixtures("_patched_calibration_dataset") +def test_calibration_state_is_restored_when_processing_raises(tmp_path: Path, monkeypatch): + """An exception mid-calibration must still restore ``experts_implementation`` + ``use_cache``. + + Both are process-global mutations owned by ``run_layerwise_quantization``; leaking them + would silently corrupt any later use of the same in-memory model. + """ + import importlib + + gptq_module = importlib.import_module("olive.passes.pytorch.gptq") + + input_model = save_tiny_moe_model(tmp_path / "input_model", "qwen3_moe") + + captured = {} + original_prepare_model = gptq_module.prepare_model + + def spy_prepare_model(model, config): + result = original_prepare_model(model, config) + captured["wrapper"] = result[0] + captured["experts_implementation"] = result[0].model.config._experts_implementation + captured["use_cache"] = result[0].model.config.use_cache + return result + + monkeypatch.setattr(gptq_module, "prepare_model", spy_prepare_model) + + calls = [] + original_process_module = Gptq.process_module + + def flaky_process_module(module, **kwargs): + calls.append(module) + if len(calls) == 2: + raise RuntimeError("injected calibration failure") + return original_process_module(module, **kwargs) + + monkeypatch.setattr(Gptq, "process_module", staticmethod(flaky_process_module)) + + p = create_pass_from_dict(Gptq, {"bits": 4, "group_size": 16, "sym": True, "moe": True}, disable_search=True) + with pytest.raises(RuntimeError, match="injected calibration failure"): + p.run(input_model, str(tmp_path / "gptq")) + + config = captured["wrapper"].model.config + assert config._experts_implementation == captured["experts_implementation"] == "eager" + assert config.use_cache == captured["use_cache"] + + +# --------------------------------------------------------------------------- +# LayerWrapper projection accessors +# --------------------------------------------------------------------------- + + +def test_get_attention_inputs_is_strict_by_default(): + """Positional consumers (rotate.py) must get an error, not a silently shorter list. + + DeepSeek-V3's MLA has no ``k_proj``/``v_proj``; with the strict default the missing + projections raise instead of shifting ``v_proj`` to another index. + """ + wrapper = ModelWrapper.from_model(build_tiny_moe_model("deepseek_v3")) + layer_wrapper = wrapper.get_layer_wrappers()[0] + + with pytest.raises(AttributeError): + layer_wrapper.get_attention_inputs() + + partial = layer_wrapper.get_attention_inputs(return_name=False, partial_ok=True) + assert len(partial) == 1 # q_proj only + + +def test_get_attention_inputs_keeps_qkv_order_for_standard_attention(): + """The positional contract (index 2 == ``v_proj``) still holds for normal attention.""" + wrapper = ModelWrapper.from_model(build_tiny_moe_model("qwen3_moe")) + _, names = wrapper.get_layer_wrappers()[0].get_attention_inputs() + assert [name.rsplit(".", 1)[-1] for name in names] == ["q_proj", "k_proj", "v_proj"] + + +@pytest.mark.parametrize("model_type", TESTED_MOE_MODEL_TYPES) +def test_get_mlp_projections_require_explicit_partial_mode_for_moe_layers(model_type: str): + """Strict consumers must fail rather than silently skip an MoE layer's projections.""" + wrapper = ModelWrapper.from_model(build_tiny_moe_model(model_type)) + moe_layers = 0 + for layer_wrapper in wrapper.get_layer_wrappers(): + if layer_wrapper.get_experts(return_name=False) is None: + continue + moe_layers += 1 + with pytest.raises(AttributeError): + layer_wrapper.get_mlp_inputs(return_name=False) + with pytest.raises(AttributeError): + layer_wrapper.get_mlp_outputs(return_name=False) + assert layer_wrapper.get_mlp_inputs(return_name=False, partial_ok=True) == [] + assert layer_wrapper.get_mlp_outputs(return_name=False, partial_ok=True) == [] + assert moe_layers > 0 + + +# --------------------------------------------------------------------------- +# RTN fallback +# --------------------------------------------------------------------------- + + +def _run_recording(model, wrapper, experts, fallback_threshold: float): + session = MoeCalibrationSession.create(wrapper, fallback_threshold=fallback_threshold) + session.start() + try: + with session.record([experts]), torch.no_grad(): + model(torch.randint(0, VOCAB_SIZE, (2, 16))) + finally: + session.finish() + return session + + +def test_rtn_fallback_when_below_threshold(): + """Below the threshold an expert is RTN-quantized: float weight kept, RTN qparams.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + _attach_quant_info(experts) + # threshold of 1.0 => an expert must see *every* calibration token to avoid the fallback, + # which top-k routing can never satisfy for more than top_k experts. + _run_recording(model, wrapper, experts, fallback_threshold=1.0) + + original = experts.gate_up_proj.data.clone() + quantizer = experts.gate_up_proj.quant_info.quantizer + with capture_logs("olive.passes.pytorch.gptq") as records: + Gptq.process_module(experts, moe_fallback_threshold=1.0) + + assert any("quantized with RTN" in message for message in records) + # RTN derives qparams straight from the float weight, and the weight is left + # fake-quantized (on-grid) so the true-sequential re-run sees the same invariant as the + # GPTQ experts. + expected_scales, expected_zp = quantizer.find_qparams(original.float()) + expected_weight = quantizer.fake_quantize(original.float(), expected_scales, expected_zp) + torch.testing.assert_close(experts.gate_up_proj.data, expected_weight.to(original.dtype)) + torch.testing.assert_close(experts.gate_up_proj.quant_info.scales, expected_scales.cpu()) + torch.testing.assert_close(experts.gate_up_proj.quant_info.zero_points, expected_zp.cpu()) + + +def test_gptq_path_used_when_above_threshold(): + """When both thresholds are satisfied, every routed expert is GPTQ-quantized (not RTN).""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + _attach_quant_info(experts) + _run_recording(model, wrapper, experts, fallback_threshold=0.005) + + original = experts.gate_up_proj.data.clone() + # min_k_multiple=0.0 isolates this test to the skew condition alone (this test's + # purpose): every routed expert must clear the *default* skew threshold on this tiny + # fixture's calibration size, but not necessarily the *default* sufficiency threshold + # (2.0x K) -- that condition is covered separately below. + with capture_logs("olive.passes.pytorch.gptq") as records: + Gptq.process_module(experts, moe_fallback_threshold=0.005, moe_fallback_min_k_multiple=0.0) + + assert not any("quantized with RTN" in message for message in records) + assert not torch.equal(experts.gate_up_proj.data, original) + num_groups = original.shape[-1] // 16 + assert experts.gate_up_proj.quant_info.scales.shape == (NUM_EXPERTS, original.shape[1], num_groups) + assert experts.gate_up_proj.quant_info.data is None + + +def test_rtn_fallback_when_below_sufficiency_threshold_even_if_skew_passes(): + """Sufficiency (min_k_multiple) must independently gate fallback, even when skew passes. + + Regression for the dual-condition OR-gate: with a wide-open skew threshold (0.0, so it + never fires) but a large ``moe_fallback_min_k_multiple``, every expert must still fall + back to RTN once its observed token count is below ``min_k_multiple * K`` -- this is the + condition added in the dual fallback-threshold commit, and it previously had zero direct + test coverage. + """ + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + _attach_quant_info(experts) + # skew_threshold = 0.0 * tokens_seen = 0, so it can never fail; sufficiency_threshold = + # 1000.0 * K is unreachable by this tiny fixture's calibration set, so every expert must + # fail on sufficiency alone. + _run_recording(model, wrapper, experts, fallback_threshold=0.0) + + original = experts.gate_up_proj.data.clone() + quantizer = experts.gate_up_proj.quant_info.quantizer + with capture_logs("olive.passes.pytorch.gptq") as records: + Gptq.process_module(experts, moe_fallback_threshold=0.0, moe_fallback_min_k_multiple=1000.0) + + assert any("quantized with RTN" in message for message in records) + assert any(f"{NUM_EXPERTS}/{NUM_EXPERTS} experts quantized with RTN" in message for message in records) + expected_scales, expected_zp = quantizer.find_qparams(original.float()) + expected_weight = quantizer.fake_quantize(original.float(), expected_scales, expected_zp) + torch.testing.assert_close(experts.gate_up_proj.data, expected_weight.to(original.dtype)) + + +def test_zero_sample_expert_falls_back_without_hessian(): + """An expert that was never routed to has no Hessian at all and must not crash GPTQ.""" + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + _attach_quant_info(experts) + + # bias routing hard toward expert 0 so at least one expert is never routed to + router = wrapper.get_layer_wrappers()[0].get_router(return_name=False) + with torch.no_grad(): + router.weight.zero_() + router.weight[0] = 10.0 + _run_recording(model, wrapper, experts, fallback_threshold=0.005) + + data = experts.gate_up_proj.quant_info.data + unseen = [idx for idx, count in enumerate(data["token_counts"]) if count == 0] + assert unseen + original = experts.gate_up_proj.data.clone() + quantizer = experts.gate_up_proj.quant_info.quantizer + + Gptq.process_module(experts, moe_fallback_threshold=0.005) + + for idx in unseen: + # no Hessian => RTN fallback: the weight is the plain fake-quantized original + expected = quantizer.fake_quantize(original[idx].float()) + torch.testing.assert_close(experts.gate_up_proj.data[idx], expected.to(original.dtype)) + + +def test_rtn_fallback_weights_are_on_grid_before_true_sequential_rerun(): + """Fallback experts must be fake-quantized (not left float) after ``process_module``. + + The true-sequential loop re-runs the layer after ``process_module``; if fallback experts + still held raw float weights there, the next layer would be calibrated against a + higher-precision layer than the one that is actually saved. + """ + model = build_tiny_moe_model("qwen3_moe") + wrapper = ModelWrapper.from_model(model) + experts = wrapper.get_layer_wrappers()[0].get_experts(return_name=False) + _attach_quant_info(experts) + _run_recording(model, wrapper, experts, fallback_threshold=1.0) # forces every expert to fall back + + quantizer = experts.gate_up_proj.quant_info.quantizer + Gptq.process_module(experts, moe_fallback_threshold=1.0) + + info = experts.gate_up_proj.quant_info + weight = experts.gate_up_proj.data.float() + # already on the quantization grid => re-applying the recorded qparams is a no-op, + # which is exactly what ``finalize`` does when it serializes the weight + on_grid = quantizer.fake_quantize(weight, info.scales.to(weight.device), info.zero_points.to(weight.device)) + torch.testing.assert_close(on_grid, weight) + + +# --------------------------------------------------------------------------- +# end to end +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("model_type", TESTED_MOE_MODEL_TYPES) +@pytest.mark.usefixtures("_patched_calibration_dataset") +def test_gptq_moe_end_to_end(tmp_path: Path, model_type: str): + """The real ``Gptq`` pass quantizes fused expert weights for every tested model.""" + input_model = save_tiny_moe_model(tmp_path / "input_model", model_type) + p = create_pass_from_dict( + Gptq, + {"bits": 4, "group_size": 16, "sym": True, "moe": True}, + disable_search=True, + ) + out = p.run(input_model, str(tmp_path / "gptq")) + assert isinstance(out, HfModelHandler) + + loaded = out.load_model() + assert loaded.config.quantization_config.moe is True + + wrapper = ModelWrapper.from_model(loaded) + moe_layers = 0 + for layer_wrapper in wrapper.get_layer_wrappers(): + experts = layer_wrapper.get_experts(return_name=False) + if experts is None: + continue + moe_layers += 1 + for pname in ("gate_up_proj", "down_proj"): + param = experts._parameters[pname] + assert isinstance(param.data, QuantTensor), f"{pname} was not quantized" + assert param.data.scales.dim() == 3 + # routers stay full precision + router = layer_wrapper.get_router(return_name=False) + assert not any(isinstance(p.data, QuantTensor) for p in router.parameters()) + assert moe_layers > 0 + + # ``finalize()`` re-serializes config.json without the custom ``experts_implementation`` + # field, so re-pin ``eager`` for the forward (the default ``grouped_mm`` transposes the + # fused weight, which a storage-only QuantTensor cannot do). + loaded.set_experts_implementation("eager") + loaded.eval() + with torch.no_grad(): + logits = loaded(torch.randint(0, VOCAB_SIZE, (1, 8))).logits + assert torch.isfinite(logits).all() + + +@pytest.mark.usefixtures("_patched_calibration_dataset") +def test_gptq_moe_disabled_leaves_experts_alone(tmp_path: Path): + """``moe=False`` (the default) keeps fused expert weights in full precision.""" + input_model = save_tiny_moe_model(tmp_path / "input_model", "qwen3_moe") + p = create_pass_from_dict(Gptq, {"bits": 4, "group_size": 16, "sym": True}, disable_search=True) + out = p.run(input_model, str(tmp_path / "gptq")) + + loaded = out.load_model() + experts = loaded.model.layers[0].mlp.experts + assert not any(isinstance(param.data, QuantTensor) for param in experts.parameters()) + assert isinstance(loaded.model.layers[0].self_attn.q_proj.weight.data, QuantTensor) + + +@pytest.mark.usefixtures("_patched_calibration_dataset") +def test_gptq_moe_accepts_model_type_outside_old_allow_list(tmp_path: Path, monkeypatch): + """A K-last architecture is accepted without a ``model_type`` allow-list entry.""" + input_model = save_tiny_moe_model(tmp_path / "input_model", "qwen3_moe") + + # Pretend the K-last model is a future architecture; only its diagnostic name changes. + original = ModelWrapper.from_model + + def patched_from_model(model): + wrapper = original(model) + wrapper.model_type = "some_future_moe_architecture" + return wrapper + + monkeypatch.setattr(ModelWrapper, "from_model", staticmethod(patched_from_model)) + + p = create_pass_from_dict(Gptq, {"bits": 4, "group_size": 16, "sym": True, "moe": True}, disable_search=True) + out = p.run(input_model, str(tmp_path / "gptq")) + assert isinstance(out, HfModelHandler) diff --git a/test/passes/pytorch/test_kquant.py b/test/passes/pytorch/test_kquant.py index 8af686a0ec..323ff724a3 100644 --- a/test/passes/pytorch/test_kquant.py +++ b/test/passes/pytorch/test_kquant.py @@ -2,21 +2,70 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +# pylint: disable=protected-access from pathlib import Path import pytest import torch from olive.common.quant.hf_utils import OliveHfQuantizationConfig -from olive.common.quant.nn import QuantEmbedding, QuantLinear +from olive.common.quant.tensor import QuantTensor from olive.common.quant.utils import WeightQuantizer, get_maxq_minq from olive.hardware.accelerator import AcceleratorSpec, Device from olive.model import HfModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.pytorch.kquant import KQuant, kquant_find_qparams +from olive.passes.pytorch.moe_support import MoeSupportError +from olive.passes.pytorch.quant_utils import prepare_model from test.utils import get_tiny_phi3 +def _save_trivial_tokenizer(save_path: Path, vocab_size: int) -> None: + """Save a local tokenizer so pass metadata serialization never needs the hub.""" + from tokenizers import Tokenizer, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast + + tokenizer = Tokenizer(models.WordLevel({f"t{i}": i for i in range(vocab_size)}, unk_token="t0")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="t0", pad_token="t0").save_pretrained(save_path) + + +def _make_local_tiny_qwen3_moe(save_path: Path) -> HfModelHandler: + """Save a tiny K-last fused-experts model without downloading a checkpoint.""" + from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM + + torch.manual_seed(0) + save_path.mkdir(parents=True, exist_ok=True) + config = Qwen3MoeConfig( # pylint: disable=unexpected-keyword-arg + vocab_size=32, + hidden_size=16, + intermediate_size=16, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_experts=2, + num_experts_per_tok=1, + decoder_sparse_step=1, + head_dim=8, + experts_implementation="eager", + ) + Qwen3MoeForCausalLM(config).save_pretrained(save_path) + _save_trivial_tokenizer(save_path, config.vocab_size) + return HfModelHandler(model_path=str(save_path)) + + +def _is_quant(module: torch.nn.Module) -> bool: + if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)): + return False + weight = module._parameters.get("weight") + return weight is not None and isinstance(weight.data, QuantTensor) + + +def _bits(module: torch.nn.Module) -> int: + return module.weight.data.bits + + @pytest.mark.parametrize("sym", [True, False]) @pytest.mark.parametrize("bits", [2, 4]) def test_kquant_find_qparams_beats_min_max_rtn(bits: int, sym: bool): @@ -50,6 +99,138 @@ def test_kquant_find_qparams_handles_constant_groups(values: float, sym: bool): assert torch.allclose(dq, weight, atol=1e-5) +@pytest.mark.parametrize("sym", [True, False]) +def test_kquant_find_qparams_3d_matches_per_expert_results(sym: bool): + torch.manual_seed(1) + weight = torch.randn(3, 5, 32, dtype=torch.float32) + group_size = 8 + maxq, minq = get_maxq_minq(4, signed=False) + + scales, zero_points = kquant_find_qparams(weight, group_size, maxq, minq, symmetric=sym) + + assert scales.shape == (3, 5, 4) + assert zero_points.shape == (3, 5, 4) + for expert_idx in range(weight.shape[0]): + expert_scales, expert_zero_points = kquant_find_qparams( + weight[expert_idx], group_size, maxq, minq, symmetric=sym + ) + torch.testing.assert_close(scales[expert_idx], expert_scales) + torch.testing.assert_close(zero_points[expert_idx], expert_zero_points) + + +@pytest.mark.parametrize( + ("layout", "message"), + [ + (True, r"\(E, K, OUT\)"), + (None, "is_transposed"), + ("False", "is_transposed"), + ("missing", "is_transposed"), + ], +) +def test_kquant_moe_rejects_unsafe_or_unverifiable_layout(tmp_path: Path, monkeypatch, layout: object, message: str): + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + + def patched_prepare_model(*args, **kwargs): + wrapper, qcfg, retie = prepare_model(*args, **kwargs) + for layer in wrapper.get_layer_wrappers(): + experts = layer.get_experts(return_name=False) + if experts is None: + continue + if layout == "missing": + if hasattr(experts, "is_transposed"): + del experts.is_transposed + else: + experts.is_transposed = layout + return wrapper, qcfg, retie + + monkeypatch.setattr("olive.passes.pytorch.kquant.prepare_model", patched_prepare_model) + quantizer = create_pass_from_dict(KQuant, {"moe": True, "group_size": -1}, disable_search=True) + + with pytest.raises(MoeSupportError, match=message): + quantizer.run(input_model, str(tmp_path / "kquant")) + + +def test_kquant_moe_module_list_without_direct_3d_parameter_is_exempt(tmp_path: Path, monkeypatch): + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + + def patched_prepare_model(*args, **kwargs): + wrapper, qcfg, retie = prepare_model(*args, **kwargs) + classic_experts = torch.nn.ModuleList([torch.nn.Linear(4, 4) for _ in range(2)]) + for layer in wrapper.get_layer_wrappers(): + layer.get_experts = lambda return_name=True, experts=classic_experts: ( + (experts, "mlp.experts") if return_name else experts + ) + return wrapper, qcfg, retie + + monkeypatch.setattr("olive.passes.pytorch.kquant.prepare_model", patched_prepare_model) + quantizer = create_pass_from_dict(KQuant, {"moe": True, "group_size": -1}, disable_search=True) + + out = quantizer.run(input_model, str(tmp_path / "kquant")) + + assert isinstance(out, HfModelHandler) + + +def test_kquant_moe_k_last_quantize_dequantize_roundtrip(tmp_path: Path): + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + original = input_model.load_model().model.layers[0].mlp.experts.gate_up_proj.detach().clone() + quantizer = create_pass_from_dict( + KQuant, + {"bits": 4, "group_size": 4, "moe": True}, + disable_search=True, + ) + + out = quantizer.run(input_model, str(tmp_path / "kquant")) + + experts = out.load_model().model.layers[0].mlp.experts + quantized = experts.gate_up_proj.data + assert isinstance(quantized, QuantTensor) + assert quantized.scales.shape == (*original.shape[:-1], original.shape[-1] // 4) + dequantized = quantized.to_dense() + assert torch.isfinite(dequantized).all() + error = (dequantized - original).abs().mean() + relative_error = error / original.abs().mean() + assert 0 < relative_error < 0.15 + + +def test_kquant_moe_false_does_not_run_layout_gate(tmp_path: Path, monkeypatch): + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + + def unexpected_gate(*args, **kwargs): + pytest.fail("MoE layout support check must not run when moe=False") + + monkeypatch.setattr("olive.passes.pytorch.kquant.check_moe_layout_support", unexpected_gate) + quantizer = create_pass_from_dict(KQuant, {"moe": False, "group_size": -1}, disable_search=True) + + out = quantizer.run(input_model, str(tmp_path / "kquant")) + + experts = out.load_model().model.layers[0].mlp.experts + assert not any(isinstance(param.data, QuantTensor) for param in experts.parameters()) + + +def test_kquant_moe_gate_ignores_prior_checkpoint_moe_flag(tmp_path: Path, monkeypatch): + """Regression test: a second KQuant pass with moe=False must not re-run the layout gate. + + ``prepare_model`` ORs a pre-existing checkpoint's ``moe`` flag into the merged + ``qcfg.moe`` (see ``quant_utils.prepare_model``), so gating on ``qcfg.moe`` would make + this second, moe=False invocation incorrectly re-run fused-experts layout validation + -- something this run never asked for. The gate must key off this invocation's own + ``config.moe`` request instead. + """ + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + first_pass = create_pass_from_dict(KQuant, {"bits": 4, "group_size": 4, "moe": True}, disable_search=True) + quantized = first_pass.run(input_model, str(tmp_path / "kquant_first")) + + def unexpected_gate(*args, **kwargs): + pytest.fail("MoE layout support check must not re-run when this invocation requests moe=False") + + monkeypatch.setattr("olive.passes.pytorch.kquant.check_moe_layout_support", unexpected_gate) + second_pass = create_pass_from_dict(KQuant, {"moe": False, "lm_head": True, "group_size": -1}, disable_search=True) + out = second_pass.run(quantized, str(tmp_path / "kquant_second")) + + loaded = out.load_model() + assert isinstance(loaded.lm_head.weight.data, QuantTensor) + + @pytest.mark.parametrize("group_size", [-1, 16]) @pytest.mark.parametrize("sym", [True, False]) @pytest.mark.parametrize("lm_head", [True, False]) @@ -80,12 +261,13 @@ def test_kquant(tmp_path: Path, group_size: int, sym: bool, lm_head: bool): assert loaded_model.config.quantization_config.symmetric is sym assert loaded_model.config.quantization_config.group_size == group_size assert loaded_model.config.quantization_config.lm_head == lm_head - assert not any(isinstance(m, torch.nn.Linear) for m in loaded_model.model.layers.modules()) - assert isinstance(loaded_model.model.layers[0].self_attn.o_proj, QuantLinear) - assert loaded_model.model.layers[0].self_attn.o_proj.quantizer.bits == 8 - assert loaded_model.model.layers[0].mlp.down_proj.quantizer.bits == 4 - assert isinstance(loaded_model.lm_head, QuantLinear) == lm_head + assert not any(isinstance(m, torch.nn.Linear) and not _is_quant(m) for m in loaded_model.model.layers.modules()) + assert _is_quant(loaded_model.model.layers[0].self_attn.o_proj) + assert _bits(loaded_model.model.layers[0].self_attn.o_proj) == 8 + assert _bits(loaded_model.model.layers[0].mlp.down_proj) == 4 + assert _is_quant(loaded_model.lm_head) == lm_head assert isinstance(loaded_model.model.embed_tokens, torch.nn.Embedding) + assert not _is_quant(loaded_model.model.embed_tokens) # compose another kquant pass to also quantize embeds and lm_head p2 = create_pass_from_dict( @@ -104,7 +286,7 @@ def test_kquant(tmp_path: Path, group_size: int, sym: bool, lm_head: bool): assert isinstance(out2, HfModelHandler) loaded_model_2 = out2.load_model() - assert isinstance(loaded_model_2.model.embed_tokens, QuantEmbedding) - assert loaded_model_2.model.embed_tokens.quantizer.bits == 8 - assert isinstance(loaded_model_2.lm_head, QuantLinear) - assert loaded_model_2.lm_head.quantizer.bits == 4 if lm_head else 8 + assert _is_quant(loaded_model_2.model.embed_tokens) + assert _bits(loaded_model_2.model.embed_tokens) == 8 + assert _is_quant(loaded_model_2.lm_head) + assert _bits(loaded_model_2.lm_head) == (4 if lm_head else 8) diff --git a/test/passes/pytorch/test_moe_support.py b/test/passes/pytorch/test_moe_support.py new file mode 100644 index 0000000000..2252a05132 --- /dev/null +++ b/test/passes/pytorch/test_moe_support.py @@ -0,0 +1,92 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for shared fused-MoE layout safety checks.""" + +import pytest +import torch + +from olive.passes.pytorch.moe_support import MoeSupportError, check_moe_layout_support + + +class _FakeExperts(torch.nn.Module): + def __init__(self, **flags): + super().__init__() + for key, value in flags.items(): + setattr(self, key, value) + + +def make_fake_experts(class_name: str = "FakeExperts", **flags) -> torch.nn.Module: + """Build a stand-in experts module whose class name is ``class_name``.""" + return type(class_name, (_FakeExperts,), {})(**flags) + + +@pytest.mark.parametrize("model_type", ["mixtral", "qwen3_moe", "some_future_moe_architecture"]) +def test_check_support_accepts_any_model_type_when_not_transposed(model_type: str): + """``is_transposed=False`` is accepted for any model_type, not just a fixed allow-list.""" + check_moe_layout_support( + [make_fake_experts(is_transposed=False)], + model_type=model_type, + operation="Test MoE quantization", + ) + + +def test_check_support_rejects_missing_capability(): + with pytest.raises(MoeSupportError, match="is_transposed") as exc_info: + check_moe_layout_support([make_fake_experts()], model_type="qwen3_moe", operation="Test MoE quantization") + assert "moe=False" in str(exc_info.value) + + +@pytest.mark.parametrize("bad_value", [None, 0, 1, "False", "", [], torch.tensor(False)]) +def test_check_support_rejects_non_bool_is_transposed(bad_value): + """Non-boolean values (e.g. a stray ``None``/tensor/string) must not be treated as falsy-safe.""" + with pytest.raises(MoeSupportError, match="is_transposed") as exc_info: + check_moe_layout_support( + [make_fake_experts(is_transposed=bad_value)], model_type="qwen3_moe", operation="Test MoE quantization" + ) + assert "moe=False" in str(exc_info.value) + + +def test_check_support_rejects_transposed_layout(): + with pytest.raises(MoeSupportError, match=r"\(E, K, OUT\)") as exc_info: + check_moe_layout_support( + [make_fake_experts(is_transposed=True)], + model_type="gpt_oss", + operation="Test MoE quantization", + ) + message = str(exc_info.value) + assert "moe=False" in message + assert "gpt_oss" in message + + +def test_check_support_exempts_module_list_experts_with_no_direct_3d_param(): + """Classic per-expert ``nn.ModuleList`` experts are exempt from the layout check. + + E.g. PhiMoE/Mixtral on transformers releases without the fused-experts refactor are + unconditionally K-last (each child is a plain 2D ``nn.Linear``), so they must not be + rejected for a missing ``is_transposed`` attribute. + """ + module_list = torch.nn.ModuleList([torch.nn.Linear(4, 4) for _ in range(2)]) + check_moe_layout_support([module_list], model_type="any_model_type", operation="Test MoE quantization") + + +def test_check_support_rejects_module_list_with_direct_3d_param(): + """A ``ModuleList`` that also owns a direct 3D parameter is not exempt. + + Olive's quantizer selection would group that parameter's last dimension directly, so it + needs the same ``is_transposed`` guarantee as any other fused-experts module -- and a + bare ``ModuleList`` has no such attribute to trust. + """ + module_list = torch.nn.ModuleList([torch.nn.Linear(4, 4)]) + module_list.fused_weight = torch.nn.Parameter(torch.zeros(2, 4, 4)) + with pytest.raises(MoeSupportError, match="is_transposed"): + check_moe_layout_support([module_list], model_type="any_model_type", operation="Test MoE quantization") + + +def test_check_support_checks_every_module_in_a_mixed_list(): + """A single transposed module among otherwise-safe ones still triggers rejection.""" + safe = make_fake_experts("SafeExperts", is_transposed=False) + unsafe = make_fake_experts("UnsafeExperts", is_transposed=True) + with pytest.raises(MoeSupportError, match="UnsafeExperts"): + check_moe_layout_support([safe, unsafe], model_type="mixtral", operation="Test MoE quantization") diff --git a/test/passes/pytorch/test_quant_utils.py b/test/passes/pytorch/test_quant_utils.py index 6c3f8a79bf..9e49fcb240 100644 --- a/test/passes/pytorch/test_quant_utils.py +++ b/test/passes/pytorch/test_quant_utils.py @@ -2,16 +2,20 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +# pylint: disable=protected-access,redefined-outer-name,not-callable import logging from copy import deepcopy from types import SimpleNamespace import pytest import torch +from torch import nn from transformers import LlamaConfig, LlamaForCausalLM from olive.common.hf.wrapper import ModelWrapper from olive.common.quant.hf_utils import OliveHfQuantizationConfig +from olive.common.quant.state_dict import install_quant_tensor_param +from olive.common.quant.tensor import QuantTensor from olive.constants import PrecisionBits from olive.model import HfModelHandler from olive.passes.pytorch import quant_utils as quant_utils_module @@ -66,6 +70,54 @@ def fake(self, exclude_load_keys=None): monkeypatch.setattr(HfModelHandler, "get_hf_model_config", fake) +def test_get_layer_inputs_cleans_up_after_forward_error(monkeypatch): + model = LlamaForCausalLM( + LlamaConfig( # pylint: disable=unexpected-keyword-arg + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + vocab_size=32, + ) + ) + wrapper = ModelWrapper.from_model(model) + first_layer = wrapper.get_layers(return_name=False)[0] + pre_layer_modules = list(wrapper.get_embeds(return_name=False)) + if rotary_embed := wrapper.get_rotary_embed(return_name=False): + pre_layer_modules.append(rotary_embed) + + to_calls = {id(module): [] for module in pre_layer_modules} + for module in pre_layer_modules: + monkeypatch.setattr( + module, + "to", + lambda device, current=module: to_calls[id(current)].append(device) or current, + ) + + monkeypatch.setattr( + quant_utils_module, + "get_calibration_dataset", + lambda *_args, **_kwargs: [{"input_ids": torch.ones((1, 2), dtype=torch.long)}], + ) + + def failing_forward(**_kwargs): + raise RuntimeError("calibration forward failed") + + monkeypatch.setattr(wrapper.model, "forward", failing_forward) + + with pytest.raises(RuntimeError, match="calibration forward failed"): + quant_utils_module.get_layer_inputs_for_calibration( + SimpleNamespace(), + wrapper, + data_config=None, + device="cpu", + ) + + assert not first_layer._forward_pre_hooks + assert all(calls == ["cpu", "cpu"] for calls in to_calls.values()) + + # --------------------------------------------------------------------------- # _quant_config_rank # --------------------------------------------------------------------------- @@ -196,9 +248,9 @@ def test_prepare_model_no_existing_quant_config_no_overrides_quantizes_all_linea for name, module in wrapper.model.named_modules(): if isinstance(module, torch.nn.Linear): if name == "lm_head": - assert not hasattr(module, "quant_info") + assert not hasattr(module.weight, "quant_info") else: - assert module.quant_info.quantizer.bits == PrecisionBits.BITS4 + assert module.weight.quant_info.quantizer.bits == PrecisionBits.BITS4 assert qcfg.overrides == {} assert qcfg.bits == PrecisionBits.BITS4 assert eligible is False @@ -236,7 +288,7 @@ def test_prepare_model_promotes_user_override_conflicts_for_qkv(input_model): for proj in ("q_proj", "k_proj", "v_proj"): assert qcfg.get_qlinear_init_args(f"model.layers.0.self_attn.{proj}") == expected attached = getattr(wrapper.model.model.layers[0].self_attn, proj) - assert attached.quant_info.quantizer.bits == PrecisionBits.BITS8 + assert attached.weight.quant_info.quantizer.bits == PrecisionBits.BITS8 def test_prepare_model_attaches_quant_info_matching_final_post_normalize_config(input_model): @@ -251,7 +303,7 @@ def test_prepare_model_attaches_quant_info_matching_final_post_normalize_config( attn = wrapper.model.model.layers[0].self_attn for proj in ("q_proj", "k_proj", "v_proj"): - attached_bits = getattr(attn, proj).quant_info.quantizer.bits + attached_bits = getattr(attn, proj).weight.quant_info.quantizer.bits cfg_bits = qcfg.get_qlinear_init_args(f"model.layers.0.self_attn.{proj}")["bits"] assert attached_bits == cfg_bits == PrecisionBits.BITS8 @@ -279,7 +331,7 @@ def test_prepare_model_drops_override_for_lm_head_when_lm_head_disabled(input_mo wrapper, qcfg, _ = prepare_model(input_model, config) assert "lm_head" not in (qcfg.overrides or {}) - assert not hasattr(wrapper.model.lm_head, "quant_info") + assert not hasattr(wrapper.model.lm_head.weight, "quant_info") def test_prepare_model_drops_embedding_override_when_embeds_disabled(input_model): @@ -321,15 +373,15 @@ def test_prepare_model_drops_qkv_overrides_for_modules_excluded_via_exclude_attn wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config(), exclude_attn_inputs=True) attention = wrapper.model.model.layers[0].self_attn - assert not hasattr(attention.q_proj, "quant_info") - assert not hasattr(attention.k_proj, "quant_info") + assert not hasattr(attention.q_proj.weight, "quant_info") + assert not hasattr(attention.k_proj.weight, "quant_info") # V is quantized and promoted to the group-wide 8-bit config. assert qcfg.get_qlinear_init_args("model.layers.0.self_attn.v_proj") == { "bits": PrecisionBits.BITS8, "symmetric": True, "group_size": 16, } - assert attention.v_proj.quant_info.quantizer.bits == PrecisionBits.BITS8 + assert attention.v_proj.weight.quant_info.quantizer.bits == PrecisionBits.BITS8 # Q/K overrides dropped; follow-up pass will rebuild them from V (locked). assert "model.layers.0.self_attn.q_proj" not in (qcfg.overrides or {}) assert "model.layers.0.self_attn.k_proj" not in (qcfg.overrides or {}) @@ -341,7 +393,7 @@ def test_prepare_model_exclude_attn_inputs_fused_qkv_does_not_create_overrides_f wrapper, qcfg, _ = prepare_model(model_handler, _baseline_pass_config(), exclude_attn_inputs=True) qkv_proj = wrapper.model.model.layers[0].self_attn.qkv_proj - assert not hasattr(qkv_proj, "quant_info") + assert not hasattr(qkv_proj.weight, "quant_info") assert "model.layers.0.self_attn.qkv_proj" not in (qcfg.overrides or {}) @@ -494,8 +546,6 @@ def test_prepare_model_locks_default_quantized_qkv_member_without_override(input for V -- that would disagree with V's on-disk weights. Instead, Q/K should be demoted to V's existing default config. """ - from olive.common.quant.nn import QuantLinear - qu = quant_utils_module existing = { @@ -513,19 +563,11 @@ def test_prepare_model_locks_default_quantized_qkv_member_without_override(input def fake_load(model_handler, **kwargs): loaded = real_loader(model_handler, **kwargs) - # Replace v_proj of layer 0 with a QuantLinear so it looks already-quantized on disk. - attn = loaded.model.layers[0].self_attn - v = attn.v_proj - attn.v_proj = QuantLinear( - in_features=v.in_features, - out_features=v.out_features, - bias=v.bias is not None, - bits=PrecisionBits.BITS4, - symmetric=False, - group_size=16, - device=v.weight.device, - dtype=v.weight.dtype, - ) + # Make v_proj of layer 0 look already-quantized on disk by swapping its weight + # for a QuantTensor at the existing config's defaults (no override entry). + v = loaded.model.layers[0].self_attn.v_proj + qt = QuantTensor.from_float(v.weight.detach(), bits=4, symmetric=False, group_size=16) + install_quant_tensor_param(v, "weight", qt) return loaded monkeypatch.setattr(qu, "load_hf_base_model", fake_load) @@ -548,3 +590,104 @@ def fake_load(model_handler, **kwargs): assert qcfg.get_qlinear_init_args(f"model.layers.0.self_attn.{proj}") == default # No new override added for V (it stays at defaults on disk). assert "model.layers.0.self_attn.v_proj" not in (qcfg.overrides or {}) + + +# --------------------------------------------------------------------------- +# State-dict helper tests (install_quant_tensor_param, 2D + 3D fused MoE) +# --------------------------------------------------------------------------- + + +class _ExpertsBlock(nn.Module): + """Fake fused-3D experts module (gpt-oss / Qwen3-MoE style).""" + + def __init__(self, num_experts: int = 4, out_features: int = 32, in_features: int = 16): + super().__init__() + self.gate_up_proj = nn.Parameter( + torch.randn(num_experts, out_features, in_features, dtype=torch.float32), + requires_grad=False, + ) + + +class TestInstallQuantTensorParam: + def test_install_3d_quant_tensor(self): + block = _ExpertsBlock() + qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=True, group_size=16) + + install_quant_tensor_param(block, "gate_up_proj", qt) + + # Parameter is a QuantTensor (tensor subclass parameters return + # the subclass instance directly from nn.Parameter.__new__). + param = block._parameters["gate_up_proj"] + assert isinstance(param, QuantTensor) + # Sibling buffers are registered and share storage with the QuantTensor. + buffers = dict(block.named_buffers()) + assert "gate_up_proj_qweight" in buffers + assert "gate_up_proj_scales" in buffers + assert "gate_up_proj_qzeros" not in buffers # symmetric → no qzeros + assert buffers["gate_up_proj_qweight"] is param.qweight + assert buffers["gate_up_proj_scales"] is param.scales + + def test_install_asymmetric_emits_qzeros(self): + block = _ExpertsBlock() + qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=False, group_size=16) + + install_quant_tensor_param(block, "gate_up_proj", qt) + + buffers = dict(block.named_buffers()) + assert "gate_up_proj_qzeros" in buffers + assert buffers["gate_up_proj_qzeros"] is block._parameters["gate_up_proj"].qzeros + + def test_state_dict_drops_quant_tensor_entry(self): + """After install, state_dict must contain only plain Tensors (no QuantTensor entry).""" + block = _ExpertsBlock() + qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=True, group_size=16) + + install_quant_tensor_param(block, "gate_up_proj", qt) + + sd = block.state_dict() + # No QuantTensor instance should appear in the state_dict. + for key, value in sd.items(): + assert not isinstance(value, QuantTensor), f"{key} should not be a QuantTensor" + # The plain ``gate_up_proj`` key (the QuantTensor parameter) is dropped; + # the buffers carry the on-disk representation. + assert "gate_up_proj" not in sd + assert "gate_up_proj_qweight" in sd + assert "gate_up_proj_scales" in sd + + def test_install_on_linear_module(self): + """Smoke test on a normal nn.Linear so that F.linear forward still works.""" + linear = nn.Linear(16, 32, bias=False) + weight = linear.weight.detach().clone() + qt = QuantTensor.from_float(weight, bits=4, symmetric=True, group_size=16) + + install_quant_tensor_param(linear, "weight", qt) + + assert isinstance(linear.weight, QuantTensor) + # forward dispatches through QuantTensor.__torch_function__ and returns a plain Tensor + x = torch.randn(2, 16) + y = linear(x) + assert y.shape == (2, 32) + assert not isinstance(y, QuantTensor) + + +def test_module_weight_has_quant_info_only_for_marked_params(): + """Regression: discovery must not pick up LayerNorm / Conv2d weights. + + GPTQ / AutoClip discover quantizable layers via + ``_module_weight_has_quant_info``. Modules that happen to expose a + ``weight`` attribute but never had ``quant_info`` stamped on it must + be left alone. + """ + from olive.common.quant.utils import WeightQuantizer + from olive.passes.pytorch.quant_utils import QuantInfo, _module_weight_has_quant_info + + ln = nn.LayerNorm(16) + conv = nn.Conv2d(3, 8, kernel_size=3) + linear_unmarked = nn.Linear(16, 32, bias=False) + linear_marked = nn.Linear(16, 32, bias=False) + linear_marked.weight.quant_info = QuantInfo(quantizer=WeightQuantizer(bits=4, symmetric=True, group_size=16)) + + assert not _module_weight_has_quant_info(ln) + assert not _module_weight_has_quant_info(conv) + assert not _module_weight_has_quant_info(linear_unmarked) + assert _module_weight_has_quant_info(linear_marked) diff --git a/test/passes/pytorch/test_rtn.py b/test/passes/pytorch/test_rtn.py index 934978ca58..82d52b19e1 100644 --- a/test/passes/pytorch/test_rtn.py +++ b/test/passes/pytorch/test_rtn.py @@ -2,20 +2,177 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +# pylint: disable=protected-access from pathlib import Path import pytest import torch from olive.common.quant.hf_utils import OliveHfQuantizationConfig -from olive.common.quant.nn import QuantEmbedding, QuantLinear +from olive.common.quant.tensor import QuantTensor from olive.hardware.accelerator import AcceleratorSpec, Device from olive.model import HfModelHandler from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.pytorch.gptq import Gptq +from olive.passes.pytorch.moe_support import MoeSupportError +from olive.passes.pytorch.quant_utils import prepare_model from olive.passes.pytorch.rtn import Rtn from test.utils import get_tiny_phi3 +def _save_trivial_tokenizer(save_path: Path, vocab_size: int) -> None: + """Save a local tokenizer so pass metadata serialization never needs the hub.""" + from tokenizers import Tokenizer, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast + + tokenizer = Tokenizer(models.WordLevel({f"t{i}": i for i in range(vocab_size)}, unk_token="t0")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="t0", pad_token="t0").save_pretrained(save_path) + + +def _is_quant(module: torch.nn.Module) -> bool: + if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)): + return False + weight = module._parameters.get("weight") + return weight is not None and isinstance(weight.data, QuantTensor) + + +def _bits(module: torch.nn.Module) -> int: + return module.weight.data.bits + + +def _make_local_tiny_mixtral(save_path): + """Save a local copy of ``yujiepan/mixtral-tiny-random`` forced to the ``eager`` experts implementation. + + The hub model defaults to ``grouped_mm`` for its MoE forward, which hits a CUDA kernel + stride-alignment limitation on this architecture's tiny (non-16-byte-aligned) hidden + dims and is unrelated to Olive's quantization code; forcing ``eager`` avoids it while + keeping this a real HF model / real Gptq+Rtn pass run. + """ + import json + + from transformers import AutoModelForCausalLM, AutoTokenizer + + save_path = Path(save_path) + save_path.mkdir(parents=True, exist_ok=True) + model = AutoModelForCausalLM.from_pretrained("yujiepan/mixtral-tiny-random") + model.save_pretrained(save_path) + AutoTokenizer.from_pretrained("yujiepan/mixtral-tiny-random").save_pretrained(save_path) + + config_path = save_path / "config.json" + config = json.loads(config_path.read_text()) + config["experts_implementation"] = "eager" + config_path.write_text(json.dumps(config, indent=2)) + + return HfModelHandler(model_path=str(save_path)) + + +def _make_local_tiny_qwen3_moe(save_path) -> HfModelHandler: + """Save a tiny K-last fused-experts model without downloading a checkpoint.""" + from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM + + torch.manual_seed(0) + save_path = Path(save_path) + save_path.mkdir(parents=True, exist_ok=True) + config = Qwen3MoeConfig( # pylint: disable=unexpected-keyword-arg + vocab_size=32, + hidden_size=16, + intermediate_size=16, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_experts=2, + num_experts_per_tok=1, + decoder_sparse_step=1, + head_dim=8, + experts_implementation="eager", + ) + Qwen3MoeForCausalLM(config).save_pretrained(save_path) + _save_trivial_tokenizer(save_path, config.vocab_size) + return HfModelHandler(model_path=str(save_path)) + + +def test_rtn_moe_refuses_transposed_layout_before_finalize(tmp_path: Path, monkeypatch): + """A transposed-layout (``is_transposed=True``) experts module fails before finalize. + + The rejection must happen before quantization or output serialization. + """ + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + + def patched_prepare_model(*args, **kwargs): + wrapper, qcfg, retie = prepare_model(*args, **kwargs) + for layer in wrapper.get_layer_wrappers(): + experts = layer.get_experts(return_name=False) + if experts is not None: + experts.is_transposed = True + return wrapper, qcfg, retie + + def unexpected_finalize(*args, **kwargs): + pytest.fail("finalize must not run after MoE layout support is rejected") + + monkeypatch.setattr("olive.passes.pytorch.rtn.prepare_model", patched_prepare_model) + monkeypatch.setattr("olive.passes.pytorch.rtn.finalize", unexpected_finalize) + + quantizer = create_pass_from_dict(Rtn, {"moe": True}, disable_search=True) + output_path = tmp_path / "rtn" + with pytest.raises(MoeSupportError, match=r"\(E, K, OUT\).*moe=False"): + quantizer.run(input_model, str(output_path)) + assert not output_path.exists() + + +def test_rtn_moe_false_does_not_run_layout_gate(tmp_path: Path, monkeypatch): + """The default dense path must not invoke fused-experts layout validation.""" + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + + def unexpected_gate(*args, **kwargs): + pytest.fail("MoE layout support check must not run when moe=False") + + monkeypatch.setattr("olive.passes.pytorch.rtn.check_moe_layout_support", unexpected_gate) + quantizer = create_pass_from_dict(Rtn, {"moe": False, "group_size": -1}, disable_search=True) + out = quantizer.run(input_model, str(tmp_path / "rtn")) + + loaded = out.load_model() + experts = loaded.model.layers[0].mlp.experts + assert not any(isinstance(param.data, QuantTensor) for param in experts.parameters()) + assert isinstance(loaded.model.layers[0].self_attn.q_proj.weight.data, QuantTensor) + + +def test_rtn_moe_gate_ignores_prior_checkpoint_moe_flag(tmp_path: Path, monkeypatch): + """Regression test: a second RTN pass with moe=False must not re-run the layout gate. + + ``prepare_model`` ORs a pre-existing checkpoint's ``moe`` flag into the merged + ``qcfg.moe`` (see ``quant_utils.prepare_model``), so gating on ``qcfg.moe`` would make + this second, moe=False invocation incorrectly re-run fused-experts layout validation + -- something this run never asked for. The gate must key off this invocation's own + ``config.moe`` request instead. + """ + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + first_pass = create_pass_from_dict(Rtn, {"moe": True, "group_size": -1}, disable_search=True) + quantized = first_pass.run(input_model, str(tmp_path / "rtn_first")) + + def unexpected_gate(*args, **kwargs): + pytest.fail("MoE layout support check must not re-run when this invocation requests moe=False") + + monkeypatch.setattr("olive.passes.pytorch.rtn.check_moe_layout_support", unexpected_gate) + second_pass = create_pass_from_dict(Rtn, {"moe": False, "lm_head": True, "group_size": -1}, disable_search=True) + out = second_pass.run(quantized, str(tmp_path / "rtn_second")) + + loaded = out.load_model() + assert isinstance(loaded.lm_head.weight.data, QuantTensor) + + +def test_rtn_moe_k_last_layout_quantizes_experts(tmp_path: Path): + """A K-last (``is_transposed=False``) fused-experts model quantizes successfully end to end.""" + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + quantizer = create_pass_from_dict(Rtn, {"moe": True, "group_size": -1}, disable_search=True) + out = quantizer.run(input_model, str(tmp_path / "rtn")) + + loaded = out.load_model() + experts = loaded.model.layers[0].mlp.experts + assert any(isinstance(param.data, QuantTensor) for param in experts.parameters()) + + @pytest.mark.parametrize("group_size", [-1, 16]) @pytest.mark.parametrize("sym", [True, False]) @pytest.mark.parametrize("lm_head", [True, False]) @@ -48,13 +205,14 @@ def test_gptq(tmp_path: Path, group_size: int, sym: bool, lm_head: bool): assert hasattr(loaded_model.config, "quantization_config") assert isinstance(loaded_model.config.quantization_config, OliveHfQuantizationConfig) assert loaded_model.config.quantization_config.group_size == group_size - assert not any(isinstance(m, torch.nn.Linear) for m in loaded_model.model.layers.modules()) - assert isinstance(loaded_model.model.layers[0].self_attn.o_proj, QuantLinear) - assert loaded_model.model.layers[0].self_attn.o_proj.quantizer.bits == 8 - assert loaded_model.model.layers[0].mlp.down_proj.quantizer.bits == 4 + assert not any(isinstance(m, torch.nn.Linear) and not _is_quant(m) for m in loaded_model.model.layers.modules()) + assert _is_quant(loaded_model.model.layers[0].self_attn.o_proj) + assert _bits(loaded_model.model.layers[0].self_attn.o_proj) == 8 + assert _bits(loaded_model.model.layers[0].mlp.down_proj) == 4 assert loaded_model.config.quantization_config.lm_head == lm_head - assert isinstance(loaded_model.lm_head, QuantLinear) == lm_head + assert _is_quant(loaded_model.lm_head) == lm_head assert isinstance(loaded_model.model.embed_tokens, torch.nn.Embedding) + assert not _is_quant(loaded_model.model.embed_tokens) # compose another rtn pass on top of the partially quantized model p2 = create_pass_from_dict( @@ -76,8 +234,270 @@ def test_gptq(tmp_path: Path, group_size: int, sym: bool, lm_head: bool): assert isinstance(out2, HfModelHandler) loaded_model_2 = out2.load_model() # check that the embed tokens layer is quantized to 8 bits - assert isinstance(loaded_model_2.model.embed_tokens, QuantEmbedding) - assert loaded_model_2.model.embed_tokens.quantizer.bits == 8 + assert _is_quant(loaded_model_2.model.embed_tokens) + assert _bits(loaded_model_2.model.embed_tokens) == 8 # check that the lm head is quantized to 8 bits if it was not quantized before - assert isinstance(loaded_model_2.lm_head, QuantLinear) - assert loaded_model_2.lm_head.quantizer.bits == 4 if lm_head else 8 + assert _is_quant(loaded_model_2.lm_head) + assert _bits(loaded_model_2.lm_head) == (4 if lm_head else 8) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Gptq requires a CUDA-capable GPU") +def test_gptq_then_rtn_moe_e2e(tmp_path: Path): + """M5: real end-to-end ``Gptq`` -> ``Rtn(moe=True, embeds=True)`` composition. + + Runs the actual ``Gptq`` pass (calibration-based) on a small, real MoE model + (``yujiepan/mixtral-tiny-random``), then runs the actual ``Rtn`` pass with + ``moe=True, embeds=True`` on top. Verifies that: + + * the ``nn.Linear`` attention/MLP-router weights that GPTQ already quantized are + *not* re-quantized by RTN (still the original GPTQ ``QuantTensor``), and + * the fused-3D MoE expert weights and the input embeddings -- which GPTQ does not + touch -- get RTN-quantized. + * a save -> reload round-trip preserves both sets of quantized weights. + """ + input_model = _make_local_tiny_mixtral(tmp_path / "tiny_mixtral") + + gptq_pass = create_pass_from_dict( + Gptq, + {"group_size": -1, "lm_head": False}, + disable_search=True, + accelerator_spec=AcceleratorSpec(accelerator_type=Device.GPU, execution_provider="CUDAExecutionProvider"), + ) + gptq_out_folder = str(tmp_path / "gptq") + gptq_out = gptq_pass.run(input_model, gptq_out_folder) + assert isinstance(gptq_out, HfModelHandler) + + gptq_loaded = gptq_out.load_model() + # GPTQ quantizes attention/MLP nn.Linear weights, not the MoE fused-3D expert params. + assert _is_quant(gptq_loaded.model.layers[0].self_attn.q_proj) + router_gate = gptq_loaded.model.layers[0].mlp.gate + assert not _is_quant(router_gate) # router (MixtralTopKRouter, not nn.Linear) stays full precision + experts = gptq_loaded.model.layers[0].mlp.experts + assert not any(isinstance(p.data, QuantTensor) for p in experts.parameters()) + assert not _is_quant(gptq_loaded.model.embed_tokens) + # snapshot the GPTQ-quantized q_proj weight for later comparison + gptq_qproj_before = gptq_loaded.model.layers[0].self_attn.q_proj.weight.data.to_dense().clone() + gptq_qproj_bits = _bits(gptq_loaded.model.layers[0].self_attn.q_proj) + + rtn_pass = create_pass_from_dict( + Rtn, + {"bits": 4, "group_size": -1, "moe": True, "embeds": True, "lm_head": True}, + disable_search=True, + accelerator_spec=AcceleratorSpec(accelerator_type=Device.GPU, execution_provider="CUDAExecutionProvider"), + ) + rtn_out_folder = str(tmp_path / "rtn") + rtn_out = rtn_pass.run(gptq_out, rtn_out_folder) + assert isinstance(rtn_out, HfModelHandler) + + rtn_loaded = rtn_out.load_model() + + # The GPTQ-quantized Linear is untouched by RTN (same bits, same dequantized values). + q_proj = rtn_loaded.model.layers[0].self_attn.q_proj + assert _is_quant(q_proj) + assert _bits(q_proj) == gptq_qproj_bits + torch.testing.assert_close(q_proj.weight.data.to_dense(), gptq_qproj_before, rtol=0, atol=0) + + # MoE experts and embeddings, which GPTQ left untouched, are now RTN-quantized. + rtn_experts = rtn_loaded.model.layers[0].mlp.experts + assert any(isinstance(p.data, QuantTensor) for p in rtn_experts.parameters()) + assert _is_quant(rtn_loaded.model.embed_tokens) + assert _bits(rtn_loaded.model.embed_tokens) == 4 + + # Save -> reload round-trip preserves both GPTQ and RTN quantization. + reload_folder = str(tmp_path / "reload") + rtn_loaded.save_pretrained(reload_folder) + reloaded_handler = HfModelHandler(model_path=reload_folder) + reloaded = reloaded_handler.load_model() + + reloaded_q_proj = reloaded.model.layers[0].self_attn.q_proj + assert _is_quant(reloaded_q_proj) + assert _bits(reloaded_q_proj) == gptq_qproj_bits + torch.testing.assert_close(reloaded_q_proj.weight.data.to_dense(), gptq_qproj_before, rtol=0, atol=0) + + reloaded_experts = reloaded.model.layers[0].mlp.experts + assert any(isinstance(p.data, QuantTensor) for p in reloaded_experts.parameters()) + assert _is_quant(reloaded.model.embed_tokens) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Rtn(moe=True) real forward needs a CUDA-capable GPU") +def test_rtn_moe_real_forward_after_reload(tmp_path: Path): + """Regression test for a save/reload round-trip bug in fused-3D MoE quantized weights. + + ``transformers``'s ``save_pretrained`` defaults to ``save_original_format=True``, + which for Mixtral-family MoE architectures round-trips the on-disk state dict + through a *legacy* per-expert ``nn.Linear``-shaped layout (splitting the fused-3D + ``experts.gate_up_proj``/``down_proj`` into ``experts.{i}.w1/w2/w3.weight`` and back + on load). That reshape machinery assumes plain float weight tensors and silently + drops the trailing group-size dimension of our quantized ``_scales``/``_qzeros`` + buffers, which crashes real ``forward()`` calls on the reloaded model (previously + undetected because no existing test called ``forward()`` on a save/reload round trip). + + This test quantizes a real MoE model with ``Rtn(moe=True)``, saves via the actual + pass output (exercising ``finalize()``'s ``save_pretrained`` call), reloads from + disk, and calls the model's real ``forward()`` -- asserting no crash, no ``NaN``, + and that the fused-3D scales buffer keeps its group-size dimension. + """ + input_model = _make_local_tiny_mixtral(tmp_path / "tiny_mixtral") + + rtn_pass = create_pass_from_dict( + Rtn, + {"bits": 4, "group_size": -1, "moe": True, "embeds": True, "lm_head": True}, + disable_search=True, + accelerator_spec=AcceleratorSpec(accelerator_type=Device.GPU, execution_provider="CUDAExecutionProvider"), + ) + rtn_out_folder = str(tmp_path / "rtn") + rtn_out = rtn_pass.run(input_model, rtn_out_folder) + assert isinstance(rtn_out, HfModelHandler) + + # ``finalize()`` re-serializes config.json and does not preserve custom, + # non-quantization-related fields; re-patch ``experts_implementation`` so the + # reloaded model uses the ``eager`` MoE forward (see ``_make_local_tiny_mixtral``). + import json + + config_path = Path(rtn_out_folder) / "config.json" + config = json.loads(config_path.read_text()) + config["experts_implementation"] = "eager" + config_path.write_text(json.dumps(config, indent=2)) + + reloaded_handler = HfModelHandler(model_path=rtn_out_folder) + reloaded = reloaded_handler.load_model() + + # Guard the exact regression: the fused-3D expert weight's scales/qzeros must keep + # their trailing group-size dimension (num_experts, out_features, num_groups) after + # the disk round trip, instead of collapsing to a 2D (num_experts, out_features). + gate_up_proj = reloaded.model.layers[0].mlp.experts.gate_up_proj + assert _is_quant(reloaded.model.embed_tokens) # sanity: quantization actually applied + assert isinstance(gate_up_proj.data, QuantTensor) + scales_shape = gate_up_proj.data.scales.shape + assert len(scales_shape) == 3, f"expected fused-3D scales, got shape {scales_shape}" + assert scales_shape[:2] == gate_up_proj.data.qweight.shape[:2] + + reloaded = reloaded.cuda().eval() + input_ids = torch.randint(0, 100, (1, 8), device="cuda") + with torch.no_grad(): + out = reloaded(input_ids) + assert not torch.isnan(out.logits).any() + assert not torch.isinf(out.logits).any() + + +def _make_local_tiny_tied_llama(save_path) -> HfModelHandler: + """Save a tiny ``LlamaForCausalLM`` with ``tie_word_embeddings=True`` (built locally, no hub access).""" + from transformers import LlamaConfig, LlamaForCausalLM + + torch.manual_seed(0) + save_path = Path(save_path) + save_path.mkdir(parents=True, exist_ok=True) + config = LlamaConfig( # pylint: disable=unexpected-keyword-arg + vocab_size=32, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + tie_word_embeddings=True, + ) + LlamaForCausalLM(config).save_pretrained(save_path) + _save_trivial_tokenizer(save_path, config.vocab_size) + return HfModelHandler(model_path=str(save_path)) + + +def _load_quant_tensor_from_disk(model_dir, pname: str, sym: bool, group_size: int) -> QuantTensor: + """Rebuild the ``QuantTensor`` for ``pname`` straight from the saved safetensors shard. + + Gives a device-independent, bit-exact reference for what ``from_pretrained`` must load + (recomputing the quantization in-process would not match bit-for-bit, since the pass + quantizes on GPU when one is available). + """ + from safetensors.torch import load_file + + from olive.common.quant.state_dict import buffer_names + + qname, sname, zname = buffer_names(pname) + weights: dict = {} + for shard in sorted(Path(model_dir).glob("*.safetensors")): + weights.update(load_file(shard)) + qweight, scales = weights[qname], weights[sname] + return QuantTensor.from_packed( + qweight=qweight, + scales=scales, + qzeros=weights.get(zname), + bits=4, + group_size=group_size, + symmetric=sym, + shape=(scales.shape[0], qweight.shape[-1] * 2), + dtype=scales.dtype, + ) + + +@pytest.mark.parametrize("sym", [True, False]) +def test_rtn_tied_word_embeddings_roundtrip(tmp_path: Path, sym: bool): + """Regression: tied ``lm_head`` / ``embed_tokens`` must survive a save -> reload round trip. + + ``tie_quant_word_embeddings`` installs the *same* ``QuantTensor`` object on both modules + and aliases ``lm_head``'s buffer dict entries to ``embed_tokens``'s buffer *objects* — a + one-time snapshot. HF's loader then replaces ``embed_tokens``'s buffer objects, leaving + ``lm_head``'s entries stale. ``refresh_quant_tensor_refs`` used to rebind the shared + QuantTensor once per hosting module (last-write-wins over ``named_modules()`` order), so + the stale ``lm_head`` buffers could win and silently zero/garbage the reloaded weight + while ``is_placeholder`` still read ``False``. + """ + group_size = 16 + input_model = _make_local_tiny_tied_llama(tmp_path / "tiny_llama_tied") + original_embed = input_model.load_model().model.embed_tokens.weight.detach().clone() + + p = create_pass_from_dict( + Rtn, + {"bits": 4, "group_size": group_size, "sym": sym, "lm_head": True, "embeds": True}, + disable_search=True, + ) + out = p.run(input_model, str(tmp_path / "quantized")) + assert isinstance(out, HfModelHandler) + + loaded = out.load_model() + assert loaded.config.quantization_config.tie_word_embeddings is True + + embed, lm_head = loaded.model.embed_tokens, loaded.lm_head + assert _is_quant(embed) + assert _is_quant(lm_head) + # Tying must be preserved end to end: same QuantTensor object and same buffer objects. + # NOTE: read the parameter out of ``_parameters`` -- ``param.data`` on a tensor subclass + # goes through ``detach()`` and returns a *fresh* QuantTensor whose inner tensors are + # views, so it cannot be used for identity assertions. + shared = embed._parameters["weight"] + assert shared is lm_head._parameters["weight"] + assert shared.qweight is embed._buffers["weight_qweight"] + assert shared.scales is embed._buffers["weight_scales"] + assert embed._buffers["weight_qweight"] is lm_head._buffers["weight_qweight"] + assert embed._buffers["weight_scales"] is lm_head._buffers["weight_scales"] + assert shared.is_placeholder is False + + # The reloaded weight must be bit-identical to what was written to disk -- not zeros, + # not stale placeholder data. + disk = _load_quant_tensor_from_disk(tmp_path / "quantized", "model.embed_tokens.weight", sym, group_size) + assert torch.equal(shared.qweight, disk.qweight) + assert torch.equal(shared.scales, disk.scales) + expected = disk.to_dense() + + embed_dense = shared.to_dense() + lm_head_dense = lm_head._parameters["weight"].to_dense() + assert torch.isfinite(embed_dense).all() + assert embed_dense.abs().sum() > 0 + torch.testing.assert_close(embed_dense, expected, rtol=0, atol=0) + torch.testing.assert_close(lm_head_dense, expected, rtol=0, atol=0) + # ... and it still approximates the original float weight within quantization error. + torch.testing.assert_close(embed_dense, original_embed, rtol=0, atol=float(disk.scales.max())) + + # A real forward pass on the reloaded model produces finite logits. + loaded.eval() + with torch.no_grad(): + logits = loaded(torch.randint(0, 32, (1, 8))).logits + assert torch.isfinite(logits).all() + + # And a second save -> reload round trip is stable (the alias buffers written to disk + # agree with the live ones). + resave_path = tmp_path / "resaved" + loaded.save_pretrained(resave_path, save_original_format=False) + input_model.save_metadata(str(resave_path)) + resaved = HfModelHandler(model_path=str(resave_path)).load_model() + torch.testing.assert_close(resaved.model.embed_tokens._parameters["weight"].to_dense(), expected, rtol=0, atol=0) + torch.testing.assert_close(resaved.lm_head._parameters["weight"].to_dense(), expected, rtol=0, atol=0) diff --git a/test/passes/pytorch/test_train_utils.py b/test/passes/pytorch/test_train_utils.py new file mode 100644 index 0000000000..4e70884d47 --- /dev/null +++ b/test/passes/pytorch/test_train_utils.py @@ -0,0 +1,29 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import inspect + +import pytest + +from olive.passes.pytorch.train_utils import get_calibration_data_config, get_calibration_dataset + + +@pytest.mark.parametrize("func", [get_calibration_dataset, get_calibration_data_config]) +def test_calibration_helpers_default_split_is_full_train(func): + # execute + default_split = inspect.signature(func).parameters["split"].default + + # assert + # measured with the gpt2 tokenizer: wikitext-2-raw-v1 'train[:1000]' is only 61,816 tokens (353 of the + # first 1000 rows are blank lines or headers) which is ~30 blocks of 2048 tokens, silently + # under-delivering the requested max_samples. The full 'train' split gives ~1,167 blocks. + assert default_split == "train" + + +def test_get_calibration_data_config_uses_default_split(): + # execute + data_config = get_calibration_data_config("dummy-model") + + # assert + assert data_config.load_dataset_config.params["split"] == "train" diff --git a/test/passes/qairt/conftest.py b/test/passes/qairt/conftest.py index f49524a323..b921d45206 100644 --- a/test/passes/qairt/conftest.py +++ b/test/passes/qairt/conftest.py @@ -13,7 +13,7 @@ def mock_container_fixture(): """Provide a pre-configured LLMContainer mock with input/output metadata and an export stub. - Tests are responsible for wiring LLMContainer.load.return_value so they can customise + Tests are responsible for wiring LLMContainer.load.return_value so they can customize container attributes before the pass runs. """ container = MagicMock() diff --git a/test/requirements-test.txt b/test/requirements-test.txt index bebf4bb260..90e57b3ea6 100644 --- a/test/requirements-test.txt +++ b/test/requirements-test.txt @@ -37,4 +37,4 @@ sentencepiece soundfile tabulate torchvision -transformers<5.4.0 # transformers 5.4.0 breaks CI. Need to investigate and update our test code +transformers>5.3.0 diff --git a/test/systems/docker/test_docker_system.py b/test/systems/docker/test_docker_system.py index 5430b68587..900fd42971 100644 --- a/test/systems/docker/test_docker_system.py +++ b/test/systems/docker/test_docker_system.py @@ -2,6 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +import json from unittest.mock import MagicMock, patch import pytest @@ -10,7 +11,7 @@ from olive.systems.system_config import DockerTargetUserConfig from test.utils import ONNX_MODEL_PATH -# pylint: disable=attribute-defined-outside-init,protected-access +# pylint: disable=attribute-defined-outside-init,duplicate-code,protected-access class TestDockerSystem: @@ -140,6 +141,68 @@ def test_run_workflow(self, mock_find_resources, mock_tempdir, mock_from_env, tm # Verify cleanup mock_container.remove.assert_called_once() + @patch("olive.systems.docker.docker_system.docker.from_env") + def test_prepare_environment_forwards_ci_to_workflow_container(self, mock_from_env, monkeypatch): + mock_docker_client = MagicMock() + mock_from_env.return_value = mock_docker_client + mock_docker_client.images.get.return_value = MagicMock() + monkeypatch.setenv("TF_BUILD", "True") + docker_config = self.get_default_docker_config() + docker_system = DockerSystem( + image_name=docker_config.image_name, + build_context_path=docker_config.build_context_path, + dockerfile=docker_config.dockerfile, + work_dir=docker_config.work_dir, + ) + + environment = docker_system._prepare_environment({}) + + assert environment["CI"] == "1" + + @patch("olive.systems.docker.docker_system.docker.from_env") + def test_prepare_environment_forwards_full_telemetry_opt_out(self, mock_from_env): + from olive.telemetry.telemetry import Telemetry + + mock_from_env.return_value.images.get.return_value = MagicMock() + docker_config = self.get_default_docker_config() + docker_system = DockerSystem( + image_name=docker_config.image_name, + build_context_path=docker_config.build_context_path, + dockerfile=docker_config.dockerfile, + work_dir=docker_config.work_dir, + ) + + with patch.object(Telemetry, "_process_disabled", True): + environment = docker_system._prepare_environment({}) + + assert environment["OLIVE_DISABLE_TELEMETRY"] == "1" + + def test_workflow_runner_disables_inner_recipe_telemetry(self, tmp_path, monkeypatch): + from olive.systems.docker import workflow_runner + + monkeypatch.delenv("HF_TOKEN", raising=False) + config = {"input_model": {"type": "ONNXModel", "model_path": "model.onnx"}} + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(config)) + + telemetry = MagicMock() + with ( + patch.object(workflow_runner, "olive_run") as mock_olive_run, + patch.object(workflow_runner.Telemetry, "get_existing_instance", return_value=telemetry), + ): + workflow_runner.runner_entry(config_path) + + mock_olive_run.assert_called_once_with( + config, + emit_error_telemetry=False, + emit_recipe_telemetry=False, + ) + telemetry.shutdown.assert_called_once_with( + timeout_millis=15_000, + callback_timeout_millis=15_000, + flush_seconds=15, + ) + @patch("olive.systems.docker.docker_system.docker.from_env") @patch("olive.systems.docker.docker_system.tempfile.TemporaryDirectory") @patch("olive.systems.docker.docker_system.find_all_resources") diff --git a/test/test_telemetry.py b/test/test_telemetry.py new file mode 100644 index 0000000000..e0883c5186 --- /dev/null +++ b/test/test_telemetry.py @@ -0,0 +1,2027 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +# pylint: disable=duplicate-code,protected-access,redefined-outer-name +"""Tests for the SQLite-backed telemetry pipeline. + +Covers enabled, full opt-out, and CI recipe-only semantics; the +ALLOWED_KEYS whitelist filtering, the durable SQLite store, the single-drainer +process lock, the background uploader's success/poison/transient handling, and +the Common Schema serialization helpers. No test touches the network or the real +user profile: the HTTP transport is stubbed and the store directory is +redirected to a temp dir. +""" + +import hashlib +import json +import os +import sqlite3 +import stat +import subprocess +import sys +import tempfile +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import olive.telemetry.deviceid._store as deviceid_store_mod +import olive.telemetry.library.transport as transport_mod +import olive.telemetry.telemetry as tmod +import olive.telemetry.utils as telemetry_utils +from olive.telemetry.library.connection_string_parser import ConnectionStringParser +from olive.telemetry.library.serialization import CommonSchemaJsonSerializationHelper as Serializer +from olive.telemetry.offline_store import OfflineEventStore +from olive.telemetry.process_lock import ProcessDrainLock +from olive.telemetry.telemetry_redaction import scrub_value_for_telemetry +from olive.telemetry.uploader import DrainOutcome, DrainResult, EventUploader + +ACTION_EVENT_NAME = tmod.ACTION_EVENT_NAME +ERROR_EVENT_NAME = tmod.ERROR_EVENT_NAME +HEARTBEAT_EVENT_NAME = tmod.HEARTBEAT_EVENT_NAME +RECIPE_EVENT_NAME = tmod.RECIPE_EVENT_NAME +Telemetry = tmod.Telemetry +is_ci_environment = tmod.is_ci_environment + +_FULL_DISABLE_VAR = "ORT_DISABLE_TELEMETRY" +_OPT_OUT_VAR = "OLIVE_DISABLE_TELEMETRY" +_CI_VARS = ( + "CI", + "TF_BUILD", + "GITHUB_ACTIONS", + "GITLAB_CI", + "CIRCLECI", + "TRAVIS", + "JENKINS_URL", + "CODEBUILD_BUILD_ID", + "BUILDKITE", + "TEAMCITY_VERSION", + "APPVEYOR", + "BITBUCKET_BUILD_NUMBER", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", +) + + +@pytest.fixture +def tenv(tmp_path, monkeypatch): + """Hermetic telemetry environment. + + Clears CI/opt-out signals so each test sets its own mode, stubs the HTTP + transport (recording every send in ``.sends``), and redirects the durable + store off the real profile. + """ + Telemetry._instance = None + Telemetry._process_disabled = False + Telemetry._heartbeat_enqueued = False + for var in (_FULL_DISABLE_VAR, _OPT_OUT_VAR, *_CI_VARS): + monkeypatch.delenv(var, raising=False) + + sends = [] + + def _record_send(self, payload, timeout_sec, item_count=1, on_send_admitted=None): + if on_send_admitted is not None: + on_send_admitted() + sends.append({"item_count": item_count, "size": len(payload), "payload": payload}) + return True, 204 + + monkeypatch.setattr(transport_mod.HttpJsonPostTransport, "send", _record_send) + monkeypatch.setattr(tmod, "get_telemetry_base_dir", lambda: tmp_path) + monkeypatch.setattr(telemetry_utils, "get_telemetry_base_dir", lambda: tmp_path) + monkeypatch.setattr(deviceid_store_mod, "get_telemetry_base_dir", lambda: tmp_path) + + yield SimpleNamespace(sends=sends, tmp_path=tmp_path) + + inst = Telemetry._instance + if inst is not None: + uploader = getattr(inst, "_uploader", None) + if uploader is not None: + uploader.stop_loop(5) + Telemetry._instance = None + Telemetry._process_disabled = False + Telemetry._heartbeat_enqueued = False + + +def _quiesce(t): + """Stop the background loop and drain its durable queue deterministically.""" + if t._uploader is not None: + t._uploader.stop_loop(5) + for _ in range(20): + if t._store is None or t._store.count() == 0: + break + t._uploader.drain_once() + + +def _sent_event_names(sends): + names = [] + for s in sends: + payload = bytes(s["payload"]) + names.extend( + token.decode() + for token in (b"OliveHeartbeat", b"OliveRecipe", b"OliveAction", b"OliveError") + if token in payload + ) + return names + + +# -------------------------------------------------------------------------- +# Full opt-out semantics +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", " yes ", "on", "y"]) +@pytest.mark.parametrize("variable", [_FULL_DISABLE_VAR, _OPT_OUT_VAR]) +def test_environment_opt_out_uses_canonical_allowlist(tenv, monkeypatch, variable, value): + monkeypatch.setenv(variable, value) + + assert tmod.is_telemetry_disabled_by_environment() is True + + +@pytest.mark.parametrize("value", ["", "0", "false", "no", "off", "unexpected"]) +@pytest.mark.parametrize("variable", [_FULL_DISABLE_VAR, _OPT_OUT_VAR]) +def test_environment_opt_out_rejects_non_allowlisted_values(tenv, monkeypatch, variable, value): + monkeypatch.setenv(variable, value) + + assert tmod.is_telemetry_disabled_by_environment() is False + + +def test_ci_is_recipe_only_with_no_heartbeat(tenv, monkeypatch): + monkeypatch.setenv("CI", "1") + t = Telemetry() + + # CI suppresses the device-id heartbeat but still persists recipe events. + assert t._store is not None + assert t.accepts_detailed_events is False + assert t._uploader._thread is None + + t.log(RECIPE_EVENT_NAME, {"recipe_name": "r", "success": True}) + t.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + + _quiesce(t) + names = _sent_event_names(tenv.sends) + assert "OliveHeartbeat" not in names + assert "OliveRecipe" in names + + +def test_ci_recipe_queue_does_not_drain_local_details(tenv, monkeypatch): + local_store = OfflineEventStore(str(tenv.tmp_path / tmod.DB_FILE_NAME)) + local_store.store(b'{"name":"OliveAction"}') + monkeypatch.setenv("CI", "1") + + telemetry = Telemetry() + telemetry.log(RECIPE_EVENT_NAME, {"recipe_name": "r", "success": True}) + _quiesce(telemetry) + + assert local_store.count() == 1 + assert "OliveAction" not in _sent_event_names(tenv.sends) + assert "OliveRecipe" in _sent_event_names(tenv.sends) + local_store.close() + + +def test_ci_shutdown_flushes_recipe_within_callback_budget(tenv, monkeypatch): + monkeypatch.setenv("CI", "1") + with patch.object(EventUploader, "start") as start: + telemetry = Telemetry() + telemetry.log(RECIPE_EVENT_NAME, {"recipe_name": "r", "success": True}) + + telemetry.shutdown(timeout_millis=0, callback_timeout_millis=1_000) + + start.assert_not_called() + assert _sent_event_names(tenv.sends) == ["OliveRecipe"] + + +def test_ci_user_opt_out_fully_disables_telemetry(tenv, monkeypatch): + monkeypatch.setenv("CI", "1") + monkeypatch.setenv(_OPT_OUT_VAR, "1") + with ( + patch.object(tmod, "OfflineEventStore") as mock_store, + patch.object(tmod, "get_hashed_device_id_and_status") as mock_device_id, + ): + telemetry = Telemetry() + + assert Telemetry._instance is None + assert telemetry._enabled is False + assert telemetry._store is None + mock_store.assert_not_called() + mock_device_id.assert_not_called() + assert tenv.sends == [] + + +def test_user_opt_out_sends_and_persists_nothing(tenv, monkeypatch): + monkeypatch.setenv(_OPT_OUT_VAR, "1") + with ( + patch.object(tmod, "OfflineEventStore") as mock_store, + patch.object(tmod, "EventUploader") as mock_uploader, + patch.object(tmod, "get_hashed_device_id_and_status") as mock_device_id, + ): + t = Telemetry() + + assert t._enabled is False + assert t._store is None + assert t._uploader is None + assert t.accepts_detailed_events is False + mock_store.assert_not_called() + mock_uploader.assert_not_called() + mock_device_id.assert_not_called() + assert tenv.sends == [] + + +def test_user_opt_out_is_latched_across_reinitialization(tenv, monkeypatch): + monkeypatch.setenv(_OPT_OUT_VAR, "true") + first = Telemetry() + first.shutdown() + monkeypatch.delenv(_OPT_OUT_VAR) + + second = Telemetry() + + assert Telemetry._instance is None + assert second is not first + assert second._enabled is False + assert second._store is None + assert second._uploader is None + assert tenv.sends == [] + + +def test_full_disable_and_ci_send_nothing(tenv, monkeypatch): + monkeypatch.setenv(_FULL_DISABLE_VAR, "1") + monkeypatch.setenv("CI", "1") + t = Telemetry() + _quiesce(t) + + # Explicit full suppression + CI: record and send nothing at all. + assert Telemetry._instance is None + assert t._enabled is False + assert t._store is None + assert tenv.sends == [] + + +def test_full_disable_sends_nothing(tenv, monkeypatch): + monkeypatch.setenv(_FULL_DISABLE_VAR, "1") + t = Telemetry() + _quiesce(t) + + assert Telemetry._instance is None + assert t._enabled is False + assert t._store is None + assert t._uploader is None + assert tenv.sends == [] + + +def test_public_disable_before_initialization_only_latches_suppression(tenv): + from olive.telemetry.telemetry_extensions import log_action, log_error, log_recipe_result + + with ( + patch.object(tmod, "OfflineEventStore") as mock_store, + patch.object(tmod, "EventUploader") as mock_uploader, + patch.object(tmod, "get_hashed_device_id_and_status") as mock_device_id, + ): + tmod.disable_telemetry() + assert Telemetry._instance is None + log_action("cli", "work", 1.0, True) + log_error("RuntimeError", "boom") + log_recipe_result("recipe", True) + assert Telemetry._instance is None + telemetry = Telemetry() + tmod.disable_telemetry() + + assert Telemetry._instance is None + mock_store.assert_not_called() + mock_uploader.assert_not_called() + mock_device_id.assert_not_called() + assert telemetry._enabled is False + assert telemetry._store is None + assert tenv.sends == [] + + +def test_environment_opt_out_helpers_do_not_publish_singleton(tenv, monkeypatch): + from olive.telemetry.telemetry_extensions import log_action, log_error, log_recipe_result + + monkeypatch.setenv(_OPT_OUT_VAR, "1") + log_action("cli", "work", 1.0, True) + log_error("RuntimeError", "boom") + log_recipe_result("recipe", True) + + assert Telemetry._instance is None + + +def test_environment_opt_out_after_initialization_retains_queue(tenv, monkeypatch): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + telemetry.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1, "success": True}) + row_count = telemetry._store.count() + + monkeypatch.setenv(_OPT_OUT_VAR, "1") + assert Telemetry.get_or_create_if_enabled() is None + + assert telemetry._disabled is True + remaining_store = OfflineEventStore(str(tenv.tmp_path / tmod.DB_FILE_NAME)) + assert remaining_store.count() == row_count + remaining_store.close() + + +def test_enabled_records_heartbeat_and_events(tenv): + import uuid + + t = Telemetry() + session_guid = uuid.UUID(t._app_session_guid) + assert session_guid.version == 4 + assert session_guid.variant == uuid.RFC_4122 + + assert t._enabled is True + assert t._store is not None + assert t.accepts_detailed_events is True + + t.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + + _quiesce(t) + names = _sent_event_names(tenv.sends) + assert "OliveHeartbeat" in names + assert "OliveAction" in names + + +def test_heartbeat_is_stored_durably_before_uploader_starts(tenv): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + + batch = telemetry._store.get_batch(10) + assert len(batch) == 1 + assert json.loads(batch[0][1])["name"] == HEARTBEAT_EVENT_NAME + assert tenv.sends == [] + + +def test_heartbeat_releases_minimal_fallback_when_enrichment_fails(tenv): + with ( + patch.object(EventUploader, "start"), + patch.object( + tmod, + "get_hashed_device_id_and_status", + return_value=("c:device", SimpleNamespace(value="Existing")), + ), + patch.object(tmod.platform, "system", side_effect=RuntimeError("unavailable")), + ): + telemetry = Telemetry() + + batch = telemetry._store.get_batch(10) + assert len(batch) == 1 + data = json.loads(batch[0][1])["data"] + assert data["deviceId"] == "c:device" + assert data["deviceIdStatus"] == "Existing" + assert "os" not in data + + +def test_disable_telemetry_stops_detailed_events(tenv): + t = Telemetry() + _quiesce(t) + t.disable_telemetry() + + assert t._enabled is False + before = t._store.count() if t._store is not None else 0 + t.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + after = t._store.count() if t._store is not None else 0 + assert after == before + t.shutdown() + assert Telemetry()._enabled is False + assert _sent_event_names(tenv.sends).count("OliveHeartbeat") == 1 + + +def test_runtime_disable_does_not_emit_an_additional_heartbeat(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + tenv.sends.clear() + telemetry.disable_telemetry() + telemetry.disable_telemetry() + + assert tenv.sends == [] + + +def test_runtime_disable_retains_all_unsent_rows(tenv): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + other_store = OfflineEventStore(telemetry._store.db_path) + other_store.store(b'{"other":1}') + telemetry.log( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}, + ) + + telemetry.disable_telemetry() + + payloads = [payload for _, payload in other_store.get_batch(10)] + assert b'{"other":1}' in payloads + assert {json.loads(payload)["name"] for payload in payloads if payload != b'{"other":1}'} == { + HEARTBEAT_EVENT_NAME, + ACTION_EVENT_NAME, + } + other_store.close() + + +def test_runtime_disable_serializes_with_inflight_log(tenv): + with patch.object(EventUploader, "start"): + telemetry = Telemetry() + entered = threading.Event() + release = threading.Event() + original_store = telemetry._store.store + + def blocked_store(payload): + entered.set() + assert release.wait(5) + return original_store(payload) + + with patch.object(telemetry._store, "store", side_effect=blocked_store): + logging_thread = threading.Thread( + target=telemetry.log, + args=( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}, + ), + ) + logging_thread.start() + assert entered.wait(5) + disable_thread = threading.Thread(target=telemetry.disable_telemetry) + disable_thread.start() + release.set() + logging_thread.join(5) + disable_thread.join(5) + + assert not logging_thread.is_alive() + assert not disable_thread.is_alive() + assert OfflineEventStore(str(tenv.tmp_path / tmod.DB_FILE_NAME)).count() == 2 + + +def test_shutdown_closes_store_without_uploader(): + t = object.__new__(Telemetry) + t._disabled = False + t._uploader = None + t._store = MagicMock() + t._initialized = True + + store = t._store + t.shutdown() + + store.close.assert_called_once() + assert t._store is None + assert t._initialized is False + + +def test_shutdown_uses_one_overall_budget(): + t = object.__new__(Telemetry) + t._disabled = False + t._uploader = MagicMock() + t._uploader.stop_loop.return_value = True + t._store = MagicMock() + uploader = t._uploader + + monotonic = MagicMock(side_effect=[100.0, 101.0, 102.0]) + with patch("olive.telemetry.telemetry.time", SimpleNamespace(monotonic=monotonic)): + t.shutdown(timeout_millis=5_000, callback_timeout_millis=5_000, flush_seconds=5) + + uploader.stop_loop.assert_called_once_with(join_timeout_seconds=4.0) + uploader.flush.assert_called_once_with(3.0) + assert t._uploader is None + assert t._store is None + assert t._initialized is False + + +def test_shutdown_does_not_wait_or_flush_after_full_disable(): + telemetry = object.__new__(Telemetry) + telemetry._initialized = True + telemetry._disabled = True + telemetry._uploader = MagicMock() + telemetry._uploader.stop_loop.return_value = True + telemetry._store = MagicMock() + uploader = telemetry._uploader + + telemetry.shutdown(timeout_millis=5_000, callback_timeout_millis=5_000, flush_seconds=5) + + uploader.stop_loop.assert_called_once_with(join_timeout_seconds=0) + uploader.flush.assert_not_called() + + +def test_reinitialization_does_not_enqueue_second_heartbeat(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + telemetry.shutdown() + + restarted = Telemetry() + restarted.log(ACTION_EVENT_NAME, {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}) + _quiesce(restarted) + + names = _sent_event_names(tenv.sends) + assert names.count("OliveHeartbeat") == 1 + assert names.count("OliveAction") == 1 + + +def test_live_uploader_keeps_store_open(): + telemetry = object.__new__(Telemetry) + telemetry._disabled = False + telemetry._uploader = MagicMock() + telemetry._uploader.stop_loop.return_value = False + telemetry._store = MagicMock() + + telemetry.shutdown(timeout_millis=0) + + telemetry._store.close.assert_not_called() + + +def test_closed_store_disables_telemetry(tenv): + closed_store = MagicMock(is_open=False) + with ( + patch.object(tmod, "OfflineEventStore", return_value=closed_store), + patch.object(tmod, "EventUploader") as mock_uploader, + ): + t = Telemetry() + + assert t._enabled is False + assert t._store is None + assert Telemetry._heartbeat_enqueued is False + mock_uploader.assert_not_called() + + +def test_closed_store_allows_initialization_retry(tenv): + closed_store = MagicMock(is_open=False) + open_store = MagicMock(is_open=True) + with ( + patch.object(tmod, "OfflineEventStore", side_effect=[closed_store, open_store]), + patch.object(tmod, "EventUploader") as mock_uploader, + ): + first = Telemetry() + assert first._initialized is False + + second = Telemetry() + + assert second is first + assert second._initialized is True + assert second._enabled is True + mock_uploader.assert_called_once_with(open_store, instrumentation_key=second._instrumentation_key) + + +# -------------------------------------------------------------------------- +# Whitelist filtering / payload building +# -------------------------------------------------------------------------- + + +def test_build_payload_drops_non_whitelisted_keys(tenv): + t = Telemetry() + _quiesce(t) + + payload = t._build_payload( + ACTION_EVENT_NAME, + { + "invoked_from": "cli", + "action_name": "WorkflowRun", + "duration_ms": 1.0, + "success": True, + "secret": "SHOULD_NOT_BE_SENT", + }, + ) + data = json.loads(payload)["data"] + assert "secret" not in data + assert data["actionName"] == "WorkflowRun" + # Defaults are stamped on every event. + assert data["appName"] == "Olive" + assert data["LibraryVersion"] + assert data["AppSessionGuid"] + assert "appVersion" not in data + assert "appSessionGuid" not in data + + +def test_build_payload_returns_none_for_unknown_event(tenv): + t = Telemetry() + _quiesce(t) + assert t._build_payload("TotallyUnknownEvent", {"k": "v"}) is None + + +def test_build_payload_heartbeat_uses_flat_os_fields(tenv): + t = Telemetry() + _quiesce(t) + + payload = t._build_payload( + HEARTBEAT_EVENT_NAME, + { + "device_id": "DEVICE", + "device_id_status": "ok", + "os": "Windows", + "os_version": "10.0.22631", + "os_release": "11", + "os_arch": "AMD64", + "leak": "DROP", + }, + ) + data = json.loads(payload)["data"] + assert data["deviceId"] == "DEVICE" + assert data["deviceIdStatus"] == "ok" + assert data["os"] == "Windows" + assert data["osVersion"] == "10.0.22631" + assert "leak" not in data + + +def test_build_payload_drops_non_scalar_whitelisted_values(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + { + "recipe_name": "optimize", + "success": True, + "config_overrides": { + "url": "https://user:secret@example.test/model", + Path("/home/alice/private"): [Path("/home/alice/model.onnx")], + }, + }, + ) + data = json.loads(payload)["data"] + + assert "configOverrides" not in data + + +def test_build_payload_drops_object_before_stringification(tenv): + class SensitiveObject: + def __str__(self): + return r"C:\Users\Alice\private.txt" + + telemetry = Telemetry() + _quiesce(telemetry) + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + {"recipe_name": "optimize", "success": True, "model_task": SensitiveObject()}, + ) + + serialized = payload.decode("utf-8") + assert "modelTask" not in json.loads(payload)["data"] + assert "Alice" not in serialized + + +def test_final_scrubber_redacts_nested_credential_aliases(): + scrubbed = scrub_value_for_telemetry( + { + "modelTask": { + "access_token": "snake-secret", + "accesstoken": "compact-secret", + "apiKey": "camel-secret", + "auth_header": "header-secret", + "credential_value": "credential-secret", + "docker_env": {"TOKEN": "env-secret"}, + "modelpath": "private/model.onnx", + "environmentVariables": {"HOME": "/home/private"}, + } + } + ) + + assert scrubbed["modelTask"] == { + "access_token": "", + "accesstoken": "", + "apiKey": "", + "auth_header": "", + "credential_value": "", + "docker_env": "", + "environmentVariables": "", + "modelpath": "", + } + + +def test_recipe_snapshot_json_remains_parseable_after_url_redaction(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + snapshot = json.dumps({"auth": "token=[secret]", "endpoint": "[path]", "z_value": "kept"}) + + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + {"recipe_name": "optimize", "success": True, "config_overrides": snapshot}, + ) + serialized_snapshot = json.loads(payload)["data"]["configOverrides"] + + assert json.loads(serialized_snapshot) == { + "auth": "", + "endpoint": "[path]", + "z_value": "kept", + } + + raw_snapshot = json.dumps( + { + "access_token": "snake-secret", + "accesstoken": "compact-secret", + "modelpath": "private/model.onnx", + } + ) + payload = telemetry._build_payload( + RECIPE_EVENT_NAME, + {"recipe_name": "optimize", "success": True, "config_overrides": raw_snapshot}, + ) + serialized_snapshot = json.loads(payload)["data"]["configOverrides"] + assert json.loads(serialized_snapshot) == { + "access_token": "", + "accesstoken": "", + "modelpath": "", + } + + +def test_final_scrubber_scrubs_text_bytes_and_drops_binary(): + scrubbed = scrub_value_for_telemetry( + { + "text": rb"C:\Users\Alice Smith\model.onnx", + "binary": b"\xff\x00", + } + ) + + assert scrubbed["text"] == "[path]" + assert scrubbed["binary"] == "[binary]" + + +def test_non_finite_event_is_rejected_without_affecting_next_event(tenv): + telemetry = Telemetry() + _quiesce(telemetry) + telemetry._uploader = None + + telemetry.log( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "bad", "duration_ms": float("nan"), "success": True}, + ) + telemetry.log( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "good", "duration_ms": 1.0, "success": True}, + ) + + assert telemetry._store.count() == 1 + + +def test_device_id_is_canonical_shared_hash(): + import olive.telemetry.deviceid.deviceid as deviceid + + raw_id = "123e4567-e89b-42d3-a456-426614174000" + with patch.dict( + deviceid._device_id_state, + {"device_id": raw_id, "status": deviceid.DeviceIdStatus.EXISTING}, + clear=True, + ): + hashed, status = deviceid.get_hashed_device_id_and_status() + + expected = hashlib.sha256(raw_id.encode("utf-8")).hexdigest() + assert hashed == f"c:{expected}" + assert len(hashed) == 66 + assert status == deviceid.DeviceIdStatus.EXISTING + + +def test_device_id_hash_matches_native_algorithm_vector(): + import olive.telemetry.deviceid.deviceid as deviceid + + raw_id = "01234567-89ab-4def-8123-456789abcdef" + with patch.dict( + deviceid._device_id_state, + {"device_id": raw_id, "status": deviceid.DeviceIdStatus.EXISTING}, + clear=True, + ): + hashed, _ = deviceid.get_hashed_device_id_and_status() + + assert hashed == f"c:{hashlib.sha256(raw_id.encode('utf-8')).hexdigest()}" + assert hashed == "c:6225bd190d6ccf87766a49c9986d174def3391fe175a61525e49a1d2334d6a43" + + +def test_global_metadata_is_merged_then_filtered(tenv): + t = Telemetry() + _quiesce(t) + + t.add_global_metadata({"app_version": "9.9.9", "app_instance_id": "mallory@example.test", "not_allowed": "DROP"}) + payload = t._build_payload( + ACTION_EVENT_NAME, + {"invoked_from": "cli", "action_name": "x", "duration_ms": 1.0, "success": True}, + ) + data = json.loads(payload)["data"] + assert data["LibraryVersion"] == tmod.VERSION + assert data["AppSessionGuid"] == t._app_session_guid + assert "mallory" not in payload.decode("utf-8") + assert "not_allowed" not in data + + +def test_event_attributes_override_metadata(tenv): + t = Telemetry() + _quiesce(t) + payload = t._build_payload( + ERROR_EVENT_NAME, + {"exception_type": "ValueError", "exception_message": "safe"}, + {"exception_message": r"C:\Users\Mallory\secret.txt"}, + ) + assert json.loads(payload)["data"]["exceptionMessage"] == "safe" + + +def test_error_event_whitelist(tenv): + t = Telemetry() + _quiesce(t) + payload = t._build_payload( + ERROR_EVENT_NAME, + {"exception_type": "RuntimeError", "exception_message": "boom", "stack": "SENSITIVE"}, + ) + data = json.loads(payload)["data"] + assert data["exceptionType"] == "RuntimeError" + assert data["exceptionMessage"] == "boom" + assert "stack" not in data + + +@pytest.mark.parametrize("event_name", sorted(tmod.ALLOWED_KEYS)) +def test_all_whitelisted_fields_use_canonical_names(tenv, event_name): + t = Telemetry() + _quiesce(t) + source = dict.fromkeys(tmod.ALLOWED_KEYS[event_name], "value") + for snapshot_key in ("config_overrides", "package_config_overrides"): + if snapshot_key in source: + source[snapshot_key] = '{"key":"value"}' + + payload = t._build_payload(event_name, source) + data = json.loads(payload)["data"] + expected = {tmod.FIELD_NAMES.get(key, key) for key in source} + expected.update({"appName", "LibraryVersion", "AppSessionGuid"}) + + assert set(data) == expected + for source_name, canonical_name in tmod.FIELD_NAMES.items(): + if source_name != canonical_name and source_name in source: + assert source_name not in data + + +# -------------------------------------------------------------------------- +# CI detection +# -------------------------------------------------------------------------- + + +def test_is_ci_environment(monkeypatch): + for var in (_FULL_DISABLE_VAR, _OPT_OUT_VAR, *_CI_VARS): + monkeypatch.delenv(var, raising=False) + assert is_ci_environment() is False + monkeypatch.setenv("GITHUB_ACTIONS", "true") + assert is_ci_environment() is True + + +@pytest.mark.parametrize("value", ["", "0", "false", " no ", "OFF"]) +def test_false_ci_values_are_not_ci(monkeypatch, value): + with patch.dict(os.environ, {"CI": value}, clear=True): + assert is_ci_environment() is False + + +# -------------------------------------------------------------------------- +# Durable SQLite store +# -------------------------------------------------------------------------- + + +def _new_store(**kwargs): + db = os.path.join(tempfile.mkdtemp(), "olive_telemetry.db") + return OfflineEventStore(db, **kwargs) + + +def test_store_is_fifo(): + store = _new_store() + for i in range(5): + store.store(f'{{"e":{i}}}'.encode()) + assert store.count() == 5 + batch = store.get_batch(3) + assert [payload for _, payload in batch] == [b'{"e":0}', b'{"e":1}', b'{"e":2}'] + + +def test_store_closes_connection_when_initialization_fails(tmp_path): + connection = MagicMock() + connection.execute.side_effect = RuntimeError("pragma failed") + + with patch("olive.telemetry.offline_store.sqlite3.connect", return_value=connection): + store = OfflineEventStore(str(tmp_path / "failed.db")) + + assert store.is_open is False + connection.close.assert_called_once() + + +def test_store_delete(): + store = _new_store() + store.store(b'{"a":1}') + store.store(b'{"b":2}') + ids = [row_id for row_id, _ in store.get_batch(10)] + store.delete(ids[:1]) + assert store.count() == 1 + + +def test_store_trims_over_capacity(): + store = _new_store(max_records=8) + for i in range(40): + store.store(f'{{"i":{i}}}'.encode()) + assert store.count() <= 8 + + +def test_store_rejects_empty_payload(): + store = _new_store() + assert store.store(b"") is False + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX permissions") +def test_store_uses_owner_only_permissions(): + store = _new_store() + db_path = Path(store.db_path) + assert stat.S_IMODE(db_path.parent.stat().st_mode) == 0o700 + assert stat.S_IMODE(db_path.stat().st_mode) == 0o600 + + +def test_empty_permission_path_does_not_chmod_cwd(): + from olive.telemetry.offline_store import _chmod_best_effort + + with patch("olive.telemetry.offline_store.Path.chmod") as mock_chmod: + _chmod_best_effort("", 0o700) + + mock_chmod.assert_not_called() + + +def test_store_sets_schema_version(): + store = _new_store() + + assert store._conn.execute("PRAGMA user_version").fetchone()[0] == 3 + + +def test_store_migrates_availability_column(tmp_path): + db_path = tmp_path / "v1.db" + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE events (id INTEGER PRIMARY KEY AUTOINCREMENT, payload BLOB NOT NULL)") + conn.execute("PRAGMA user_version=1") + conn.commit() + conn.close() + + store = OfflineEventStore(str(db_path)) + + columns = {row[1] for row in store._conn.execute("PRAGMA table_info(events)")} + assert columns == {"id", "payload", "available_at", "acknowledged"} + assert store._conn.execute("PRAGMA user_version").fetchone()[0] == 3 + assert store.store(b'{"migrated":1}') is True + + +def test_reserved_event_is_hidden_until_release(): + store = _new_store() + row_id = store.reserve(b'{"minimal":1}', 60) + + assert row_id is not None + assert store.count() == 1 + assert store.get_batch(10) == [] + + assert store.release(row_id, b'{"enriched":1}') is True + assert store.get_batch(10) == [(row_id, b'{"enriched":1}')] + + +def test_failed_delete_is_reported_and_rolled_back(): + store = _new_store() + store.store(b'{"a":1}') + row_id = store.get_batch(1)[0][0] + store._conn.execute("CREATE TRIGGER fail_delete BEFORE DELETE ON events BEGIN SELECT RAISE(FAIL, 'blocked'); END") + store._conn.commit() + + assert store.delete([row_id]) is False + assert store.count() == 1 + + +def test_store_operations_bound_busy_timeout_to_deadline(): + store = _new_store() + store.store(b'{"a":1}') + row_id = store.get_batch(1)[0][0] + statements = [] + store._conn.set_trace_callback(statements.append) + + with patch("olive.telemetry.offline_store.time.monotonic", return_value=100.0): + assert store.delete([row_id], deadline=100.025) + + bounded = [statement for statement in statements if statement.startswith("PRAGMA busy_timeout=")] + assert bounded[0] in {"PRAGMA busy_timeout=24", "PRAGMA busy_timeout=25"} + assert bounded[-1] == "PRAGMA busy_timeout=3000" + + +# -------------------------------------------------------------------------- +# Single-drainer process lock +# -------------------------------------------------------------------------- + + +def _lock_path(): + return os.path.join(tempfile.mkdtemp(), "olive_telemetry.db.lock") + + +def test_lock_is_mutually_exclusive(): + path = _lock_path() + a = ProcessDrainLock(path) + b = ProcessDrainLock(path) + assert a.acquire() is True + assert b.acquire() is False # held by a + a.release() + assert b.acquire() is True # released + b.release() + + +def test_lock_reacquire_is_idempotent(): + a = ProcessDrainLock(_lock_path()) + assert a.acquire() is True + assert a.acquire() is True # already held + assert a.held is True + a.release() + assert a.held is False + + +# -------------------------------------------------------------------------- +# Uploader drain classification (no real network) +# -------------------------------------------------------------------------- + + +def _store_and_uploader(): + db = os.path.join(tempfile.mkdtemp(), "olive_telemetry.db") + store = OfflineEventStore(db) + uploader = EventUploader(store, instrumentation_key="abc-def") + return store, uploader + + +def test_uploader_deletes_on_success(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = lambda *a, **k: (True, 204) + result = uploader.drain_once() + assert (result.delivered, result.left, result.outcome) == (1, 0, DrainOutcome.PROGRESS) + assert store.count() == 0 + + +def test_uploader_retention_latch_preserves_inflight_success(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + entered = threading.Event() + release = threading.Event() + result = [] + + def blocked_send(*_args, **_kwargs): + _kwargs["on_send_admitted"]() + entered.set() + assert release.wait(5) + return True, 204 + + uploader._transport.send = blocked_send + drain_thread = threading.Thread(target=lambda: result.append(uploader.drain_once())) + drain_thread.start() + assert entered.wait(5) + + uploader.retain_queued_rows() + release.set() + drain_thread.join(5) + + assert not drain_thread.is_alive() + assert result[0].outcome is DrainOutcome.TRANSPORT_RETRY + assert store.count() == 1 + + +def test_uploader_retention_latch_prevents_new_send(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock() + + uploader.retain_queued_rows() + + assert uploader.drain_once().outcome is DrainOutcome.TRANSPORT_RETRY + uploader._transport.send.assert_not_called() + assert store.count() == 1 + + +def test_uploader_retries_failed_delete_without_reposting(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + original_delete = store.delete + delete_attempts = 0 + + def flaky_delete(ids, deadline=None): + nonlocal delete_attempts + delete_attempts += 1 + return False if delete_attempts == 1 else original_delete(ids, deadline) + + store.delete = flaky_delete + + assert uploader.drain_once().outcome is DrainOutcome.DELETE_RETRY + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert store.count() == 0 + uploader._transport.send.assert_called_once() + + +def test_uploader_retries_failed_acknowledgement_without_reposting(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + original_acknowledge = store.acknowledge + acknowledge_attempts = 0 + + def flaky_acknowledge(ids, deadline=None): + nonlocal acknowledge_attempts + acknowledge_attempts += 1 + return False if acknowledge_attempts == 1 else original_acknowledge(ids, deadline) + + store.acknowledge = flaky_acknowledge + + assert uploader.drain_once().outcome is DrainOutcome.ACKNOWLEDGE_RETRY + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert store.count() == 0 + uploader._transport.send.assert_called_once() + + +def test_acknowledged_rows_are_deleted_without_reposting_after_restart(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + store.delete = MagicMock(return_value=False) + + assert uploader.drain_once().outcome is DrainOutcome.DELETE_RETRY + assert store._conn.execute("SELECT acknowledged FROM events").fetchone()[0] == 1 + db_path = store.db_path + store.close() + + reopened = OfflineEventStore(db_path) + restarted = EventUploader(reopened, instrumentation_key="abc-def") + restarted._transport.send = MagicMock() + assert restarted.drain_once().outcome is DrainOutcome.PROGRESS + assert reopened.count() == 0 + restarted._transport.send.assert_not_called() + + +def test_uploader_reports_storage_read_failure(): + store, uploader = _store_and_uploader() + store.get_batch_for_upload = MagicMock(return_value=None) + + assert uploader.drain_once().outcome is DrainOutcome.STORAGE_RETRY + + +def test_uploader_uses_only_remaining_deadline(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(False, None)) + + with patch("olive.telemetry.uploader.time.monotonic", return_value=100.75): + result = uploader.drain_once(deadline=101.0) + + assert (result.delivered, result.left, result.outcome) == (0, 1, DrainOutcome.TRANSPORT_RETRY) + assert uploader._transport.send.call_args.args[1] == pytest.approx(0.25) + + +def test_request_drain_only_wakes_lock_holder(): + _, uploader = _store_and_uploader() + uploader._wake = MagicMock() + uploader._drain_lock = MagicMock(held=False) + + uploader.request_drain() + uploader._wake.set.assert_not_called() + + uploader._drain_lock.held = True + uploader.request_drain() + uploader._wake.set.assert_called_once() + + +def test_uploader_drops_poison_4xx(): + store, uploader = _store_and_uploader() + store.store(b'{"bad":1}') + uploader._transport.send = lambda *a, **k: (False, 400) + uploader.drain_once() + assert store.count() == 0 # dropped, not retried forever + + +def test_uploader_retains_transient_5xx(): + store, uploader = _store_and_uploader() + store.store(b'{"later":1}') + uploader._transport.send = lambda *a, **k: (False, 503) + result = uploader.drain_once() + assert (result.delivered, result.left, result.outcome) == (0, 1, DrainOutcome.TRANSPORT_RETRY) + assert store.count() == 1 # kept for retry + + +@pytest.mark.parametrize("status", [507, 520, 599]) +def test_uploader_retains_all_server_errors(status): + store, uploader = _store_and_uploader() + store.store(b'{"later":1}') + uploader._transport.send = lambda *a, **k: (False, status) + + assert uploader.drain_once().outcome is DrainOutcome.TRANSPORT_RETRY + assert store.count() == 1 + + +@pytest.mark.parametrize("status", [400, 413, 422]) +def test_uploader_isolates_rejected_event(status): + store, uploader = _store_and_uploader() + store.store(b'{"bad":1}') + store.store(b'{"valid":1}') + + def send(payload, timeout, item_count=1, on_send_admitted=None): + if on_send_admitted is not None: + on_send_admitted() + if item_count > 1 or b'"bad"' in payload: + return (False, status) + return (True, 204) + + uploader._transport.send = MagicMock(side_effect=send) + + assert uploader.drain_once().outcome is DrainOutcome.SPLIT + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert uploader.drain_once().outcome is DrainOutcome.PROGRESS + assert store.count() == 0 + assert uploader._transport.send.call_count == 3 + + +@pytest.mark.parametrize("status", [400, 413, 422]) +def test_flush_completes_poison_isolation_within_one_process(status): + store, uploader = _store_and_uploader() + store.store(b'{"bad":1}') + store.store(b'{"valid":1}') + + def send(payload, timeout, item_count=1, on_send_admitted=None): + if on_send_admitted is not None: + on_send_admitted() + if item_count > 1 or b'"bad"' in payload: + return (False, status) + return (True, 204) + + uploader._transport.send = MagicMock(side_effect=send) + uploader.flush(1) + + assert store.count() == 0 + assert uploader._transport.send.call_count == 3 + + +def test_flush_retries_acknowledged_delete_without_reposting(): + store, uploader = _store_and_uploader() + store.store(b'{"ok":1}') + uploader._transport.send = MagicMock(return_value=(True, 204)) + original_delete = store.delete + delete_attempts = 0 + + def flaky_delete(ids, deadline=None): + nonlocal delete_attempts + delete_attempts += 1 + return False if delete_attempts == 1 else original_delete(ids, deadline) + + store.delete = flaky_delete + uploader.flush(1) + + assert store.count() == 0 + uploader._transport.send.assert_called_once() + + +def test_flush_does_not_touch_lock_while_thread_is_alive(): + _, uploader = _store_and_uploader() + uploader._thread = MagicMock() + uploader._thread.is_alive.return_value = True + uploader._drain_lock.acquire = MagicMock() + uploader._drain_lock.release = MagicMock() + + uploader.flush(0.01) + + uploader._drain_lock.acquire.assert_not_called() + uploader._drain_lock.release.assert_not_called() + + +def test_uploader_backs_off_after_lock_contention(): + _, uploader = _store_and_uploader() + uploader._drain_lock = MagicMock() + uploader._drain_lock.acquire.return_value = False + uploader._stop = MagicMock() + uploader._stop.is_set.side_effect = [False, True] + uploader._wake = MagicMock() + + uploader._run() + + uploader._wake.wait.assert_called_once_with(uploader._idle_backoff) + + +def test_uploader_backs_off_after_drain_exception(): + _, uploader = _store_and_uploader() + uploader._drain_lock = MagicMock() + uploader._drain_lock.acquire.return_value = True + uploader._stop = MagicMock() + uploader._stop.is_set.side_effect = [False, False, True] + uploader._wake = MagicMock() + uploader.drain_once = MagicMock(side_effect=RuntimeError("transient failure")) + + uploader._run() + + uploader._wake.wait.assert_called_once_with(uploader._idle_backoff) + uploader._drain_lock.release.assert_called_once() + + +def test_uploader_backs_off_after_storage_failure(): + _, uploader = _store_and_uploader() + uploader._drain_lock = MagicMock() + uploader._drain_lock.acquire.return_value = True + uploader._stop = MagicMock() + uploader._stop.is_set.side_effect = [False, False, True] + uploader._wake = MagicMock() + uploader.drain_once = MagicMock(return_value=DrainResult(0, 0, DrainOutcome.STORAGE_RETRY)) + + uploader._run() + + uploader._wake.wait.assert_called_once_with(uploader._idle_backoff) + uploader._drain_lock.release.assert_called_once() + + +# -------------------------------------------------------------------------- +# Serialization + connection string parsing +# -------------------------------------------------------------------------- + + +def test_serialize_basic_types(): + assert Serializer.serialize_value(None) is None + assert Serializer.serialize_value(True) is True + assert Serializer.serialize_value(42) == 42 + assert Serializer.serialize_value("hello") == "hello" + assert Serializer.serialize_value([1, "two", 3.0]) == [1, "two", 3.0] + assert Serializer.serialize_value({"k": "v"}) == {"k": "v"} + assert Serializer.serialize_value({0: "zero", "": "skip"}) == {"0": "zero"} + assert Serializer.serialize_value({False: "false"}) == {"False": "false"} + + +def test_redacted_mapping_key_collisions_are_dropped_deterministically(): + first = {"/first": "a", "/second": "b", "safe": "kept"} + second = dict(reversed(list(first.items()))) + + assert scrub_value_for_telemetry(first) == {"safe": "kept"} + assert scrub_value_for_telemetry(second) == {"safe": "kept"} + assert Serializer.serialize_value(first) == {"safe": "kept"} + assert Serializer.serialize_value(second) == {"safe": "kept"} + + +def test_snapshot_key_collisions_are_dropped_deterministically(): + from olive.telemetry.telemetry_redaction import scrub_config_snapshot_for_telemetry + + first = {"/first": "a", "/second": "b", "safe": "kept"} + second = dict(reversed(list(first.items()))) + + assert scrub_config_snapshot_for_telemetry(first) == {"safe": "kept"} + assert scrub_config_snapshot_for_telemetry(second) == {"safe": "kept"} + + +def test_create_event_envelope(): + envelope = Serializer.create_event_envelope( + event_name="TestEvent", + timestamp=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc), + ikey="o:test-key", + data={"key": "value"}, + ) + assert envelope["name"] == "TestEvent" + assert envelope["iKey"] == "o:test-key" + assert envelope["data"] == {"key": "value"} + + +def test_connection_string_parser(): + assert ConnectionStringParser("InstrumentationKey=abc-def-ghi").instrumentation_key == "abc-def-ghi" + with pytest.raises(ValueError, match="Connection string cannot be empty"): + ConnectionStringParser("") + with pytest.raises(ValueError, match="InstrumentationKey"): + ConnectionStringParser("SomeOtherKey=value") + + +def test_http_retry_shares_one_timeout_budget(): + import urllib.error + import urllib.request + + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + with ( + patch( + "olive.telemetry.library.transport.urllib.request.urlopen", + side_effect=urllib.error.URLError("offline"), + ) as urlopen, + patch("olive.telemetry.library.transport.time.monotonic", side_effect=[10.0, 10.0, 10.75]), + ): + assert transport._do_request(request, 1.0) == (False, None) + + assert urlopen.call_count == 2 + assert urlopen.call_args_list[0].kwargs["timeout"] == pytest.approx(1.0) + assert urlopen.call_args_list[1].kwargs["timeout"] == pytest.approx(0.25) + + +def test_http_error_body_is_not_read_during_bounded_send(): + import urllib.error + import urllib.request + + response = MagicMock() + http_error = urllib.error.HTTPError("https://example.invalid", 503, "unavailable", {}, response) + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + + with patch("olive.telemetry.library.transport.urllib.request.urlopen", side_effect=http_error): + assert transport._do_request(request, 1.0) == (False, 503) + + response.read.assert_not_called() + response.close.assert_called_once() + + +def test_http_deadline_bounds_blocked_resolution(): + import urllib.request + + entered = threading.Event() + release = threading.Event() + response = MagicMock(status=204) + response.__enter__.return_value = response + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + + def blocked_urlopen(*_args, **_kwargs): + entered.set() + release.wait(5) + return response + + start = time.perf_counter() + with patch("olive.telemetry.library.transport.urllib.request.urlopen", side_effect=blocked_urlopen): + assert transport._do_request(request, 0.05) == (False, None) + assert entered.wait(1) + elapsed = time.perf_counter() - start + release.set() + + assert elapsed < 0.5 + + +def test_http_timeout_does_not_start_replacement_worker(): + import urllib.request + + entered = threading.Event() + release = threading.Event() + response = MagicMock(status=204) + response.__enter__.return_value = response + request = urllib.request.Request("https://example.invalid", data=b"{}", method="POST") + transport = transport_mod.HttpJsonPostTransport( + endpoint="https://example.invalid", + ikey="key", + compression=transport_mod.CompressionType.NO_COMPRESSION, + ) + + def blocked_urlopen(*_args, **_kwargs): + entered.set() + release.wait(5) + return response + + with patch( + "olive.telemetry.library.transport.urllib.request.urlopen", + side_effect=blocked_urlopen, + ) as urlopen: + assert transport._do_request(request, 0.05) == (False, None) + assert entered.wait(1) + assert transport._do_request(request, 0.05) == (False, None) + assert urlopen.call_count == 1 + release.set() + transport._inflight_worker.join(1) + assert transport._do_request(request, 0.05) == (True, 204) + assert urlopen.call_count == 1 + + +@pytest.mark.parametrize("status", [500, 507, 520, 599]) +def test_all_server_errors_are_retryable(status): + assert transport_mod.HttpJsonPostTransport.is_retryable(status) + + +def test_serialization_orders_maps_and_sets_deterministically(): + helper = Serializer + timestamp = datetime(2025, 1, 1, tzinfo=timezone.utc) + first = helper.create_event_envelope("event", timestamp, "o:key", {"z": {"b", "a"}, "a": 1}) + second = helper.create_event_envelope("event", timestamp, "o:key", {"a": 1, "z": {"a", "b"}}) + + assert helper.serialize_to_json_bytes(first) == helper.serialize_to_json_bytes(second) + + +# -------------------------------------------------------------------------- +# Exception-message path redaction (privacy) +# -------------------------------------------------------------------------- + + +def test_redact_paths_and_general_length_contract(): + from olive.telemetry.telemetry_redaction import MAX_TELEMETRY_STRING_LENGTH, scrub_string_for_telemetry + + assert scrub_string_for_telemetry(r"C:\Users\alice\model.onnx") == "[path]" + assert scrub_string_for_telemetry(r"\secret.onnx") == "[path]" + assert scrub_string_for_telemetry(r"failed \secret.onnx") == "failed [path]" + assert scrub_string_for_telemetry("/var/data/run/output.log") == "[path]" + assert scrub_string_for_telemetry("/secret.onnx") == "[path]" + # Last segment is a directory/username (no extension) -> fully redacted. + assert scrub_string_for_telemetry("/home/bob") == "[path]" + # UNC paths are redacted too. + assert scrub_string_for_telemetry(r"\\server\share\secret") == "[path]" + assert scrub_string_for_telemetry(r"failed C:\Users\Alice Smith\models\phi.onnx") == "failed [path]" + assert scrub_string_for_telemetry("failed /home/Alice Smith/models/phi.onnx") == "failed [path]" + assert scrub_string_for_telemetry("a/b/c") == "[path]" + assert scrub_string_for_telemetry(r"Load Users\bob\model.onnx failed") == "Load [path]" + assert scrub_string_for_telemetry("Users/Alice Smith/models/phi.onnx") == "[path]" + assert scrub_string_for_telemetry("Users/Alice Smith/model.onnx") == "[path]" + assert scrub_string_for_telemetry("models/foo.onnx") == "models/foo.onnx" + assert scrub_string_for_telemetry("ratio 3/4 and and/or") == "ratio 3/4 and and/or" + assert scrub_string_for_telemetry("before /home/alice/model.onnx\nafter") == "before [path]" + assert scrub_string_for_telemetry("https://example.test/model?token=secret") == "[path]" + assert scrub_string_for_telemetry("download(s3://private-bucket/model)") == "download([path]" + assert scrub_string_for_telemetry("GET https://host/x?path=/home/alice/model.onnx failed") == "GET [path]" + assert scrub_string_for_telemetry("open ./private/model.onnx") == "open [path]" + assert scrub_string_for_telemetry(r"open ..\private\model.onnx") == "open [path]" + assert scrub_string_for_telemetry("request token=supersecret") == "request token=[secret]" + assert scrub_string_for_telemetry("token=first&api_key=second") == "token=[secret]" + assert scrub_string_for_telemetry("fetch example.test?access_token=top-secret") == ( + "fetch example.test?access_token=[secret]" + ) + assert scrub_string_for_telemetry("redirect#access_token=top-secret") == ("redirect#access_token=[secret]") + assert scrub_string_for_telemetry("auth.token=top-secret") == "auth.token=[secret]" + assert scrub_string_for_telemetry("auth=top-secret") == "auth=[secret]" + assert scrub_string_for_telemetry("refreshToken=top-secret") == "refreshToken=[secret]" + assert scrub_string_for_telemetry("AWS_SECRET_ACCESS_KEY=top-secret") == ("AWS_SECRET_ACCESS_KEY=[secret]") + assert scrub_string_for_telemetry("PWD=odbc-value") == "PWD=[secret]" + assert scrub_string_for_telemetry("Authorization: ******") == "Authorization: [secret]" + assert scrub_string_for_telemetry("failure --api-key top-secret") == "failure --api-key [secret]" + assert scrub_string_for_telemetry("Command ['tool', '--api-key', 'top-secret']") == ( + "Command ['tool', '--api-key', '[secret]" + ) + assert scrub_string_for_telemetry('--password "-abc"') == '--password "[secret]' + assert scrub_string_for_telemetry("Command ['tool', '--api-key', '-abc']") == ( + "Command ['tool', '--api-key', '[secret]" + ) + assert scrub_string_for_telemetry("--password -abc") == "--password -abc" + assert scrub_string_for_telemetry("model dir /_apikey:top-secret") == "model dir [path]" + assert scrub_string_for_telemetry("arg /1token=top-secret") == "arg [path]" + assert scrub_string_for_telemetry("connect /2fa_token=top-secret") == "connect [path]" + assert scrub_string_for_telemetry("failure /password:top-secret") == "failure /password:[secret]" + assert scrub_string_for_telemetry("connect user:password@localhost/model") == "connect [secret]" + assert scrub_string_for_telemetry("n/a read/write domain\\user") == "n/a read/write domain\\user" + assert scrub_string_for_telemetry("meta-llama/Llama-3.1-8B-Instruct") == "meta-llama/Llama-3.1-8B-Instruct" + assert scrub_string_for_telemetry("tokenizer=enabled oauth=enabled") == "tokenizer=enabled oauth=enabled" + assert ( + len(scrub_string_for_telemetry("x" * (MAX_TELEMETRY_STRING_LENGTH + 100)).encode("utf-8")) + == MAX_TELEMETRY_STRING_LENGTH + ) + assert scrub_string_for_telemetry("x" * (MAX_TELEMETRY_STRING_LENGTH - 1) + "€") == ( + "x" * (MAX_TELEMETRY_STRING_LENGTH - 1) + ) + redacted_at_limit = scrub_string_for_telemetry("x" * (MAX_TELEMETRY_STRING_LENGTH - 5) + " /a/b/c") + assert len(redacted_at_limit.encode("utf-8")) == MAX_TELEMETRY_STRING_LENGTH + assert redacted_at_limit.endswith("[path]") + + +def test_secret_scanner_advances_past_rejected_key_tokens(): + from olive.telemetry import telemetry_redaction + + value = "-" + "field-" * 6_800 + "value" + with patch( + "olive.telemetry.telemetry_redaction.is_sensitive_config_key_for_telemetry", + wraps=telemetry_redaction.is_sensitive_config_key_for_telemetry, + ) as is_sensitive_key: + scrubbed = telemetry_redaction.scrub_string_for_telemetry(value) + + assert scrubbed == value + assert is_sensitive_key.call_count == 1 + + +def test_error_messages_are_capped_at_40960_utf8_bytes(): + from olive.telemetry.telemetry_extensions import log_error + from olive.telemetry.telemetry_redaction import MAX_ERROR_MESSAGE_LENGTH + + telemetry = MagicMock() + with patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry): + log_error("RuntimeError", "x" * (MAX_ERROR_MESSAGE_LENGTH + 100)) + truncated = telemetry.log.call_args.args[1]["exception_message"] + assert len(truncated.encode("utf-8")) == MAX_ERROR_MESSAGE_LENGTH + + log_error("RuntimeError", "x" * (MAX_ERROR_MESSAGE_LENGTH - 1) + "€") + multibyte = telemetry.log.call_args.args[1]["exception_message"] + assert multibyte == "x" * (MAX_ERROR_MESSAGE_LENGTH - 1) + + log_error("RuntimeError", "x" * (MAX_ERROR_MESSAGE_LENGTH - 5) + " /a/b/c") + redacted_at_limit = telemetry.log.call_args.args[1]["exception_message"] + assert len(redacted_at_limit.encode("utf-8")) == MAX_ERROR_MESSAGE_LENGTH + assert redacted_at_limit.endswith("[path]") + + +def test_error_payload_preserves_error_specific_size_limit(tenv): + from olive.telemetry.telemetry_redaction import MAX_ERROR_MESSAGE_LENGTH + + telemetry = Telemetry() + _quiesce(telemetry) + payload = telemetry._build_payload( + ERROR_EVENT_NAME, + { + "exception_type": "RuntimeError", + "exception_message": "x" * (MAX_ERROR_MESSAGE_LENGTH + 100), + }, + ) + + assert len(json.loads(payload)["data"]["exceptionMessage"].encode("utf-8")) == MAX_ERROR_MESSAGE_LENGTH + + +def test_format_exception_message_redacts_paths_in_message(): + from olive.telemetry.telemetry_extensions import _format_exception_message + + exc = RuntimeError(r"failed to read C:\Users\alice\secret\weights.bin") + message = _format_exception_message(exc, exc.__traceback__) + assert "alice" not in message + assert "[path]" in message + + +def test_format_exception_message_handles_unprintable_exception(): + from olive.telemetry.telemetry_extensions import _format_exception_message + + class UnprintableError(Exception): + def __str__(self): + raise RuntimeError("cannot render") + + assert _format_exception_message(UnprintableError()).endswith("UnprintableError: ") + + +def test_public_helpers_never_propagate_failures(): + from olive.telemetry.telemetry_extensions import log_action, log_error + + with patch("olive.telemetry.telemetry_extensions._get_logger", side_effect=RuntimeError("telemetry failed")): + log_action("test", "work", 1.0, True, metadata=["not", "a", "dict"]) + log_error("RuntimeError", "boom", metadata=["not", "a", "dict"]) + + +def _raise_error_called_with_source_secret(_secret): + raise RuntimeError("boom") + + +def test_format_exception_message_omits_source_code(): + from olive.telemetry.telemetry_extensions import _format_exception_message + + try: + _raise_error_called_with_source_secret("source-secret") + except RuntimeError as ex: + message = _format_exception_message(ex, ex.__traceback__) + + assert "source-secret" not in message + assert __file__ not in message + assert 'File "[path]"' in message + assert "in _raise_error_called_with_source_secret" in message + assert message.endswith("RuntimeError: boom") + + +def test_device_id_store_uses_owner_only_creation_mode(tmp_path): + with ( + patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path), + patch.object(Path, "mkdir") as mock_mkdir, + ): + deviceid_store_mod.Store().store_id("test-device-id") + + mock_mkdir.assert_called_once_with(mode=0o700, parents=True, exist_ok=True) + + +def test_relative_cache_environment_uses_absolute_home(tmp_path, monkeypatch): + telemetry_utils.get_telemetry_base_dir.cache_clear() + monkeypatch.setenv("XDG_CACHE_HOME", "relative-cache") + monkeypatch.setenv("HOME", str(tmp_path)) + with patch.object(telemetry_utils.platform, "system", return_value="Linux"): + base_dir = telemetry_utils.get_telemetry_base_dir() + telemetry_utils.get_telemetry_base_dir.cache_clear() + + assert base_dir == tmp_path / ".cache" / telemetry_utils.ORT_SUPPORT_DIR + assert base_dir.is_absolute() + + +def test_relative_windows_appdata_uses_absolute_home(tmp_path, monkeypatch): + telemetry_utils.get_telemetry_base_dir.cache_clear() + monkeypatch.setenv("LOCALAPPDATA", "relative-local") + monkeypatch.setenv("APPDATA", "relative-roaming") + monkeypatch.setenv("HOME", "relative-home") + with ( + patch.object(telemetry_utils.platform, "system", return_value="Windows"), + patch.object(telemetry_utils.Path, "home", return_value=tmp_path), + ): + base_dir = telemetry_utils.get_telemetry_base_dir() + telemetry_utils.get_telemetry_base_dir.cache_clear() + + assert base_dir == tmp_path / "AppData" / "Local" / telemetry_utils.ORT_SUPPORT_DIR + assert base_dir.is_absolute() + + +def test_windows_telemetry_base_dir_uses_canonical_developer_tools_path(tmp_path, monkeypatch): + telemetry_utils.get_telemetry_base_dir.cache_clear() + monkeypatch.setenv("LOCALAPPDATA", str(tmp_path)) + with patch.object(telemetry_utils.platform, "system", return_value="Windows"): + base_dir = telemetry_utils.get_telemetry_base_dir() + telemetry_utils.get_telemetry_base_dir.cache_clear() + + assert base_dir == tmp_path / telemetry_utils.ORT_SUPPORT_DIR + + +def test_home_resolution_fails_without_absolute_per_user_directory(monkeypatch): + fake_pwd = SimpleNamespace(getpwuid=MagicMock(side_effect=KeyError("missing"))) + monkeypatch.setenv("HOME", "relative-home") + with ( + patch.object(telemetry_utils.platform, "system", return_value="Linux"), + patch.object(telemetry_utils.Path, "home", return_value=Path("relative-home")), + patch.dict(sys.modules, {"pwd": fake_pwd}), + pytest.raises(RuntimeError, match="No absolute per-user"), + ): + telemetry_utils._resolve_home_dir() + + +def test_device_id_is_ephemeral_when_per_user_storage_is_unavailable(): + import olive.telemetry.deviceid.deviceid as deviceid + + deviceid._device_id_state.update({"device_id": None, "status": deviceid.DeviceIdStatus.NEW}) + with ( + patch.object(deviceid.platform, "system", return_value="Linux"), + patch.object(deviceid, "Store", side_effect=RuntimeError("no home")), + ): + generated = deviceid.get_device_id() + + assert deviceid._is_valid_device_id(generated) + assert deviceid._device_id_state["status"] == deviceid.DeviceIdStatus.FAILED + + +def test_missing_device_id_raises_file_not_found(tmp_path): + with ( + patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path), + pytest.raises(FileNotFoundError), + ): + _ = deviceid_store_mod.Store().retrieve_id + + +def test_corrupted_device_id_is_atomically_repaired(tmp_path): + import olive.telemetry.deviceid.deviceid as deviceid + + device_id_path = tmp_path / "deviceid" + device_id_path.write_bytes(b"\xff\xfe") + deviceid._device_id_state.update({"device_id": None, "status": deviceid.DeviceIdStatus.NEW}) + + with ( + patch.object(deviceid.platform, "system", return_value="Linux"), + patch.object(deviceid, "get_telemetry_base_dir", return_value=tmp_path), + patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path), + ): + repaired = deviceid.get_device_id() + + assert deviceid._is_valid_device_id(repaired) + assert device_id_path.read_text(encoding="utf-8") == repaired + assert deviceid._device_id_state["status"] == deviceid.DeviceIdStatus.CORRUPTED + + +def test_file_store_does_not_overwrite_concurrent_winner(tmp_path): + with patch.object(deviceid_store_mod, "get_telemetry_base_dir", return_value=tmp_path): + store = deviceid_store_mod.Store() + assert store.store_id("first") is True + assert store.store_id("second") is False + + assert (tmp_path / "deviceid").read_text(encoding="utf-8") == "first" + + +def test_concurrent_processes_publish_one_device_id(tmp_path): + script = ( + "import platform; " + "import olive.telemetry.deviceid.deviceid as d; " + "import olive.telemetry.utils as u; " + "platform.system=lambda:'Linux'; " + "u.get_telemetry_base_dir.cache_clear(); " + "print(d.get_device_id())" + ) + env = os.environ.copy() + env["XDG_CACHE_HOME"] = str(tmp_path) + + def run_process(_): + return subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env=env, + timeout=90, + check=False, + ) + + with ThreadPoolExecutor(max_workers=6) as executor: + completed = list(executor.map(run_process, range(6))) + results = [] + for process in completed: + assert process.returncode == 0, process.stderr + results.append(process.stdout.strip()) + + assert len(set(results)) == 1 + + +@pytest.mark.skipif(os.name != "nt", reason="Windows named mutex") +def test_windows_device_id_protocol_serializes_cross_process_publication(tmp_path): + script = """ +import os +from pathlib import Path +import olive.telemetry.deviceid.deviceid as d + +class FileBackedWindowsStore: + @property + def retrieve_id(self): + path = Path(os.environ["OLIVE_TEST_DEVICE_ID"]) + if not path.exists(): + raise FileNotFoundError + return path.read_text(encoding="utf-8") + + def store_id(self, device_id, replace_existing=False): + Path(os.environ["OLIVE_TEST_DEVICE_ID"]).write_text(device_id, encoding="utf-8") + return True + +d.platform.system = lambda: "Windows" +d.WindowsStore = FileBackedWindowsStore +d.get_telemetry_base_dir = lambda: Path(os.environ["OLIVE_TEST_DEVICE_DIR"]) +print(d.get_device_id()) +""" + env = os.environ.copy() + env["OLIVE_TEST_DEVICE_DIR"] = str(tmp_path) + env["OLIVE_TEST_DEVICE_ID"] = str(tmp_path / "registry-deviceid") + + def run_process(_): + return subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env=env, + timeout=90, + check=False, + ) + + with ThreadPoolExecutor(max_workers=6) as executor: + completed = list(executor.map(run_process, range(6))) + results = [] + for process in completed: + assert process.returncode == 0, process.stderr + results.append(process.stdout.strip()) + + assert len(set(results)) == 1 + + +def test_windows_device_id_store_uses_least_privilege_access(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_QUERY_VALUE=0x0001, + KEY_SET_VALUE=0x0002, + KEY_CREATE_SUB_KEY=0x0004, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + ) + winreg.QueryValueEx.side_effect = FileNotFoundError + key_handle = object() + winreg.CreateKeyEx.return_value.__enter__.return_value = key_handle + + with patch.dict("sys.modules", {"winreg": winreg}): + deviceid_store_mod.WindowsStore().store_id("test-device-id") + + winreg.CreateKeyEx.assert_called_once_with( + winreg.HKEY_CURRENT_USER, + deviceid_store_mod.REGISTRY_PATH, + reserved=0, + access=winreg.KEY_QUERY_VALUE | winreg.KEY_SET_VALUE | winreg.KEY_CREATE_SUB_KEY | winreg.KEY_WOW64_64KEY, + ) + winreg.SetValueEx.assert_called_once_with( + key_handle, + deviceid_store_mod.REGISTRY_KEY, + 0, + winreg.REG_SZ, + "test-device-id", + ) + + +def test_windows_device_id_store_preserves_existing_value(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_QUERY_VALUE=0x0001, + KEY_SET_VALUE=0x0002, + KEY_CREATE_SUB_KEY=0x0004, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + ) + winreg.QueryValueEx.return_value = ("existing-device-id", winreg.REG_SZ) + + with patch.dict("sys.modules", {"winreg": winreg}): + deviceid_store_mod.WindowsStore().store_id("new-device-id") + + winreg.SetValueEx.assert_not_called() + + +def test_windows_device_id_store_replaces_existing_value_when_requested(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_QUERY_VALUE=0x0001, + KEY_SET_VALUE=0x0002, + KEY_CREATE_SUB_KEY=0x0004, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + ) + key_handle = object() + winreg.CreateKeyEx.return_value.__enter__.return_value = key_handle + + with patch.dict("sys.modules", {"winreg": winreg}): + deviceid_store_mod.WindowsStore().store_id("new-device-id", replace_existing=True) + + winreg.QueryValueEx.assert_not_called() + winreg.SetValueEx.assert_called_once_with( + key_handle, + deviceid_store_mod.REGISTRY_KEY, + 0, + winreg.REG_SZ, + "new-device-id", + ) + + +def test_windows_device_id_store_rejects_wrong_registry_type(): + winreg = MagicMock( + HKEY_CURRENT_USER=object(), + KEY_READ=0x0001, + KEY_WOW64_64KEY=0x0100, + REG_SZ=1, + REG_BINARY=3, + ) + winreg.QueryValueEx.return_value = (b"not-a-string", winreg.REG_BINARY) + + with ( + patch.dict("sys.modules", {"winreg": winreg}), + pytest.raises(ValueError, match="not a string"), + ): + _ = deviceid_store_mod.WindowsStore().retrieve_id + + +def test_nested_actions_log_error_once(): + from olive.telemetry.telemetry_extensions import action + + telemetry = MagicMock(accepts_detailed_events=True) + + @action + @action + def fail(): + raise ValueError("boom") + + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions.log_error") as mock_log_error, + pytest.raises(ValueError, match="boom"), + ): + fail() + + mock_log_error.assert_called_once() + + +def test_positional_function_uses_function_action_name(): + from olive.telemetry.telemetry_extensions import action + + telemetry = MagicMock(accepts_detailed_events=True) + + @action + def work(value): + return value + + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.log_action") as mock_log_action, + ): + assert work("value") == "value" + + assert mock_log_action.call_args.kwargs["action_name"] == "work" + + +def test_action_context_without_start_time_reports_zero_duration(): + from olive.telemetry.telemetry_extensions import ActionContext + + telemetry = MagicMock(accepts_detailed_events=True) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from", return_value="test"), + patch("olive.telemetry.telemetry_extensions.time.perf_counter", return_value=100.0), + patch("olive.telemetry.telemetry_extensions.log_action") as mock_log_action, + ): + context = ActionContext("work") + context.__exit__(None, None, None) + + assert mock_log_action.call_args.kwargs["duration_ms"] == 0 + + +def test_disabled_action_skips_stack_inspection(): + from olive.telemetry.telemetry_extensions import action + + telemetry = MagicMock(accepts_detailed_events=False) + + @action + def work(): + return 42 + + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from") as mock_resolve, + ): + assert work() == 42 + + mock_resolve.assert_not_called() + + +def test_disabled_action_context_skips_stack_inspection(): + from olive.telemetry.telemetry_extensions import ActionContext + + telemetry = MagicMock(accepts_detailed_events=False) + with ( + patch("olive.telemetry.telemetry_extensions._get_logger", return_value=telemetry), + patch("olive.telemetry.telemetry_extensions._resolve_invoked_from") as mock_resolve, + patch("olive.telemetry.telemetry_extensions.log_action") as mock_log_action, + ActionContext("work"), + ): + result = 42 + + assert result == 42 + mock_resolve.assert_not_called() + mock_log_action.assert_not_called() diff --git a/test/utils.py b/test/utils.py index 778f472e55..a7d908c1dd 100644 --- a/test/utils.py +++ b/test/utils.py @@ -398,6 +398,8 @@ def get_wikitext_data_config( load_dataset_config={ "data_name": "Salesforce/wikitext", "subset": "wikitext-2-raw-v1", + # intentionally a small slice: these tests only need max_samples=1, so reading the full + # split would just slow them down "split": "train[:1000]", }, pre_process_data_config={ diff --git a/test/workflows/test_workflow_run.py b/test/workflows/test_workflow_run.py index 82cc4980bf..c482a69ace 100644 --- a/test/workflows/test_workflow_run.py +++ b/test/workflows/test_workflow_run.py @@ -1,10 +1,20 @@ +# pylint: disable=protected-access +import json import sys from copy import deepcopy from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest +from olive.telemetry.recipe_telemetry import ( + _NO_OVERRIDE, + _build_recipe_hash, + _classify_input_model_source, + _classify_run_config_source, + _extract_config_overrides, + _sanitize_config_snapshot, +) from olive.workflows import run as olive_run from test.utils import ( get_pytorch_model, @@ -125,3 +135,432 @@ def test_run_packages(): # cleanup requirements_file_path.unlink() + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +@patch("olive.telemetry.recipe_telemetry.is_ci_environment", return_value=False) +def test_run_logs_recipe_result_success(_, mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [{"device": "gpu", "execution_providers": ["CUDAExecutionProvider"]}], + } + }, + "engine": {"target": "local_system"}, + "passes": {"dynamic_quant": {"type": "OnnxDynamicQuantization"}}, + } + expected_output = object() + mock_run_engine.return_value = expected_output + + output = olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + assert output is expected_output + mock_log_recipe_result.assert_called_once() + assert mock_log_recipe_result.call_args.args[0] == "Quantize" + assert mock_log_recipe_result.call_args.kwargs["success"] is True + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert metadata["recipe_command"] == "Quantize" + assert metadata["recipe_source"] == "generated_cli" + assert metadata["recipe_format"] == "generated" + assert metadata["workflow_id"] == "default_workflow" + assert metadata["input_model_type"] == "hfmodel" + assert metadata["input_model_source"] == "string_name" + assert metadata["model_task"] == "text-generation" + assert metadata["target_system_type"] == "LocalSystem" + assert metadata["target_device"] == "gpu" + assert metadata["target_execution_provider"] == "CUDAExecutionProvider" + assert metadata["target_execution_providers"] == "CUDAExecutionProvider" + assert metadata["host_system_type"] == "LocalSystem" + assert "host_device" not in metadata + assert "host_execution_provider" not in metadata + assert "host_execution_providers" not in metadata + assert metadata["pass_types"] == "onnxdynamicquantization" + assert metadata["pass_count"] == 1 + assert metadata["data_config_count"] == 0 + assert metadata["search_enabled"] is False + assert metadata["package_config_provided"] is False + assert metadata["is_ci"] is False + assert metadata["recipe_hash"] + assert "input_model_name_hash" not in metadata + assert "config_overrides" not in metadata + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_config_overrides_when_recipe_metadata_provides_overrides(mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + } + } + mock_run_engine.return_value = object() + + olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "WorkflowRun", + "config_overrides": { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + }, + "engine": {"target": "local_system"}, + "data_path": Path("data"), + }, + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + config_overrides = json.loads(metadata["config_overrides"]) + assert config_overrides["input_model"]["model_path"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert config_overrides["engine"]["target"] == "" + assert config_overrides["data_path"] == "" + + +def test_missing_baseline_keys_are_not_reported_as_overrides(): + value = {"passes": {"optimize": {"type": "OrtTransformersOptimization"}}} + baseline = { + "passes": {"optimize": {"type": "OrtTransformersOptimization"}}, + "extra_dependencies": {"cpu": ["onnxruntime"]}, + } + + assert _extract_config_overrides(value, baseline) is _NO_OVERRIDE + + +def test_config_snapshot_redacts_unknown_object_type_names(): + class ProjectSpecificObject: + pass + + assert _sanitize_config_snapshot({"value": ProjectSpecificObject()}) == {"value": ""} + + +def test_recipe_hash_does_not_fingerprint_credentials_or_environment_values(): + first = { + "input_model": {"type": "HfModel", "model_path": "Qwen/Qwen2.5-0.5B-Instruct"}, + "environmentVariables": {"API_KEY": "first-secret"}, + "service": {"clientSecret": "first-client-secret"}, + "auth": "first-auth-secret", + "auth_header": "first-header-secret", + "credential_value": "first-credential-secret", + "docker_env": {"TOKEN": "first-env-secret"}, + } + second = deepcopy(first) + second["environmentVariables"]["API_KEY"] = "second-secret" + second["service"]["clientSecret"] = "second-client-secret" + second["auth"] = "second-auth-secret" + second["auth_header"] = "second-header-secret" + second["credential_value"] = "second-credential-secret" + second["docker_env"]["TOKEN"] = "second-env-secret" + + assert _build_recipe_hash(first) == _build_recipe_hash(second) + + +def test_recipe_hash_does_not_fingerprint_paths_under_generic_keys(): + first = {"script_args": [r"C:\Users\Alice\private.json"], "label": "same"} + second = {"script_args": ["/home/bob/private.json"], "label": "same"} + + assert _build_recipe_hash(first) == _build_recipe_hash(second) + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_recipe_metadata_redacts_credentials_from_preformatted_overrides(mock_run_engine, mock_log_recipe_result): + mock_run_engine.return_value = object() + olive_run( + {"input_model": {"type": "HfModel", "model_path": "Qwen/Qwen2.5-0.5B-Instruct"}}, + recipe_telemetry_metadata={ + "config_overrides": json.dumps( + { + "auth": "token=private", + "endpoint": "https://private.example/model", + "environmentVariables": {"API_KEY": "top-secret"}, + "service": { + "apiKey": "api-secret", + "aws_access_key_id": "aws-secret", + "batch_size": 1, + "clientSecret": "client-secret", + "serviceCredential": "credential-secret", + }, + "modelPath": "private/model.onnx", + } + ), + "package_config_overrides": json.dumps({"feed": {"access_token": "package-secret"}}), + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + serialized = metadata["config_overrides"] + overrides = json.loads(serialized) + assert overrides["environmentVariables"] == "" + assert overrides["auth"] == "" + assert overrides["endpoint"] == "[path]" + assert overrides["modelPath"] == "" + assert overrides["service"] == { + "apiKey": "", + "aws_access_key_id": "", + "batch_size": 1, + "clientSecret": "", + "serviceCredential": "", + } + assert "top-secret" not in serialized + for secret in ("api-secret", "aws-secret", "client-secret", "credential-secret", "private/model.onnx"): + assert secret not in serialized + assert json.loads(metadata["package_config_overrides"]) == {"feed": {"access_token": ""}} + assert "package-secret" not in metadata["package_config_overrides"] + + +@patch("olive.workflows.run.run.log_error") +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_recipe_result_failure(mock_run_engine, mock_log_recipe_result, mock_log_error): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + }, + "passes": {"dynamic_quant": {"type": "OnnxDynamicQuantization"}}, + } + mock_run_engine.side_effect = ValueError("recipe failed") + + with pytest.raises(ValueError, match="recipe failed"): + olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + mock_log_recipe_result.assert_called_once() + assert mock_log_recipe_result.call_args.args[0] == "Quantize" + assert mock_log_recipe_result.call_args.kwargs["success"] is False + assert "exception_type" not in mock_log_recipe_result.call_args.kwargs + mock_log_error.assert_called_once() + assert mock_log_error.call_args.kwargs["exception_type"] == "ValueError" + assert "recipe failed" in mock_log_error.call_args.kwargs["exception_message"] + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_skips_recipe_result_when_recipe_telemetry_is_not_emitted(mock_run_engine, mock_log_recipe_result): + expected_output = object() + mock_run_engine.return_value = expected_output + + output = olive_run( + { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + } + }, + emit_recipe_telemetry=False, + ) + + assert output is expected_output + mock_log_recipe_result.assert_not_called() + + +@patch("olive.workflows.run.run.is_ci_environment", return_value=True) +@patch("olive.workflows.run.run.Telemetry.get_existing_instance") +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_programmatic_ci_run_flushes_recipe_with_bounded_shutdown( + mock_run_engine, mock_log_recipe_result, mock_get_telemetry, _mock_is_ci +): + expected_output = object() + telemetry = Mock() + mock_run_engine.return_value = expected_output + mock_get_telemetry.return_value = telemetry + + output = olive_run( + { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + } + } + ) + + assert output is expected_output + mock_log_recipe_result.assert_called_once() + telemetry.shutdown.assert_called_once_with(timeout_millis=2_000, callback_timeout_millis=2_000) + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.systems.system_config.SystemConfig.create_system") +def test_run_logs_single_parent_recipe_result_for_docker_host(mock_create_system, mock_log_recipe_result): + expected_output = object() + docker_system = Mock() + + def run_workflow(container_run_config): + container_run_config.engine.host = container_run_config.engine.target + return expected_output + + docker_system.run_workflow.side_effect = run_workflow + mock_create_system.return_value = docker_system + config = { + "input_model": {"type": "ONNXModel", "model_path": "model.onnx"}, + "systems": { + "docker_system": { + "type": "Docker", + "config": { + "dockerfile": "Dockerfile", + "build_context_path": "build_context", + "image_name": "test-image:latest", + "work_dir": "/olive-ws", + }, + }, + "local_system": {"type": "LocalSystem"}, + }, + "engine": {"host": "docker_system", "target": "local_system"}, + } + + output = olive_run(config) + + assert output is expected_output + mock_log_recipe_result.assert_called_once() + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert metadata["host_system_type"] == "Docker" + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_recipe_host_metadata_without_explicit_target(mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + }, + "systems": { + "host_system": { + "type": "LocalSystem", + "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], + } + }, + "engine": {"host": "host_system"}, + } + mock_run_engine.return_value = object() + + olive_run( + config, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert "target_system_type" not in metadata + assert "target_device" not in metadata + assert "target_execution_provider" not in metadata + assert "target_execution_providers" not in metadata + assert metadata["host_system_type"] == "LocalSystem" + assert metadata["host_device"] == "cpu" + assert metadata["host_execution_provider"] == "CPUExecutionProvider" + assert metadata["host_execution_providers"] == "CPUExecutionProvider" + + +@patch("olive.workflows.run.run.log_recipe_result") +@patch("olive.workflows.run.run.run_engine") +def test_run_logs_package_config_overrides_when_package_config_provided(mock_run_engine, mock_log_recipe_result): + config = { + "input_model": { + "type": "HfModel", + "model_path": "Qwen/Qwen2.5-0.5B-Instruct", + "task": "text-generation", + "load_kwargs": {"attn_implementation": "eager"}, + } + } + mock_run_engine.return_value = object() + + olive_run( + config, + package_config={ + "passes": { + "AddOliveMetadata": { + "module_path": "olive.passes.onnx.add_metadata.AddOliveMetadata", + "supported_providers": ["CPUExecutionProvider"], + } + }, + "extra_dependencies": {"custom_accelerator": ["custom-package"]}, + }, + recipe_telemetry_metadata={ + "recipe_name": "Quantize", + "recipe_command": "Quantize", + "recipe_source": "generated_cli", + "recipe_format": "generated", + }, + ) + + metadata = mock_log_recipe_result.call_args.kwargs["metadata"] + assert metadata["package_config_provided"] is True + package_config_overrides = json.loads(metadata["package_config_overrides"]) + assert package_config_overrides["passes"][0]["supported_providers"] == ["CPUExecutionProvider"] + assert "module_path" not in package_config_overrides["passes"][0] + assert package_config_overrides["extra_dependencies"]["custom_accelerator"] == ["custom-package"] + + +def test_classify_run_config_source_handles_non_pathlike_object(): + assert _classify_run_config_source(object()) == ("config_object", "object") + + +def test_classify_input_model_source_does_not_depend_on_local_filesystem(tmp_path, monkeypatch): + assert _classify_input_model_source("Qwen/Qwen2.5-0.5B-Instruct") == "string_name" + + monkeypatch.chdir(tmp_path) + (tmp_path / "bert-base-uncased").mkdir() + + assert _classify_input_model_source("bert-base-uncased") == "string_name" + assert _classify_input_model_source("./model.onnx") == "local_file" + assert _classify_input_model_source("model.onnx") == "local_file" + + +def test_recipe_hash_does_not_depend_on_local_model_path_presence(tmp_path, monkeypatch): + config = { + "input_model": {"type": "HfModel", "config": {"model_path": "bert-base-uncased"}}, + "engine": {"output_dir": "output"}, + } + recipe_hash = _build_recipe_hash(config) + + monkeypatch.chdir(tmp_path) + (tmp_path / "bert-base-uncased").mkdir() + + assert _build_recipe_hash(config) == recipe_hash + + +def test_recipe_hash_handles_path_values(): + config = { + "input_model": {"type": "HfModel", "config": {"model_path": Path("model")}}, + "custom_value": Path("custom"), + } + + assert _build_recipe_hash(config)