double precision mode

TICO-DCO-1.0-Signed-off-by: s.malakhov <s.malakhov@partner.samsung.com>

Author: stamalakhov

tico/quantization/algorithm/gptq/gptq.py

@@ -49,6 +49,7 @@ def configure(
         maxshrink=0.8,
         trits=False,
         sensitivity=None,
+        mse_tolerance=1e-5,
     ):
         self.maxq = torch.tensor(2**bits - 1)
         self.perchannel = perchannel
@@ -58,6 +59,7 @@ def configure(
         self.grid = grid
         self.maxshrink = maxshrink
         self.sensitivity = sensitivity
+        self.mse_tolerance = mse_tolerance
         if trits:
             self.maxq = torch.tensor(-1)
 
@@ -209,7 +211,10 @@ def _update_best_params(self, best, err, scale1, zero1):
         Returns:
             Updated best error values
         """
-        tmp = err < best
+        # Relative tolerance: only update if new error is significantly better.
+        # This prevents tiny float noise (from batch-size-dependent Hessian
+        # rounding) from flipping the winning scale/zero.
+        tmp = err < best * (1 - self.mse_tolerance)
         if torch.any(tmp):
             best[tmp] = err[tmp]
             self.scale[tmp] = scale1[tmp]
@@ -226,7 +231,7 @@ def _grid_search(self, x, xmin, xmax, compute_error_fn):
             compute_error_fn: Function that takes (x, scale1, zero1) and returns error tensor
         """
         dev = x.device
-        best = torch.full([x.shape[0]], float("inf"), device=dev)
+        best = torch.full([x.shape[0]], float("inf"), device=dev, dtype=x.dtype)
         for i in range(int(self.maxshrink * self.grid)):
             p = 1 - i / self.grid
             scale1, zero1 = self._compute_shrink_params(p, xmin, xmax)

tico/quantization/algorithm/gptq/llama_quantizer.py

@@ -107,6 +107,63 @@ def __init__(self, config: GPTQConfig):
         # Reference to original model for use_orig_model_inference and GPTQv2
         self.orig_model: Optional[torch.nn.Module] = None
 
+    @staticmethod
+    def _cast_to_double(obj):
+        """Recursively cast all tensors in a list/dict/tensor to float64."""
+        if isinstance(obj, torch.Tensor):
+            return obj.double()
+        if isinstance(obj, (list, tuple)):
+            return type(obj)(GPTQQuantizer._cast_to_double(o) for o in obj)
+        if isinstance(obj, dict):
+            return {k: GPTQQuantizer._cast_to_double(v) for k, v in obj.items()}
+        return obj
+
+    def _run_layer_forward_double_precision(
+        self,
+        layer: torch.nn.Module,
+        args: List[Any],
+        kwargs: Dict[str, Any],
+        double_precision: bool,
+    ):
+        """Run a layer forward, optionally in float64 for batch-size-independence.
+
+        When double_precision=True, temporarily casts the layer and its inputs to float64
+        so that GPU matmul tiling differences don't produce different results
+        for different batch sizes. Outputs are kept in float64 to avoid losing
+        precision when they become the next layer's input.
+        """
+        if not double_precision:
+            return layer(*args, **kwargs)
+
+        # Save original dtypes and cast layer to double
+        orig_dtypes: Dict[str, torch.dtype] = {}
+        for name, param in layer.named_parameters():
+            orig_dtypes[name] = param.dtype
+            param.data = param.data.double()
+        for name, buf in layer.named_buffers():
+            orig_dtypes[f"buf:{name}"] = buf.dtype
+            buf.data = buf.data.double()
+
+        # Cast inputs to double
+        args_d = self._cast_to_double(args)
+        kwargs_d = self._cast_to_double(kwargs)
+
+        outs = layer(*args_d, **kwargs_d)
+
+        # Keep outputs in float64 — casting to float32 would introduce
+        # batch-size-dependent rounding when these outputs become the
+        # next layer's input, defeating the purpose of double precision.
+        if isinstance(outs, tuple):
+            outs = outs[0]
+
+        # Restore original dtypes
+        for name, param in layer.named_parameters():
+            param.data = param.data.to(orig_dtypes[name])
+        for name, buf in layer.named_buffers():
+            buf.data = buf.data.to(orig_dtypes[f"buf:{name}"])
+
+        return outs
+
     def _resolve_weight_bits(
         self,
         gptq_conf: GPTQConfig,
@@ -335,7 +392,12 @@ def convert(self, model):
                     cache_kwargs_batch = move_to_device(cache_kwargs_batch, device)
 
                     orig_layer = orig_layers[l_idx].to(device)
-                    orig_layer(*cache_args_batch, **cache_kwargs_batch)
+                    if gptq_conf.double_precision:
+                        self._run_layer_forward_double_precision(
+                            orig_layer, cache_args_batch, cache_kwargs_batch, True
+                        )
+                    else:
+                        orig_layer(*cache_args_batch, **cache_kwargs_batch)
                     orig_layer.cpu()
 
                 fp_inputs_cache.clear_hook()
@@ -346,10 +408,10 @@ def convert(self, model):
 
                 gptq: Dict[str, GPTQ] = {}
                 for name in subset:
-                    gptq[name] = GPTQ(subset[name])
+                    full_module_name = module_name[subset[name]]
+                    gptq[name] = GPTQ(subset[name], double_precision=gptq_conf.double_precision, layer_name=full_module_name)
                     gptq[name].saturation_threshold = gptq_conf.saturation_threshold
                     gptq[name].saturation_min_batches = gptq_conf.saturation_min_batches
-                    full_module_name = module_name[subset[name]]
                     weight_bits = self._resolve_weight_bits(
                         gptq_conf,
                         full_module_name=full_module_name,
@@ -369,6 +431,7 @@ def convert(self, model):
                         sym=gptq_conf.symmetric,
                         mse=gptq_conf.mse,
                         sensitivity=cur_sensitivity,
+                        mse_tolerance=gptq_conf.mse_tolerance,
                     )
 
                     # GPTQv2: Assign native_inp from FPInputsCache
@@ -413,7 +476,12 @@ def _hook(_, inp, out):
                     )
                     cache_kwargs_batch = move_to_device(cache_kwargs_batch, device)
 
-                    layer(*cache_args_batch, **cache_kwargs_batch)
+                    if gptq_conf.double_precision:
+                        self._run_layer_forward_double_precision(
+                            layer, cache_args_batch, cache_kwargs_batch, True
+                        )
+                    else:
+                        layer(*cache_args_batch, **cache_kwargs_batch)
 
                 # Remove handles
                 for h in handles:
@@ -435,6 +503,7 @@ def _hook(_, inp, out):
                         adaptive_percdamp=gptq_conf.adaptive_percdamp,
                         cond_threshold_good=gptq_conf.cond_threshold_good,
                         use_iterate=gptq_conf.use_iterate,
+                        actorder_precision=gptq_conf.actorder_precision,
                     )
                     quantizers[full_module_name] = gptq[name].quantizer
                     gptq[name].free()
