diff --git a/application/backend/app/api/routers/jobs.py b/application/backend/app/api/routers/jobs.py index 50db5cd8665..e93912791d9 100644 --- a/application/backend/app/api/routers/jobs.py +++ b/application/backend/app/api/routers/jobs.py @@ -101,6 +101,7 @@ async def submit_job( job_id=job_id, max_calibration_subset_size=job_request.parameters.max_calibration_subset_size, max_drop=job_request.parameters.max_drop, + max_num_iterations=job_request.parameters.max_num_iterations, ), ) case JobType.PREPARE_DATASET_FOR_IMPORT: diff --git a/application/backend/app/api/schemas/jobs/quantization.py b/application/backend/app/api/schemas/jobs/quantization.py index 2f457964ca2..ab8eb5756bf 100644 --- a/application/backend/app/api/schemas/jobs/quantization.py +++ b/application/backend/app/api/schemas/jobs/quantization.py @@ -26,6 +26,12 @@ class QuantizationRequestParams(BaseModel): description="Maximum allowed accuracy drop, expressed as a percentage in the range [0-1] " "where e.g. 0.03 means 3%", ) + max_num_iterations: int | None = Field( + 10, + gt=0, + description="Maximum number of iterations for accuracy-aware quantization. " + "A value of null means unlimited. Only used when max_drop is set.", + ) model_config = { "json_schema_extra": { @@ -33,6 +39,7 @@ class QuantizationRequestParams(BaseModel): "model_id": "6b7bb928-5d6f-46ea-8fd2-5ce80dd1e12b", "max_calibration_subset_size": 100, "max_drop": 0.01, + "max_num_iterations": 10, } } } @@ -53,6 +60,7 @@ class QuantizationRequest(BaseJobRequest): "model_id": "6b7bb928-5d6f-46ea-8fd2-5ce80dd1e12b", "max_calibration_subset_size": 100, "max_drop": 0.01, + "max_num_iterations": 10, }, } } @@ -85,6 +93,9 @@ class QuantizationMetadata(BaseModel): model_variant: QuantizationModelVariantMetadata = Field(..., description="Model variant created by quantization") max_calibration_subset_size: int = Field(..., description="Maximum calibration subset size") max_drop: float | None = Field(None, description="Maximum allowed accuracy drop") + max_num_iterations: int | None = Field( + None, description="Maximum number of iterations for accuracy-aware quantization (null means unlimited)" + ) @model_validator(mode="before") @classmethod @@ -96,5 +107,6 @@ def populate_metadata(cls, data: object) -> object: "model_variant": QuantizationModelVariantMetadata(id=data.params.model_variant_id), "max_calibration_subset_size": data.params.max_calibration_subset_size, "max_drop": data.params.max_drop, + "max_num_iterations": data.params.max_num_iterations, } return data diff --git a/application/backend/app/execution/quantization/getitune_quantizer.py b/application/backend/app/execution/quantization/getitune_quantizer.py index ca26f63da65..af534cf4466 100644 --- a/application/backend/app/execution/quantization/getitune_quantizer.py +++ b/application/backend/app/execution/quantization/getitune_quantizer.py @@ -238,6 +238,7 @@ def run_quantization( ov_engine: OVEngine, subset_size: int, max_drop: float | None = None, + max_num_iterations: int | None = None, ) -> Path: """Execute the quantization process using nncf.quantize() via OVEngine.optimize(). @@ -245,12 +246,18 @@ def run_quantization( ov_engine: The OVEngine instance. subset_size: Maximum calibration subset size. max_drop: Optional maximum accuracy drop for accuracy-aware quantization. + max_num_iterations: Optional maximum number of iterations for accuracy-aware + quantization. ``None`` means unlimited. Only used when ``max_drop`` is set. Returns: Path to the quantized model XML file. """ logger.info("Running quantization with max_calibration_subset_size={}", subset_size) - quantized_model_path = ov_engine.optimize(max_data_subset_size=subset_size, max_drop=max_drop) + quantized_model_path = ov_engine.optimize( + max_data_subset_size=subset_size, + max_drop=max_drop, + max_num_iterations=max_num_iterations, + ) logger.info("Quantization completed. Model saved at {}", quantized_model_path) return quantized_model_path @@ -328,16 +335,23 @@ def execute(self, params: QuantizationJobParams) -> None: datamodule = self.prepare_calibration_dataset(params=params, model=model) ov_engine = self.initialize_engine(params=params, model=model, datamodule=datamodule) + + # max_num_iterations only applies to accuracy-aware quantization (max_drop set). + # For standard PTQ it must always be None regardless of the configured value. + max_num_iterations = params.max_num_iterations if params.max_drop is not None else None + quantized_model_path = self.run_quantization( ov_engine=ov_engine, subset_size=params.max_calibration_subset_size, max_drop=params.max_drop, + max_num_iterations=max_num_iterations, ) quantization_info = { "type": "PTQ" if params.max_drop is None else "Accuracy-aware PTQ", "max_calibration_subset_size": params.max_calibration_subset_size, "max_drop": params.max_drop, + "max_num_iterations": max_num_iterations, } with self._db_session_factory() as db: diff --git a/application/backend/app/models/jobs/quantization_job.py b/application/backend/app/models/jobs/quantization_job.py index 1f18a0e7081..3ba72baa178 100644 --- a/application/backend/app/models/jobs/quantization_job.py +++ b/application/backend/app/models/jobs/quantization_job.py @@ -18,6 +18,10 @@ class QuantizationJobParams(JobParams): model_variant_id: UUID = Field(default_factory=uuid4) max_calibration_subset_size: int = Field(default=100, description="Max samples for calibration") max_drop: float | None = Field(default=None, description="Max accuracy drop for accuracy-aware quantization") + max_num_iterations: int | None = Field( + default=10, + description="Max number of iterations for accuracy-aware quantization (None means unlimited)", + ) class QuantizationJob(ProjectJob[QuantizationJobParams]): diff --git a/application/backend/tests/unit/api/routers/test_jobs.py b/application/backend/tests/unit/api/routers/test_jobs.py index 7fcb58740db..46947f5da4a 100644 --- a/application/backend/tests/unit/api/routers/test_jobs.py +++ b/application/backend/tests/unit/api/routers/test_jobs.py @@ -129,6 +129,7 @@ def test_submit_quantize_job(self, tmp_path, fxt_client, fxt_jobs_queue, fxt_pro "model_id": model_id, "max_calibration_subset_size": 200, "max_drop": 0.01, + "max_num_iterations": 5, }, } ) @@ -145,6 +146,7 @@ def test_submit_quantize_job(self, tmp_path, fxt_client, fxt_jobs_queue, fxt_pro assert submitted_job.params.model_id == model_id assert submitted_job.params.max_calibration_subset_size == 200 assert submitted_job.params.max_drop == 0.01 + assert submitted_job.params.max_num_iterations == 5 assert submitted_job.project_id == project.id def test_submit_quantize_job_defaults(self, tmp_path, fxt_client, fxt_jobs_queue, fxt_project_service): @@ -177,6 +179,7 @@ def test_submit_quantize_job_defaults(self, tmp_path, fxt_client, fxt_jobs_queue assert submitted_job.params.model_id == model_id assert submitted_job.params.max_calibration_subset_size == 100 assert submitted_job.params.max_drop is None + assert submitted_job.params.max_num_iterations == 10 def test_list_jobs(self, fxt_client, fxt_jobs_queue, fxt_job): fxt_jobs_queue.list_all.return_value = [fxt_job(), fxt_job()] diff --git a/application/backend/tests/unit/execution/quantization/test_getitune_quantizer.py b/application/backend/tests/unit/execution/quantization/test_getitune_quantizer.py index 04376676b6d..215395d5b72 100644 --- a/application/backend/tests/unit/execution/quantization/test_getitune_quantizer.py +++ b/application/backend/tests/unit/execution/quantization/test_getitune_quantizer.py @@ -343,6 +343,7 @@ def test_run_quantization_standard_ptq( mock_engine.optimize.assert_called_once_with( max_data_subset_size=150, max_drop=None, + max_num_iterations=None, ) def test_run_quantization_accuracy_aware( @@ -359,12 +360,14 @@ def test_run_quantization_accuracy_aware( ov_engine=mock_engine, subset_size=200, max_drop=0.01, + max_num_iterations=5, ) assert result == expected_path mock_engine.optimize.assert_called_once_with( max_data_subset_size=200, max_drop=0.01, + max_num_iterations=5, ) def test_run_quantization_propagates_engine_error( @@ -869,6 +872,7 @@ def test_execute_standard_ptq( mock_engine.optimize.assert_called_once_with( max_data_subset_size=50, max_drop=None, + max_num_iterations=None, ) fxt_model_service.create_variant.assert_called_once_with( model_revision_id=model_id, @@ -878,6 +882,7 @@ def test_execute_standard_ptq( "type": "PTQ", "max_calibration_subset_size": 50, "max_drop": None, + "max_num_iterations": None, }, model_variant_id=model_variant_id, ) @@ -927,6 +932,7 @@ def test_execute_accuracy_aware_ptq( model_variant_id=model_variant_id, max_calibration_subset_size=200, max_drop=0.02, + max_num_iterations=5, ) model = _make_model_revision(model_id) @@ -981,11 +987,13 @@ def test_execute_accuracy_aware_ptq( mock_engine.optimize.assert_called_once_with( max_data_subset_size=200, max_drop=0.02, + max_num_iterations=5, ) fxt_model_service.create_variant.assert_called_once() call_kwargs = fxt_model_service.create_variant.call_args.kwargs assert call_kwargs["quantization_info"]["type"] == "Accuracy-aware PTQ" assert call_kwargs["quantization_info"]["max_drop"] == 0.02 + assert call_kwargs["quantization_info"]["max_num_iterations"] == 5 def test_execute_cleans_workspace_on_failure( self, diff --git a/application/ui/package-lock.json b/application/ui/package-lock.json index 256d72410d0..6c04e395da6 100644 --- a/application/ui/package-lock.json +++ b/application/ui/package-lock.json @@ -3088,6 +3088,10 @@ "resolved": "packages/smart-tools", "link": true }, + "node_modules/@geti/ui": { + "resolved": "packages/ui", + "link": true + }, "node_modules/@humanfs/core": { "version": "0.19.1", "resolved": "https://registry.npmjs.org/@humanfs/core/-/core-0.19.1.tgz", @@ -5492,7 +5496,6 @@ "integrity": "sha512-o4PXJQidqJl82ckFaXUeoAW+XysPLauYI43Abki5hABd853iMhitooc6znOnczgbTYmEP6U6/y1ZyKAIsvMKGg==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@babel/code-frame": "^7.10.4", "@babel/runtime": "^7.12.5", @@ -5593,8 +5596,7 @@ "resolved": "https://registry.npmjs.org/@types/aria-query/-/aria-query-5.0.4.tgz", "integrity": "sha512-rfT93uj5s0PRL7EzccGMs3brplhcrghnDoV26NqKhCAS1hVo+WdNsPvE/yb6ilfr5hi2MEk6d5EWJTKdxg8jVw==", "dev": true, - "license": "MIT", - "peer": true + "license": "MIT" }, "node_modules/@types/babel__core": { "version": "7.20.5", @@ -8175,8 +8177,7 @@ "resolved": "https://registry.npmjs.org/dom-accessibility-api/-/dom-accessibility-api-0.5.16.tgz", "integrity": "sha512-X7BJ2yElsnOJ30pZF4uIIDfBEVgF4XEBxL9Bxhy6dnrm5hkzqmsWHGTiHqRiITNhMyFLyAiWndIJP7Z1NTteDg==", "dev": true, - "license": "MIT", - "peer": true + "license": "MIT" }, "node_modules/dom-helpers": { "version": "5.2.1", @@ -10766,7 +10767,6 @@ "integrity": "sha512-h5bgJWpxJNswbU7qCrV0tIKQCaS3blPDrqKWx+QxzuzL1zGUzij9XCWLrSLsJPu5t+eWA/ycetzYAO5IOMcWAQ==", "dev": true, "license": "MIT", - "peer": true, "bin": { "lz-string": "bin/bin.js" } @@ -11725,7 +11725,6 @@ "integrity": "sha512-Qb1gy5OrP5+zDf2Bvnzdl3jsTf1qXVMazbvCoKhtKqVs4/YK4ozX4gKQJJVyNe+cajNPn0KoC0MC3FUmaHWEmQ==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "ansi-regex": "^5.0.1", "ansi-styles": "^5.0.0", @@ -11741,7 +11740,6 @@ "integrity": "sha512-Cxwpt2SfTzTtXcfOlzGEee8O+c+MmUgGrNiBcXnuWxuFJHe6a5Hz7qwhwe5OgaSYI0IJvkLqWX1ASG+cJOkEiA==", "dev": true, "license": "MIT", - "peer": true, "engines": { "node": ">=10" }, @@ -11883,6 +11881,16 @@ "react-dom": "^16.8.0 || ^17.0.0-rc.1 || ^18.0.0 || ^19.0.0-rc.1" } }, + "node_modules/react-colorful": { + "version": "5.7.0", + "resolved": "https://registry.npmjs.org/react-colorful/-/react-colorful-5.7.0.tgz", + "integrity": "sha512-fuesYIemttah97XmsIHmz4OORDHiSFzyc9HMAIrCHJou2jaRQmL8cFJ76K4zQhhj8jzwOBlOi4BaGTjjOZCfTg==", + "license": "MIT", + "peerDependencies": { + "react": ">=16.8.0", + "react-dom": ">=16.8.0" + } + }, "node_modules/react-cool-virtual": { "version": "0.7.0", "resolved": "https://registry.npmjs.org/react-cool-virtual/-/react-cool-virtual-0.7.0.tgz", @@ -11940,8 +11948,7 @@ "version": "17.0.2", "resolved": "https://registry.npmjs.org/react-is/-/react-is-17.0.2.tgz", "integrity": "sha512-w2GsyukL62IJnlaff/nRegPQR94C/XXamvMWmSHRJ4y7Ts/4ocGRmTHvOs8PSE6pB3dWOrD/nueuU5sduBsQ4w==", - "license": "MIT", - "peer": true + "license": "MIT" }, "node_modules/react-redux": { "version": "9.2.0", @@ -14855,7 +14862,8 @@ }, "packages/config": { "name": "@geti/config", - "version": "1.0.0" + "version": "1.0.0", + "dev": true }, "packages/smart-tools": { "name": "@geti/smart-tools", @@ -14985,6 +14993,127 @@ "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", "dev": true, "license": "MIT" + }, + "packages/ui": { + "name": "@geti/ui", + "version": "1.0.0", + "dependencies": { + "@adobe/react-spectrum": "^3.43.0", + "@spectrum-icons/illustrations": "^3.6.24", + "@spectrum-icons/workflow": "^4.2.23", + "lodash-es": "^4.18.1", + "react-aria-components": "^1.11.0", + "react-colorful": "^5.6.1", + "sonner": "^2.0.7", + "usehooks-ts": "^3.1.1" + }, + "devDependencies": { + "@geti/config": "*", + "@testing-library/dom": "^10.4.1", + "@testing-library/jest-dom": "^6.6.4", + "@testing-library/react": "^16.3.0", + "@testing-library/user-event": "^14.6.1", + "@types/lodash-es": "^4.17.12", + "eslint": "^9.25.1", + "eslint-plugin-header": "^3.1.1", + "eslint-plugin-import": "^2.32.0", + "eslint-plugin-jsx-a11y": "^6.10.2", + "eslint-plugin-playwright": "^2.2.0", + "eslint-plugin-react": "^7.37.5", + "eslint-plugin-react-hooks": "^5.2.0", + "typescript": "^5.8.3" + } + }, + "packages/ui/node_modules/brace-expansion": { + "version": "1.1.15", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.15.tgz", + "integrity": "sha512-EwOCDEex4quD37XhqM3omwtMoJjr//isUZz1JopUNWms+4Z2ViyM/k1YIRePpoVNnQhENnxtFjLaxNHrT7xIUg==", + "dev": true, + "license": "MIT", + "dependencies": { + "balanced-match": "^1.0.0", + "concat-map": "0.0.1" + } + }, + "packages/ui/node_modules/debug": { + "version": "3.2.7", + "resolved": "https://registry.npmjs.org/debug/-/debug-3.2.7.tgz", + "integrity": "sha512-CFjzYYAi4ThfiQvizrFQevTTXHtnCqWfe7x1AhgEscTz6ZbLbfoLRLPugTQyBth6f8ZERVUSyWHFD/7Wu4t1XQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "ms": "^2.1.1" + } + }, + "packages/ui/node_modules/eslint-plugin-import": { + "version": "2.32.0", + "resolved": "https://registry.npmjs.org/eslint-plugin-import/-/eslint-plugin-import-2.32.0.tgz", + "integrity": "sha512-whOE1HFo/qJDyX4SnXzP4N6zOWn79WhnCUY/iDR0mPfQZO8wcYE4JClzI2oZrhBnnMUCBCHZhO6VQyoBU95mZA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@rtsao/scc": "^1.1.0", + "array-includes": "^3.1.9", + "array.prototype.findlastindex": "^1.2.6", + "array.prototype.flat": "^1.3.3", + "array.prototype.flatmap": "^1.3.3", + "debug": "^3.2.7", + "doctrine": "^2.1.0", + "eslint-import-resolver-node": "^0.3.9", + "eslint-module-utils": "^2.12.1", + "hasown": "^2.0.2", + "is-core-module": "^2.16.1", + "is-glob": "^4.0.3", + "minimatch": "^3.1.2", + "object.fromentries": "^2.0.8", + "object.groupby": "^1.0.3", + "object.values": "^1.2.1", + "semver": "^6.3.1", + "string.prototype.trimend": "^1.0.9", + "tsconfig-paths": "^3.15.0" + }, + "engines": { + "node": ">=4" + }, + "peerDependencies": { + "eslint": "^2 || ^3 || ^4 || ^5 || ^6 || ^7.2.0 || ^8 || ^9" + } + }, + "packages/ui/node_modules/eslint-plugin-react-hooks": { + "version": "5.2.0", + "resolved": "https://registry.npmjs.org/eslint-plugin-react-hooks/-/eslint-plugin-react-hooks-5.2.0.tgz", + "integrity": "sha512-+f15FfK64YQwZdJNELETdn5ibXEUQmW1DZL6KXhNnc2heoy/sg9VJJeT7n8TlMWouzWqSWavFkIhHyIbIAEapg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=10" + }, + "peerDependencies": { + "eslint": "^3.0.0 || ^4.0.0 || ^5.0.0 || ^6.0.0 || ^7.0.0 || ^8.0.0-0 || ^9.0.0" + } + }, + "packages/ui/node_modules/minimatch": { + "version": "3.1.5", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.5.tgz", + "integrity": "sha512-VgjWUsnnT6n+NUk6eZq77zeFdpW2LWDzP6zFGrCbHXiYNul5Dzqk2HHQ5uFH2DNW5Xbp8+jVzaeNt94ssEEl4w==", + "dev": true, + "license": "ISC", + "dependencies": { + "brace-expansion": "^1.1.7" + }, + "engines": { + "node": "*" + } + }, + "packages/ui/node_modules/semver": { + "version": "6.3.1", + "resolved": "https://registry.npmjs.org/semver/-/semver-6.3.1.tgz", + "integrity": "sha512-BR7VvDCVHO+q2xBEWskxS6DJE1qRnb7DxzUrogb71CWoSficBxYsiAGd+Kl0mmq/MprG9yArRkyrQxTO6XjMzA==", + "dev": true, + "license": "ISC", + "bin": { + "semver": "bin/semver.js" + } } } } diff --git a/application/ui/src/features/models/model-listing/model-variants/quantization-dialog/quantization-dialog.component.tsx b/application/ui/src/features/models/model-listing/model-variants/quantization-dialog/quantization-dialog.component.tsx index 8dff3ad46c3..898edbe7f09 100644 --- a/application/ui/src/features/models/model-listing/model-variants/quantization-dialog/quantization-dialog.component.tsx +++ b/application/ui/src/features/models/model-listing/model-variants/quantization-dialog/quantization-dialog.component.tsx @@ -14,6 +14,7 @@ import { CalibrationDatasetSizeField, DEFAULT_QUANTIZATION_PARAMETERS, MaxAccuracyDropField, + MaxNumIterationsField, } from './quantization-fields.component'; const useDatasetItemsCount = () => { @@ -42,6 +43,7 @@ export const QuantizationDialog = ({ modelId, onClose }: QuantizationDialogProps const [hasNoMaxAccuracyDrop, setHasNoMaxAccuracyDrop] = useState( DEFAULT_QUANTIZATION_PARAMETERS.hasNoMaxAccuracyDrop ); + const [maxNumIterations, setMaxNumIterations] = useState(DEFAULT_QUANTIZATION_PARAMETERS.maxNumIterations); const [calibrationSize, setCalibrationSize] = useState(DEFAULT_QUANTIZATION_PARAMETERS.calibrationSize); const [usesFullCalibrationDataset, setUsesFullCalibrationDataset] = useState( DEFAULT_QUANTIZATION_PARAMETERS.usesFullCalibrationDataset @@ -65,6 +67,7 @@ export const QuantizationDialog = ({ modelId, onClose }: QuantizationDialogProps parameters: { model_id: modelId, max_drop: hasNoMaxAccuracyDrop ? null : accuracyDrop / 100, + max_num_iterations: hasNoMaxAccuracyDrop ? null : maxNumIterations, max_calibration_subset_size: usesFullCalibrationDataset ? totalCount : effectiveCalibrationSize, }, }, @@ -105,6 +108,13 @@ export const QuantizationDialog = ({ modelId, onClose }: QuantizationDialogProps onReset={() => setAccuracyDrop(DEFAULT_QUANTIZATION_PARAMETERS.accuracyDrop)} /> + setMaxNumIterations(DEFAULT_QUANTIZATION_PARAMETERS.maxNumIterations)} + /> + void; + isDisabled: boolean; + onReset: () => void; +}; + +export const MaxNumIterationsField = ({ value, onChange, isDisabled, onReset }: MaxNumIterationsFieldProps) => { + return ( + + Max number of iterations + + + Maximum number of iterations of accuracy-aware quantization. +
+
+ Accuracy-aware quantization iteratively removes model layers from the quantization scope until the + max accuracy drop criteria is met. Limiting the number of iterations can significantly reduce the + time it takes to quantize models with many layers. +
+
+ This parameter is only used when a max accuracy drop is set. + + + + + + + + ); +}; + type CalibrationDatasetSizeFieldProps = { value: number; onChange: (value: number) => void; diff --git a/application/ui/tests/models/model-details.spec.ts b/application/ui/tests/models/model-details.spec.ts index e6873157ed4..77f6ad776c4 100644 --- a/application/ui/tests/models/model-details.spec.ts +++ b/application/ui/tests/models/model-details.spec.ts @@ -279,6 +279,7 @@ test.describe('Model Details', () => { await modelsPage.getNoMaximumCheckbox().click(); await modelsPage.getAccuracyDropInput().fill('5'); + await modelsPage.getMaxNumIterationsInput().fill('7'); await modelsPage.getCalibrationSizeInput().fill('300'); await modelsPage.submitQuantization(); @@ -291,6 +292,7 @@ test.describe('Model Details', () => { parameters: { model_id: 'model-1', max_drop: 0.05, + max_num_iterations: 7, max_calibration_subset_size: 300, }, }); diff --git a/application/ui/tests/models/models-page.ts b/application/ui/tests/models/models-page.ts index bdc6511f932..8db89169587 100644 --- a/application/ui/tests/models/models-page.ts +++ b/application/ui/tests/models/models-page.ts @@ -212,6 +212,10 @@ export class ModelsPage { return this.getQuantizationDialog().getByRole('textbox', { name: 'Change Max calibration size' }); } + getMaxNumIterationsInput() { + return this.getQuantizationDialog().getByRole('textbox', { name: 'Change Max number of iterations' }); + } + getNoMaximumCheckbox() { return this.getQuantizationDialog().getByLabel('No maximum'); } diff --git a/library/src/getitune/backend/openvino/engine.py b/library/src/getitune/backend/openvino/engine.py index 205f856e066..1213e5ebf3a 100644 --- a/library/src/getitune/backend/openvino/engine.py +++ b/library/src/getitune/backend/openvino/engine.py @@ -412,6 +412,7 @@ def optimize( datamodule: DataModule | None = None, max_data_subset_size: int | None = None, max_drop: float | None = None, + max_num_iterations: int | None = None, ) -> Path: """Apply Post-Training Quantization (PTQ) to optimize the model. @@ -428,6 +429,8 @@ def optimize( Defaults to None. max_drop (float | None, optional): Maximum accuracy drop allowed for accuracy-aware quantization. Defaults to None. + max_num_iterations (int | None, optional): Maximum number of iterations for accuracy-aware + quantization. ``None`` means unlimited. Only used when ``max_drop`` is set. Defaults to None. Returns: Path: Path to the optimized model. @@ -457,6 +460,8 @@ def optimize( ptq_config["subset_size"] = max_data_subset_size if max_drop is not None: ptq_config["max_drop"] = max_drop + if max_num_iterations is not None: + ptq_config["max_num_iterations"] = max_num_iterations logger.debug(f"PTQ configuration: {ptq_config}") return model.optimize( diff --git a/library/src/getitune/backend/openvino/models/base.py b/library/src/getitune/backend/openvino/models/base.py index 4327af9246d..a656ddb5b1f 100644 --- a/library/src/getitune/backend/openvino/models/base.py +++ b/library/src/getitune/backend/openvino/models/base.py @@ -337,9 +337,17 @@ def check_if_quantized(model: openvino.Model) -> bool: quantization_dataset = nncf.Dataset(train_dataset, self.transform_fn) + # max_num_iterations is not a direct argument of the NNCF quantization APIs; it is configured through + # AdvancedAccuracyRestorerParameters. Pop it here so it is never forwarded as an invalid kwarg. + max_num_iterations = ptq_config.pop("max_num_iterations", None) + if ptq_config.get("max_drop") is not None: validation_dataset = nncf.Dataset(data_module.val_dataloader(), self.transform_fn) validation_fn = self._create_validation_fn(data_module) + if max_num_iterations is not None: + ptq_config["advanced_accuracy_restorer_parameters"] = nncf.AdvancedAccuracyRestorerParameters( + max_num_iterations=max_num_iterations + ) compressed_model = nncf.quantize_with_accuracy_control( model=ov_model, calibration_dataset=quantization_dataset, diff --git a/library/tests/unit/backend/openvino/test_engine.py b/library/tests/unit/backend/openvino/test_engine.py index 21045057667..b3f3e537abd 100644 --- a/library/tests/unit/backend/openvino/test_engine.py +++ b/library/tests/unit/backend/openvino/test_engine.py @@ -220,6 +220,16 @@ def test_optimizing_model(self, fxt_engine, mocker) -> None: fxt_engine.optimize(max_data_subset_size=100, checkpoint="path/to/exported_model.xml") assert mock_ov_model.return_value.optimize.call_args[0][2]["subset_size"] == 100 + # With max_drop and max_num_iterations (accuracy-aware quantization) + fxt_engine.optimize(max_drop=0.01, max_num_iterations=5, checkpoint="path/to/exported_model.xml") + ptq_config = mock_ov_model.return_value.optimize.call_args[0][2] + assert ptq_config["max_drop"] == 0.01 + assert ptq_config["max_num_iterations"] == 5 + + # Without max_num_iterations it should not be part of the PTQ config + fxt_engine.optimize(max_drop=0.01, checkpoint="path/to/exported_model.xml") + assert "max_num_iterations" not in mock_ov_model.return_value.optimize.call_args[0][2] + def test_optimize_rejects_onnx(self, fxt_engine, fxt_onnx_model_path) -> None: """Test that optimize() raises RuntimeError for ONNX models.""" with pytest.raises(RuntimeError, match="does not support ONNX models"):