From 0819192139bfbaee1525c9cf79020a16db2e7353 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Tue, 16 Dec 2025 23:55:46 -0500 Subject: [PATCH 1/3] update --- tests/e2e/e2e_utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/e2e/e2e_utils.py b/tests/e2e/e2e_utils.py index 765d864cc..c9cda5afc 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -8,6 +8,7 @@ from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier from tests.test_timer.timer_utils import log_time from tests.testing_utils import process_dataset +from transformers import DefaultDataCollator def load_model(model: str, model_class: str, device_map: str | None = None): @@ -94,7 +95,11 @@ def data_collator(batch): ) # Apply quantization. + logger.info("ONESHOT KWARGS", oneshot_kwargs) + + oneshot_kwargs["shuffle_calibration_samples"] = True + oneshot_kwargs["data_collator"] = DefaultDataCollator() _run_oneshot(**oneshot_kwargs) return oneshot_kwargs["model"], processor From 11bf7ffc3d77d92c11298fe5743e2c3284ea4132 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Wed, 17 Dec 2025 09:52:33 -0500 Subject: [PATCH 2/3] clean-up --- tests/e2e/e2e_utils.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/e2e/e2e_utils.py b/tests/e2e/e2e_utils.py index c9cda5afc..eb2546052 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -1,14 +1,15 @@ +from typing import Callable + import torch import transformers from datasets import load_dataset from loguru import logger -from transformers import AutoProcessor +from transformers import AutoProcessor, DefaultDataCollator from llmcompressor import oneshot from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier from tests.test_timer.timer_utils import log_time from tests.testing_utils import process_dataset -from transformers import DefaultDataCollator def load_model(model: str, model_class: str, device_map: str | None = None): @@ -35,9 +36,12 @@ def run_oneshot_for_e2e_testing( dataset_config: str, scheme: str, quant_type: str, + shuffle_calibration_samples: bool = True, + data_collator: str | Callable = DefaultDataCollator(), ): # Load model. oneshot_kwargs = {} + oneshot_kwargs["data_collator"] = data_collator loaded_model = load_model(model=model, model_class=model_class) processor = AutoProcessor.from_pretrained(model) @@ -75,6 +79,7 @@ def data_collator(batch): oneshot_kwargs["data_collator"] = data_collator oneshot_kwargs["model"] = loaded_model + oneshot_kwargs["shuffle_calibration_samples"] = shuffle_calibration_samples if recipe: oneshot_kwargs["recipe"] = recipe else: @@ -95,11 +100,8 @@ def data_collator(batch): ) # Apply quantization. - + breakpoint() logger.info("ONESHOT KWARGS", oneshot_kwargs) - - oneshot_kwargs["shuffle_calibration_samples"] = True - oneshot_kwargs["data_collator"] = DefaultDataCollator() _run_oneshot(**oneshot_kwargs) return oneshot_kwargs["model"], processor From cc25d2d6435f5c11628dd994b6fed63625b4eb3b Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Thu, 18 Dec 2025 18:56:07 -0500 Subject: [PATCH 3/3] Apply suggestion from @brian-dellabetta Co-authored-by: Brian Dellabetta Signed-off-by: Dipika Sikka --- tests/e2e/e2e_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/e2e/e2e_utils.py b/tests/e2e/e2e_utils.py index eb2546052..05bc3b623 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -100,7 +100,6 @@ def data_collator(batch): ) # Apply quantization. - breakpoint() logger.info("ONESHOT KWARGS", oneshot_kwargs) _run_oneshot(**oneshot_kwargs)