@@ -463,9 +532,14 @@ def _hook(_, inp, out):
                     fp_cache_args_batch = gather_single_batch_from_list(fp_inps, batch_idx)
                     fp_cache_args_batch = move_to_device(fp_cache_args_batch, device)
                     orig_layer = orig_layers[l_idx].to(device)
-                    fp_outs = orig_layer(*fp_cache_args_batch, **cache_kwargs_batch)
+                    if gptq_conf.double_precision:
+                        fp_outs = self._run_layer_forward_double_precision(
+                            orig_layer, fp_cache_args_batch, cache_kwargs_batch, True
+                        )
+                    else:
+                        fp_outs = orig_layer(*fp_cache_args_batch, **cache_kwargs_batch)
+                        fp_outs = fp_outs[0] if isinstance(fp_outs, tuple) else fp_outs
                     orig_layer.cpu()
-                    fp_outs = fp_outs[0] if isinstance(fp_outs, tuple) else fp_outs
                     # Update inputs for next iteration.
                     if len(fp_inps) > 0:
                         if hasattr(fp_outs, "to") and hasattr(
@@ -477,16 +551,28 @@ def _hook(_, inp, out):
                         else:
                             fp_inps[0][batch_idx] = fp_outs
 
-                if orig_layers is None or gptq_conf.gptq_v2 is True:
-                    outs = layer(*cache_args_batch, **cache_kwargs_batch)
+                if gptq_conf.double_precision:
+                    if orig_layers is None or gptq_conf.gptq_v2 is True:
+                        outs = self._run_layer_forward_double_precision(
+                            layer, cache_args_batch, cache_kwargs_batch, True
+                        )
+                    else:
+                        orig_layer = orig_layers[l_idx].to(device)
+                        outs = self._run_layer_forward_double_precision(
+                            orig_layer, cache_args_batch, cache_kwargs_batch, True
+                        )
+                        orig_layer.cpu()
                 else:
-                    orig_layer = orig_layers[l_idx].to(device)
-                    outs = orig_layer(*cache_args_batch, **cache_kwargs_batch)
-                    orig_layer.cpu()
-                # LLaMA's decoder layer return type differs across Transformers versions:
-                # some return a tuple (hidden_states, ...), others return just a tensor.
-                # This line ensures we always take the first element when it's a tuple.
-                outs = outs[0] if isinstance(outs, tuple) else outs
+                    if orig_layers is None or gptq_conf.gptq_v2 is True:
+                        outs = layer(*cache_args_batch, **cache_kwargs_batch)
+                    else:
+                        orig_layer = orig_layers[l_idx].to(device)
+                        outs = orig_layer(*cache_args_batch, **cache_kwargs_batch)
+                        orig_layer.cpu()
+                    # LLaMA's decoder layer return type differs across Transformers versions:
+                    # some return a tuple (hidden_states, ...), others return just a tensor.
+                    # This line ensures we always take the first element when it's a tuple.
+                    outs = outs[0] if isinstance(outs, tuple) else outs
                 # Update inputs for next iteration.
                 if len(self.cache_args) > 0:
                     if hasattr(outs, "to") and hasattr(
@@ -557,8 +643,8 @@ def _quantize_lm_head(self, model, quantizers):
                 self.cache_args[0][batch_idx] = move_to_cpu(hidden_states)
 
         layer = model.lm_head
-        gptq = GPTQ(layer)
         full_module_name = "lm_head"
+        gptq = GPTQ(layer, double_precision=gptq_conf.double_precision, layer_name=full_module_name)
         weight_bits = self._resolve_weight_bits(
             gptq_conf,
             full_module_name=full_module_name,
@@ -578,6 +664,7 @@ def _quantize_lm_head(self, model, quantizers):
             sym=gptq_conf.symmetric,
             mse=gptq_conf.mse,
             sensitivity=cur_sensitivity,
+                        mse_tolerance=gptq_conf.mse_tolerance,
         )
 
         # Hook to collect (inp, out) for GPTQ with optional weights
@@ -628,6 +715,7 @@ def _hook(_, inp, out):
             adaptive_percdamp=gptq_conf.adaptive_percdamp,
             cond_threshold_good=gptq_conf.cond_threshold_good,
             use_iterate=gptq_conf.use_iterate,
+                        actorder_precision=gptq_conf.actorder_precision,
         )
         quantizers[f"lm_head"] = gptq.quantizer
         gptq.free()

tico/quantization/algorithm/gptq/quant.py

@@ -60,12 +60,24 @@ class GPTQConfig(BaseConfig):
     mse: str | None = None
     sensitivity: dict[str, torch.Tensor] | None = None
 
+    # Relative tolerance for MSE grid search: a new scale/zero candidate only
+    # wins if its error is at least mse_tolerance better (relatively) than the
+    # current best. Prevents tiny float noise from flipping the winning
+    # scale/zero, making results more batch-size-independent.
+    mse_tolerance: float = 1e-5
+
     # GPTQ.fasterquant params (algorithm hyperparams)
     percdamp: float = 0.01
     groupsize: int = -1
     actorder: bool = True
     static_groups: bool = False
 
+    # Precision threshold for actorder sorting. Diagonal values of the Hessian
+    # are rounded to this precision grid before argsort, so differences smaller
+    # than this threshold are treated as noise (ties broken by original index
+    # via stable=True). This makes actorder more stable
+    actorder_precision: float = 1e-2
+
     # use this option to stabilize GPTQ for deep models
     use_orig_model_inference: bool = False
 
@@ -94,6 +106,12 @@ class GPTQConfig(BaseConfig):
     # is checked. Ensures r_eff has stabilized past the initial transient.
     saturation_min_batches: int = 4
 
+    # Use float64 (double) for Hessian accumulation to make GPTQ results
+    # batch-size-independent. Float32 accumulation causes different rounding
+    # depending on how samples are grouped into batches. Default: False
+    # (float32, backward compatible). Enable for reproducible results.
+    double_precision: bool = False
+
     @property
     def name(self) -> str:
         return "gptq"

tico/quantization/algorithm/gptq/quantizer.py

@@ -69,12 +69,24 @@ class LlamaGPTQConfig(BaseConfig):
     mse: str | None = None
     sensitivity: dict[str, torch.Tensor] | None = None
 
+    # Relative tolerance for MSE grid search: a new scale/zero candidate only
+    # wins if its error is at least mse_tolerance better (relatively) than the
+    # current best. Prevents tiny float noise from flipping the winning
+    # scale/zero, making results more batch-size-independent.
+    mse_tolerance: float = 1e-5
+
     # GPTQ.fasterquant params (algorithm hyperparams)
     percdamp: float = 0.01
     groupsize: int = -1
     actorder: bool = True
     static_groups: bool = False
 
+    # Precision threshold for actorder sorting. Diagonal values of the Hessian
+    # are rounded to this precision grid before argsort, so differences smaller
+    # than this threshold are treated as noise (ties broken by original index
+    # via stable=True). This makes actorder more stable
+    actorder_precision: float = 1e-2
+
     # use this option to stabilize GPTQ for deep models
     use_orig_model_inference: bool = False
 
@@ -111,6 +123,12 @@ class LlamaGPTQConfig(BaseConfig):
     # is checked. Ensures r_eff has stabilized past the initial transient.
     saturation_min_batches: int = 4
 
+    # Use float64 (double) for Hessian accumulation to make GPTQ results
+    # batch-size-independent. Float32 accumulation causes different rounding
+    # depending on how samples are grouped into batches. Default: False
+    # (float32, backward compatible). Enable for reproducible results.
+    double_precision: bool = False
+
     @property
     def name(self) -> str:
         return "llama_gptq"

tico/quantization/config/gptq.py

@@ -480,6 +480,17 @@ def build_parser() -> argparse.ArgumentParser:
         ),
     )
 
