Add max_num_iterations parameter for accuracy-aware quantization

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:

application/backend/app/api/schemas/jobs/quantization.py

@@ -26,13 +26,20 @@ 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": {
             "example": {
                 "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

application/backend/app/execution/quantization/getitune_quantizer.py

@@ -238,19 +238,26 @@ 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().
 
         Args:
             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
 
@@ -332,12 +339,14 @@ def execute(self, params: QuantizationJobParams) -> None:
             ov_engine=ov_engine,
             subset_size=params.max_calibration_subset_size,
             max_drop=params.max_drop,
+            max_num_iterations=params.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": params.max_num_iterations if params.max_drop is not None else None,
         }
 
         with self._db_session_factory() as db:

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]):

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()]

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=10,
         )
         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,

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)}
                             />
 
+                            <MaxNumIterationsField
+                                value={maxNumIterations}
+                                onChange={setMaxNumIterations}
+                                isDisabled={hasNoMaxAccuracyDrop}
+                                onReset={() => setMaxNumIterations(DEFAULT_QUANTIZATION_PARAMETERS.maxNumIterations)}
+                            />
+
                             <CalibrationDatasetSizeField
                                 value={effectiveCalibrationSize}
                                 onChange={setCalibrationSize}

application/ui/src/features/models/model-listing/model-variants/quantization-dialog/quantization-fields.component.tsx

@@ -3,12 +3,24 @@
 
 import { ReactNode, useState } from 'react';
 
-import { ActionButton, Checkbox, Content, ContextualHelp, Flex, Grid, NumberField, Slider, Text } from '@geti-ui/ui';
+import {
+    ActionButton,
+    Checkbox,
+    Content,
+    ContextualHelp,
+    Flex,
+    Grid,
+    NumberField,
+    Slider,
+    Text,
+    View,
+} from '@geti-ui/ui';
 import { Refresh } from '@geti-ui/ui/icons';
 
 export const DEFAULT_QUANTIZATION_PARAMETERS = {
     accuracyDrop: 1.0,
     calibrationSize: 200,
+    maxNumIterations: 10,
     hasNoMaxAccuracyDrop: true,
     usesFullCalibrationDataset: false,
 };
@@ -97,6 +109,48 @@ export const MaxAccuracyDropField = ({
     );
 };
 
+type MaxNumIterationsFieldProps = {
+    value: number;
+    onChange: (value: number) => void;
+    isDisabled: boolean;
+    onReset: () => void;
+};
+
+export const MaxNumIterationsField = ({ value, onChange, isDisabled, onReset }: MaxNumIterationsFieldProps) => {
+    return (
+        <QuantizationFieldLayout onReset={onReset}>
+            <Text>Max number of iterations</Text>
+            <ContextualHelp>
+                <Content>
+                    Maximum number of iterations of accuracy-aware quantization.
+                    <br />
+                    <br />
+                    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.
+                    <br />
+                    <br />
+                    This parameter is only used when a max accuracy drop is set.
+                </Content>
+            </ContextualHelp>
+            <Flex gap={'size-100'}>
+                <NumberField
+                    hideStepper
+                    step={1}
+                    value={value}
+                    minValue={1}
+                    onChange={onChange}
+                    isDisabled={isDisabled}
+                    aria-label={'Change Max number of iterations'}
+                    formatOptions={{ maximumFractionDigits: 0 }}
+                    flex={1}
+                />
+            </Flex>
+            <View />
+        </QuantizationFieldLayout>
+    );
+};
+
 type CalibrationDatasetSizeFieldProps = {
     value: number;
     onChange: (value: number) => void;

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,
                 },
             });

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');
     }

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(

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,

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"):