Merge pull request #13 from sotanengel/feature/quantization-progress-eta

feat: quantization progress logs with ETA

Author: FKKimura

onecomp/qep/_quantize_with_qep.py

@@ -36,6 +36,8 @@ def run_quantize_with_qep(
     quantizer: Quantizer,
     qep_config: QEPConfig,
     calibration_config: CalibrationConfig,
+    *,
+    quantization_progress: bool = True,
 ):
     """Run quantization with Quantization Error Propagation (QEP).
 
@@ -51,6 +53,7 @@ def run_quantize_with_qep(
         qep_config (QEPConfig): Configuration for QEP
             (percdamp, perccorr, exclude_layer_keywords).
         calibration_config (CalibrationConfig): Calibration parameters.
+        quantization_progress (bool): When True, log ``[progress]`` with ETA per layer.
 
     """
     model = model_config.load_model()
@@ -80,6 +83,17 @@ def run_quantize_with_qep(
 
     logger.info("Quantizing the model using %s", quantizer.name)
 
+    progress = None
+    if quantization_progress:
+        # pylint: disable-next=import-outside-toplevel
+        from onecomp.utils.quantization_progress import QuantizationProgressTracker
+
+        progress = QuantizationProgressTracker(
+            logger,
+            len(quantizer.module_to_name),
+            "QEP quantization (general, per layer)",
+        )
+
     # 2. For each target layer, perform the following sequentially
     for module, name in quantizer.module_to_name.items():
 
@@ -114,6 +128,8 @@ def run_quantize_with_qep(
 
         # 2-4. Free memory
         del quant_input_activation
+        if progress is not None:
+            progress.step_complete(name)
 
     del original_input_activations
     quantizer.execute_post_processing()

onecomp/qep/_quantize_with_qep_arch.py

@@ -143,6 +143,7 @@ def compute_hessian_and_crossterm(
     def make_hook(name):
         def hook(module, inp, out):
             dest[name] = inp[0] if isinstance(inp, tuple) else inp
+
         return hook
 
     handlers = [
@@ -213,6 +214,7 @@ def _compute_per_module_hessians(
     def _make_hook(key):
         def hook(_, inp, __):
             dest[key] = inp[0] if isinstance(inp, tuple) else inp
+
         return hook
 
     handlers = [m.register_forward_hook(_make_hook(i)) for i, m in enumerate(modules)]
@@ -251,10 +253,7 @@ def hook(_, inp, __):
     for h in handlers:
         h.remove()
 
-    return {
-        modules[i]: (hessians[i] if nsamples[i] > 0 else None)
-        for i in range(len(modules))
-    }
+    return {modules[i]: (hessians[i] if nsamples[i] > 0 else None) for i in range(len(modules))}
 
 
 @torch.no_grad()
@@ -263,6 +262,8 @@ def run_quantize_with_qep_arch(
     quantizer: Quantizer,
     qep_config: QEPConfig,
     calibration_config: CalibrationConfig,
+    *,
+    quantization_progress: bool = True,
 ):
     """Run architecture-aware quantization with QEP.
 
@@ -279,6 +280,7 @@ def run_quantize_with_qep_arch(
         qep_config (QEPConfig): Configuration for QEP
             (percdamp, perccorr, exclude_layer_keywords).
         calibration_config (CalibrationConfig): Calibration parameters.
+        quantization_progress (bool): When True, log ``[progress]`` with ETA per target layer.
 
     """
 
@@ -318,6 +320,17 @@ def run_quantize_with_qep_arch(
         name for module, name in quantizer.module_to_name.items() if module in block_modules
     }
 
+    progress = None
+    if quantization_progress:
+        # pylint: disable-next=import-outside-toplevel
+        from onecomp.utils.quantization_progress import QuantizationProgressTracker
+
+        progress = QuantizationProgressTracker(
+            logger,
+            len(remaining_targets),
+            "QEP quantization (architecture-aware)",
+        )
+
     # 2. For each target transformer block, perform the following sequentially
     for block_idx, block in enumerate(blocks):
 
@@ -365,9 +378,7 @@ def run_quantize_with_qep_arch(
             targets = [m for m in group_q if m in quantizer.module_to_name]
             if not targets:
                 continue
-            is_expert = any(
-                ".experts." in quantizer.module_to_name[m] for m in targets
-            )
+            is_expert = any(".experts." in quantizer.module_to_name[m] for m in targets)
             if is_expert:
                 expert_modules_q.extend(targets)
             else:
@@ -442,6 +453,8 @@ def run_quantize_with_qep_arch(
                         name,
                     )
                 remaining_targets.discard(name)
+                if progress is not None:
+                    progress.step_complete(name)
 
         # 4. Process MoE expert layers with per-module Hessians (no cross-term)
         if expert_modules_q:
@@ -451,7 +464,12 @@ def run_quantize_with_qep_arch(
                 len(expert_modules_q),
             )
             expert_hessians = _compute_per_module_hessians(
-                block_q, expert_modules_q, inps_q, kwargs, batch_size, device,
+                block_q,
+                expert_modules_q,
+                inps_q,
+                kwargs,
+                batch_size,
+                device,
             )
             for module_q in expert_modules_q:
                 name = quantizer.module_to_name[module_q]
@@ -462,6 +480,8 @@ def run_quantize_with_qep_arch(
                         name,
                     )
                     remaining_targets.discard(name)
+                    if progress is not None:
+                        progress.step_complete(f"{name} (skipped, no tokens)")
                     continue
 
                 logger.info(
@@ -489,6 +509,8 @@ def run_quantize_with_qep_arch(
                         name,
                     )
                 remaining_targets.discard(name)
+                if progress is not None:
+                    progress.step_complete(name)
 
         # forward input to the next block
         inps_q = forward_input(inps_q, block_q, kwargs, batch_size, device)

onecomp/runner.py

@@ -86,6 +86,7 @@ def __init__(
         multi_gpu=False,
         gpu_ids=None,
         post_processes=None,
+        quantization_progress=True,
     ):
         """__init__ method
 
@@ -130,6 +131,11 @@ def __init__(
                 a quantized model on CPU (built via
                 ``create_quantized_model``) and may modify it in-place.
                 Processes are executed in order.  Default is None.
+            quantization_progress (bool):
+                When ``True`` (default), emit ``[progress]`` log lines with
+                completed steps, elapsed time, and a linear ETA estimate
+                during long quantization (calibration, chunked, multi-GPU,
+                QEP).  Set to ``False`` for quiet runs (e.g. CI).
 
         Note:
             For zero-config quantization (VRAM auto-estimation +
@@ -215,6 +221,7 @@ def __init__(
         self.lpcd_config = None
         if lpcd:
             self.lpcd_config = lpcd_config if lpcd_config is not None else LPCDConfig()
+        self.quantization_progress = quantization_progress
 
     def check(self):
         """Check the settings
@@ -323,27 +330,19 @@ def _exclude_moe_router_if_needed(self):
         config = self.model_config.load_config()
         num_experts = (
             getattr(config, "num_experts", 0)
-            or getattr(
-                getattr(config, "text_config", None), "num_experts", 0
-            ) or 
-            0
+            or getattr(getattr(config, "text_config", None), "num_experts", 0)
+            or 0
         )
         if num_experts == 0:
             return
 
         keyword = "router"
-        target_quantizers = (
-            self.quantizers
-            if self.quantizers is not None
-            else [self.quantizer]
-        )
+        target_quantizers = self.quantizers if self.quantizers is not None else [self.quantizer]
         for q in target_quantizers:
             if q.exclude_layer_keywords is None:
                 q.exclude_layer_keywords = [keyword]
             elif keyword not in q.exclude_layer_keywords:
-                q.exclude_layer_keywords = list(q.exclude_layer_keywords) + [
-                    keyword
-                ]
+                q.exclude_layer_keywords = list(q.exclude_layer_keywords) + [keyword]
 
         self.logger.info(
             "MoE model (num_experts=%d): excluding '%s' layers from "
@@ -520,12 +519,9 @@ def auto_run(
             uniform_bit = max(valid_wbits)
             if save_dir == "auto":
                 model_name = model_id.rstrip("/").split("/")[-1]
-                save_dir = (
-                    f"{model_name}-gptq-{uniform_bit}bit"
-                )
+                save_dir = f"{model_name}-gptq-{uniform_bit}bit"
             logger.warning(
-                "Gemma 4 detected → falling back to uniform GPTQ %d-bit "
-                "(target wbits=%.2f)",
+                "Gemma 4 detected → falling back to uniform GPTQ %d-bit " "(target wbits=%.2f)",
                 uniform_bit,
                 wbits,
             )
@@ -536,6 +532,7 @@ def auto_run(
                 save_dir = f"{model_name}-autobit-{wbits}bit"
 
             from .quantizer.autobit import AutoBitQuantizer
+
             candidate_quantizers = [
                 GPTQ(wbits=b, groupsize=groupsize, **kwargs) for b in candidate_bits
             ]
@@ -604,8 +601,30 @@ def quantize_with_calibration(self):
 
         # Register hooks to all linear layers
         handles = []
+        progress = None
+        if self.quantization_progress:
+            # pylint: disable-next=import-outside-toplevel
+            from .utils.quantization_progress import QuantizationProgressTracker
+
+            progress = QuantizationProgressTracker(
+                logger,
+                len(self.quantizer.module_to_name),
+                "Calibration quantization layers",
+            )
+
+        if progress:
+            quantize_bound = self.quantizer.quantize
+
+            def _quantize_hook(module, input, output):  # pylint: disable=redefined-builtin
+                quantize_bound(module, input, output)
+                progress.step_complete(self.quantizer.module_to_name[module])
+
+            hook_fn = _quantize_hook
+        else:
+            hook_fn = self.quantizer.quantize
+
         for module in self.quantizer.module_to_name.keys():
-            handle = module.register_forward_hook(self.quantizer.quantize)
+            handle = module.register_forward_hook(hook_fn)
             handles.append(handle)
 
         logger.info("Quantizing the model using %s", self.quantizer.name)
@@ -645,6 +664,7 @@ def quantize_with_calibration_chunked(self):
             model_config=self.model_config,
             quantizers=self.quantizers if self.quantizers is not None else [self.quantizer],
             calibration_config=self.calibration_config,
+            quantization_progress=self.quantization_progress,
         )
 
     def quantize_with_calibration_on_multi_gpu(self):
@@ -673,6 +693,7 @@ def quantize_with_calibration_on_multi_gpu(self):
             quantizer=self.quantizer,
             calibration_config=self.calibration_config,
             gpu_ids=self.gpu_ids,
+            quantization_progress=self.quantization_progress,
         )
 
         # Store results in quantizer.results
@@ -699,8 +720,20 @@ def quantize_without_calibration(self):
             "Quantizing the model without calibration using %s",
             self.quantizer.name,
         )
+        progress = None
+        if self.quantization_progress:
+            # pylint: disable-next=import-outside-toplevel
+            from .utils.quantization_progress import QuantizationProgressTracker
+
+            progress = QuantizationProgressTracker(
+                logger,
+                len(self.quantizer.module_to_name),
+                "Quantization without calibration (layers)",
+            )
         for module in self.quantizer.module_to_name.keys():
             self.quantizer.quantize(module, None, None)
+            if progress:
+                progress.step_complete(self.quantizer.module_to_name[module])
 
         self.quantizer.execute_post_processing()
 
@@ -721,6 +754,7 @@ def quantize_with_qep(self):
             quantizer=self.quantizer,
             qep_config=self.qep_config,
             calibration_config=self.calibration_config,
+            quantization_progress=self.quantization_progress,
         )
 
         if self.qep_config.general:
@@ -868,7 +902,7 @@ def prepare_calibration_dataset(self, device, model=None):
 
         Args:
             device (torch.device): Device to place tensors on (CPU or GPU)
-            model: Model instance (optional). Add model-specific fields 
+            model: Model instance (optional). Add model-specific fields
             (e.g. mm_token_type_ids for Gemma 4).
 
         Returns:
@@ -992,10 +1026,7 @@ def save_quantization_statistics(self, path: str, quantizer=None):
 
         logger.info("Saving the quantization statistics to %s", path)
 
-        statistics = {
-            key: result.get_statistics()
-            for key, result in quantizer.results.items()
-        }
+        statistics = {key: result.get_statistics() for key, result in quantizer.results.items()}
 
         with open(path, "w", encoding="utf-8") as f:
             json.dump(statistics, f, indent=4)
@@ -1674,15 +1705,10 @@ def create_quantized_model(self, pack_weights: bool = True, quantizer=None, use_
         # cf) https://docs.vllm.ai/en/stable/features/quantization/#implementing-a-quantized-moe-method
         num_experts = (
             getattr(model.config, "num_experts", None)
-            or getattr(
-                getattr(model.config, "text_config", None), "num_experts", None
-            )
+            or getattr(getattr(model.config, "text_config", None), "num_experts", None)
             or 0
         )
-        if (
-            quant_config.get("quant_method") == "gptq"
-            and num_experts > 0
-        ):
+        if quant_config.get("quant_method") == "gptq" and num_experts > 0:
             quant_config["quant_method"] = "mixed_gptq"
             self.logger.info(
                 "MoE model detected (num_experts=%d): "
@@ -1706,15 +1732,13 @@ def _patch_k_eq_v_for_vllm(self, model, quant_config: dict) -> None:
         Gemma4 full-attention layers with attention_k_eq_v=True have no
         v_proj weight — the model reuses key states as value states.
         vLLM fuses q/k/v into a single qkv_proj and requires all shards
-        to share the same quantization status.  
+        to share the same quantization status.
         """
         text_cfg = getattr(model.config, "text_config", None)
         if text_cfg is None or not getattr(text_cfg, "attention_k_eq_v", False):
             return
         layer_types = getattr(text_cfg, "layer_types", [])
-        k_eq_v_indices = {
-            i for i, lt in enumerate(layer_types) if lt == "full_attention"
-        }
+        k_eq_v_indices = {i for i, lt in enumerate(layer_types) if lt == "full_attention"}
         if not k_eq_v_indices:
             return
 
@@ -1752,9 +1776,7 @@ def _patch_k_eq_v_for_vllm(self, model, quant_config: dict) -> None:
                 and "self_attn.k_proj" in layer_cfg
                 and "self_attn.v_proj" not in layer_cfg
             ):
-                layer_cfg["self_attn.v_proj"] = copy.deepcopy(
-                    layer_cfg["self_attn.k_proj"]
-                )
+                layer_cfg["self_attn.v_proj"] = copy.deepcopy(layer_cfg["self_attn.k_proj"])
 
         for key in ("modules_in_block_to_quantize", "quantized_layer_names"):
             names = quant_config.get(key, [])
@@ -1825,6 +1847,7 @@ def save_quantized_model(self, save_directory: str, pack_weights: bool = True):
         if src_dir and not os.path.isdir(src_dir):
             # when the model_id is specified, the path is modifed to the local directory
             from huggingface_hub import snapshot_download
+
             src_dir = snapshot_download(src_dir, local_files_only=True)
         if src_dir and os.path.isdir(src_dir):
             for fname in ("processor_config.json", "preprocessor_config.json"):