+    parser.add_argument(
+        "--gptq_double_precision",
+        action="store_true",
+        default=False,
+        help=(
+            "Use float64 (double) for Hessian accumulation to make GPTQ results "
+            "stable. Float32 accumulation causes different rounding "
+            "depending on how samples are grouped into batches. Default: False "
+            "(float32, backward compatible). Enable for exact/reproducible results."
+        ),
+    )
     parser.add_argument(
         "--gptq_saturation_threshold",
         type=float,
@@ -972,6 +983,7 @@ def build_gptq_config(
             use_subgroup_runner=args.llama_gptq_use_subgroup_runner,
             sample_weights=sample_weights,
             saturation_threshold=args.gptq_saturation_threshold,
+            double_precision=args.gptq_double_precision,
         )
         return config
     else:
@@ -991,6 +1003,7 @@ def build_gptq_config(
             use_iterate=args.gptq_use_iterate,
             sample_weights=sample_weights,
             saturation_threshold=args.gptq_saturation_threshold,
+            double_precision=args.gptq_double_precision,
         )
         return config
 
@@ -1001,6 +1014,10 @@ def save_model_to(
     """
     Export and save the whole quantized model in circle format.
     """
+    if not hasattr(q_m, "wrapped"):
+        print("Saving whole model circle is supported only for PTQ quantized model")
+        return
+
     q_m.eval()
     q_m.cpu()
     model_name = "model_prefill" if prefill_decode else "model"
@@ -3184,6 +3201,7 @@ def get_ptq_model_name(model, args):
 
     # --- Calibration options ------------------------------------------------
     parts.append(str(args.nsamples_for_qcalibration))
+    parts.append(f"bs{args.batch}")
     if (
         args.calibration_samples_to_use is not None
         and args.calibration_samples_to_use != args.nsamples_for_qcalibration
@@ -4513,8 +4531,12 @@ def save_requested_artifacts(q_m, tokenizer, calib_inputs, args, sample_weights=
         print(f"Saving calibration dataset to {save_path.resolve()}")
         torch.save({"calib_inputs": calib_inputs, "sample_weights": sample_weights}, save_path)
 
+    # When --no_PTQ is used, q_m has no PTQ wrapper so there is no .wrapped
+    # attribute.  Use q_m directly in that case.
+    inner = q_m if args.no_PTQ else q_m.wrapped
+
     if should_save(args, "ptq_checkpoint"):
-        save_path = output_dir / get_ptq_model_name(q_m.wrapped, args)
+        save_path = output_dir / get_ptq_model_name(inner, args)
         print(f"Saving PTQ checkpoint to {save_path.resolve()}")
         torch.save(q_m, save_path)
 
@@ -4528,7 +4550,7 @@ def save_requested_artifacts(q_m, tokenizer, calib_inputs, args, sample_weights=
         )
 
     if should_save(args, "circle_per_layer"):
-        max_seq_len = args.max_seq_len or q_m.wrapped.config.max_position_embeddings
+        max_seq_len = args.max_seq_len or inner.config.max_position_embeddings
         save_layers_to(
             q_m,
             max_seq_len,