From 3560e8b999682cf18861b8973d3437f15aa05d85 Mon Sep 17 00:00:00 2001 From: Jan Lasek Date: Wed, 27 Nov 2024 17:03:31 +0100 Subject: [PATCH] Make export_config mandatory to get path & set default quantization_config Signed-off-by: Jan Lasek --- nemo/collections/llm/api.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nemo/collections/llm/api.py b/nemo/collections/llm/api.py index 4bafdd97ba21..adf98747059c 100644 --- a/nemo/collections/llm/api.py +++ b/nemo/collections/llm/api.py @@ -265,12 +265,11 @@ def validate( @run.cli.entrypoint(name="ptq", namespace="llm") def ptq( nemo_checkpoint: str, + export_config: ExportConfig, calib_tp: int = 1, calib_pp: int = 1, quantization_config: Annotated[Optional[QuantizationConfig], run.Config[QuantizationConfig]] = None, - export_config: Optional[Union[ExportConfig, run.Config[ExportConfig]]] = None, ) -> Path: - # TODO: Fix "nemo_run.cli.cli_parser.CLIException: An unexpected error occurred (Argument: , Context: {})" """ Applies Post-Training Quantization (PTQ) for a model using the specified quantization and export configs. It runs calibration for a small dataset to collect scaling factors low-precision GEMMs used by desired quantization method. @@ -297,6 +296,9 @@ def ptq( Returns: Path: The path where the quantized checkpoint has been saved after calibration. """ + if not quantization_config: + quantization_config = QuantizationConfig() + if export_config.path is None: raise ValueError("The export_config.path needs to be specified, got None.")