diff --git a/compress_all.py b/compress_all.py new file mode 100644 index 0000000000..0f94a88f49 --- /dev/null +++ b/compress_all.py @@ -0,0 +1,235 @@ +import argparse +import sys +from pathlib import Path + +import torch +from compressed_tensors.offload import init_dist +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationScheme, + QuantizationStrategy, + QuantizationType, +) +from compressed_tensors.quantization.quant_args import FP8_E4M3_DATA +from transformers import AutoModelForCausalLM, AutoTokenizer + +from llmcompressor import oneshot +from llmcompressor.modifiers.quantization import QuantizationModifier +from llmcompressor.utils import load_context + +# ── Constants ──────────────────────────────────────────────────────── + +DEFAULT_MODELS = [ + "meta-llama/Meta-Llama-3-8B-Instruct", + "meta-llama/Meta-Llama-3-70B-Instruct", +] + +OUTPUT_BASE = Path("/tmp/fouroversix-sanity") + +COMMON = dict( + num_bits=4, + type=QuantizationType.FLOAT, + strategy=QuantizationStrategy.TENSOR_GROUP, + symmetric=True, + dynamic=False, + group_size=16, + scale_dtype=FP8_E4M3_DATA.dtype, + zp_dtype=FP8_E4M3_DATA.dtype, +) + +# ── Observer configs ───────────────────────────────────────────────── + +CONFIGS = { + "fouroversix": QuantizationArgs( + **COMMON, + observer="fouroversix", + ), + "mse-1x-1.5x": QuantizationArgs( + **COMMON, + observer="memoryless_mse", + observer_kwargs={ + "expand": 1.5, + "grid": 3, + "maxshrink": 0.67, + "norm": 2.0, + "patience": 100000, + }, + ), + "nvfp4_expanded_mse": QuantizationArgs( + **COMMON, + observer="nvfp4_expanded_mse", + ), + "expand-3.4": QuantizationArgs( + **COMMON, + observer="memoryless_mse", + observer_kwargs={ + "expand": 3.4, + "maxshrink": round(1 - 0.8 / 3.4, 4), + "grid": 200.0, + "patience": 200, + }, + ), + "expanded-norm1.8": QuantizationArgs( + **COMMON, + observer="nvfp4_expanded_mse", + observer_kwargs={"norm": 1.8}, + ), + "expanded-norm2.0": QuantizationArgs( + **COMMON, + observer="nvfp4_expanded_mse", + observer_kwargs={"norm": 2.0}, + ), + "expanded-norm2.2": QuantizationArgs( + **COMMON, + observer="nvfp4_expanded_mse", + observer_kwargs={"norm": 2.2}, + ), + "expanded-norm2.4": QuantizationArgs( + **COMMON, + observer="nvfp4_expanded_mse", + observer_kwargs={"norm": 2.4}, + ), + "default-mse": QuantizationArgs( + **COMMON, + observer="memoryless_mse", + ), + "minmax": QuantizationArgs( + **COMMON, + observer="memoryless_minmax", + ), + "expanded-gs-prior": QuantizationArgs( + **COMMON, + observer="nvfp4_expanded_mse", + observer_kwargs={"use_global_scale_prior": True}, + ), +} + +CONFIG_DESCRIPTIONS = { + "fouroversix": "FourOverSix: per-block M=6/M=4 adaptive scaling, gs_max=256", + "mse-1x-1.5x": "MSE 1x+1.5x: 2-point search matching FourOverSix search space", + "nvfp4_expanded_mse": "NVFP4 Expanded MSE: 1.8x→0.8x range, 112 steps (default norm=2.4)", + "expand-3.4": "Original ablation: expand=3.4, maxshrink=0.7647, grid=200, patience=200", + "expanded-norm1.8": "NVFP4 Expanded MSE: norm=1.8", + "expanded-norm2.0": "NVFP4 Expanded MSE: norm=2.0", + "expanded-norm2.2": "NVFP4 Expanded MSE: norm=2.2", + "expanded-norm2.4": "NVFP4 Expanded MSE: norm=2.4 (explicit)", + "default-mse": "Default MSE: expand=1.0, maxshrink=0.20, grid=100, norm=2.4", + "minmax": "MinMax: no MSE search, simple min/max scaling", + "expanded-gs-prior": "NVFP4 Expanded MSE: with global_scale prior in search", +} + + +# ── Helpers ────────────────────────────────────────────────────────── + + +def model_short_name(model_id: str) -> str: + return model_id.rstrip("/").split("/")[-1] + + +def compress( + model_id: str, + config_key: str, + output_base: Path, + force: bool = False, +): + name = model_short_name(model_id) + out = output_base / f"{name}-{config_key}" + + if out.exists() and not force: + print(f" SKIP (exists): {out}") + return out + + print(f"\n{'='*70}") + print(f" {out}") + print(f" model: {model_id}") + print(f" config: {config_key} — {CONFIG_DESCRIPTIONS.get(config_key, '')}") + print(f"{'='*70}\n") + + tokenizer = AutoTokenizer.from_pretrained(model_id) + with load_context(): + model = AutoModelForCausalLM.from_pretrained( + model_id, device_map="auto_offload" + ) + + recipe = QuantizationModifier( + config_groups={ + "group_0": QuantizationScheme( + targets=["Linear"], + weights=CONFIGS[config_key], + ) + }, + ignore=["lm_head"], + ) + + oneshot( + model=model, + recipe=recipe, + output_dir=str(out), + ) + tokenizer.save_pretrained(out) + del model + torch.cuda.empty_cache() + print(f" DONE: {out}\n") + return out + + +# ── Main ───────────────────────────────────────────────────────────── + + +def main(): + parser = argparse.ArgumentParser( + description="Compress models with NVFP4 observer configs" + ) + parser.add_argument( + "--models", nargs="+", default=DEFAULT_MODELS, + help="HuggingFace model IDs to compress", + ) + parser.add_argument( + "--configs", nargs="+", default=list(CONFIGS.keys()), + help="Observer configs to run", + ) + parser.add_argument( + "--output-dir", type=str, default=str(OUTPUT_BASE), + help="Base directory for compressed models", + ) + parser.add_argument( + "--force", action="store_true", + help="Recompress even if output exists", + ) + parser.add_argument( + "--list", action="store_true", + help="List configs and exit", + ) + args = parser.parse_args() + + if args.list: + print("Available configs:") + for key, desc in CONFIG_DESCRIPTIONS.items(): + print(f" {key:<25s} {desc}") + print(f"\nDefault models: {', '.join(DEFAULT_MODELS)}") + sys.exit(0) + + for key in args.configs: + if key not in CONFIGS: + print(f"Unknown config: {key}") + print(f"Available: {', '.join(CONFIGS.keys())}") + sys.exit(1) + + output_base = Path(args.output_dir) + output_base.mkdir(parents=True, exist_ok=True) + + init_dist() + + total = len(args.models) * len(args.configs) + idx = 0 + for model_id in args.models: + for config_key in args.configs: + idx += 1 + print(f"\n[{idx}/{total}] {model_short_name(model_id)} / {config_key}") + compress(model_id, config_key, output_base, args.force) + + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/eval_all.py b/eval_all.py new file mode 100644 index 0000000000..bbc3f9d762 --- /dev/null +++ b/eval_all.py @@ -0,0 +1,158 @@ +import argparse +import sys +from pathlib import Path + +import torch +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoTokenizer + +# ── Constants ──────────────────────────────────────────────────────── + +DEFAULT_MODELS = [ + "meta-llama/Meta-Llama-3-8B-Instruct", + "meta-llama/Meta-Llama-3-70B-Instruct", +] + +DEFAULT_CONFIGS = [ + "fouroversix", + "mse-1x-1.5x", + "nvfp4_expanded_mse", + "expand-3.4", +] + +COMPRESSED_DIR = Path("/tmp/fouroversix-sanity") + +CONFIG_LABELS = { + "fouroversix": "FourOverSix", + "mse-1x-1.5x": "MSE 1x+1.5x", + "nvfp4_expanded_mse": "nvfp4_expanded_mse", + "expand-3.4": "expand=3.4 (original ablation)", + "default-mse": "Default MSE", + "minmax": "MinMax", + "expanded-gs-prior": "Expanded MSE + gs prior", +} + + +# ── Evaluation ─────────────────────────────────────────────────────── + + +def evaluate_ppl(model, tokenizer, max_length=2048, stride=512): + testdata = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test") + text = "\n\n".join(testdata["text"]) + encodings = tokenizer(text, return_tensors="pt") + seq_len = encodings.input_ids.size(1) + + nlls = [] + prev_end_loc = 0 + for begin_loc in range(0, seq_len, stride): + end_loc = min(begin_loc + max_length, seq_len) + trg_len = end_loc - prev_end_loc + device = next(model.parameters()).device + input_ids = encodings.input_ids[:, begin_loc:end_loc].to(device) + target_ids = input_ids.clone() + target_ids[:, :-trg_len] = -100 + with torch.no_grad(): + outputs = model(input_ids, labels=target_ids) + nlls.append(outputs.loss) + prev_end_loc = end_loc + if end_loc == seq_len: + break + + return torch.exp(torch.stack(nlls).mean()).item() + + +def load_and_eval(path: str, label: str) -> float: + print(f" Loading {label} from {path}...", flush=True) + tokenizer = AutoTokenizer.from_pretrained(path) + model = AutoModelForCausalLM.from_pretrained( + path, device_map="auto", torch_dtype="auto" + ) + model.eval() + ppl = evaluate_ppl(model, tokenizer) + print(f" {label}: {ppl:.3f}") + del model + torch.cuda.empty_cache() + return ppl + + +# ── Main ───────────────────────────────────────────────────────────── + + +def model_short_name(model_id: str) -> str: + return model_id.rstrip("/").split("/")[-1] + + +def main(): + parser = argparse.ArgumentParser( + description="PPL evaluation for NVFP4 observer comparison" + ) + parser.add_argument( + "--models", nargs="+", default=DEFAULT_MODELS, + help="HuggingFace model IDs (baselines)", + ) + parser.add_argument( + "--configs", nargs="+", default=DEFAULT_CONFIGS, + help="Observer config names to evaluate", + ) + parser.add_argument( + "--compressed-dir", type=str, default=str(COMPRESSED_DIR), + help="Directory containing compressed models", + ) + parser.add_argument( + "--no-baseline", action="store_true", + help="Skip unquantized baseline evaluation", + ) + args = parser.parse_args() + + compressed_dir = Path(args.compressed_dir) + all_results = {} + + for model_id in args.models: + name = model_short_name(model_id) + print(f"\n{'='*60}") + print(f" {name}") + print(f"{'='*60}") + + results = [] + baseline_ppl = None + + if not args.no_baseline: + baseline_ppl = load_and_eval(model_id, "Unquantized") + results.append(("Unquantized", baseline_ppl)) + + for config_key in args.configs: + path = compressed_dir / f"{name}-{config_key}" + label = CONFIG_LABELS.get(config_key, config_key) + if not path.exists(): + print(f" SKIP (not found): {path}") + results.append((label, None)) + continue + ppl = load_and_eval(str(path), label) + results.append((label, ppl)) + + all_results[name] = (baseline_ppl, results) + + # ── Print markdown tables ──────────────────────────────────── + print(f"\n\n{'='*60}") + print("Results (markdown)") + print(f"{'='*60}\n") + + for name, (baseline_ppl, results) in all_results.items(): + print(f"### {name}\n") + print("| Config | word_perplexity | delta vs unquantized |") + print("|---|---|---|") + for label, ppl in results: + if ppl is None: + print(f"| {label} | — | — |") + elif label == "Unquantized": + print(f"| {label} | **{ppl:.3f}** | — |") + elif baseline_ppl is not None: + delta = ppl - baseline_ppl + print(f"| {label} | {ppl:.3f} | +{delta:.3f} |") + else: + print(f"| {label} | {ppl:.3f} | — |") + print() + + +if __name__ == "__main__": + main() diff --git a/src/llmcompressor/observers/__init__.py b/src/llmcompressor/observers/__init__.py index b5cb48a960..9c921f43a0 100644 --- a/src/llmcompressor/observers/__init__.py +++ b/src/llmcompressor/observers/__init__.py @@ -15,3 +15,4 @@ from .min_max import * from .mse import * from .imatrix import * +from .fouroversix import * diff --git a/src/llmcompressor/observers/fouroversix.py b/src/llmcompressor/observers/fouroversix.py new file mode 100644 index 0000000000..e06306a260 --- /dev/null +++ b/src/llmcompressor/observers/fouroversix.py @@ -0,0 +1,178 @@ +import torch +from compressed_tensors.quantization import QuantizationStrategy +from compressed_tensors.quantization.lifecycle import fake_quantize +from compressed_tensors.quantization.quant_args import ( + FloatArgs, + round_to_quantized_type_dtype, +) +from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam + +from llmcompressor.observers.base import Observer, QParamsDict + +__all__ = ["FourOverSixObserver"] + + +class _FP8ScaleData256(FloatArgs): + """FP8 E4M3 with max capped at 256 for FourOverSix global scale. + + Standard NVFP4 uses 448 (full FP8 E4M3 range) as the maximum scale factor. + FourOverSix uses 256 so that blocks containing a tensor's largest values can + still select M=4, because 256 * (6/4) = 384 fits in FP8 E4M3. + """ + + exponent = 4 + mantissa = 3 + bits = 8 + max = 256.0 + min = -256.0 + + +@Observer.register("fouroversix") +class FourOverSixObserver(Observer): + """ + Adaptive block scaling observer for NVFP4 (Four Over Six / 4/6). + + For each block, quantizes with both M=6 (standard NVFP4 full range) and + M=4 (reduced range that represents near-maximal values more accurately), + then selects the per-block scale that minimizes quantization error. + + FP4 E2M1 has non-uniform step sizes (0.5 below 2, 1.0 between 2-4, 2.0 + between 4-6), creating large error for values near 5. Scaling some blocks + to M=4 makes the representable value distribution more uniform, reducing + worst-case error for near-maximal values. + + The global scale uses M_FP8=256 instead of 448 so that *every* block, + including those containing the tensor's largest values, can benefit from + M=4 (since 256 * 6/4 = 384, which is representable in FP8 E4M3). + + Configurable ``observer_kwargs``: + scale_selection (str): Error metric for per-block selection. + "mse" - mean squared error (default, best for PTQ) + "mae" - mean absolute error (best for pre-training) + "abs_max" - maximum absolute error + + Reference + --------- + Cook et al., "Four Over Six: More Accurate NVFP4 Quantization with + Adaptive Block Scaling", arXiv:2512.02010, 2025. + """ + + SCALE_EXPANSION = 1.5 # 6 / 4 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + kw = self.args.observer_kwargs or {} + self.scale_selection: str = kw.get("scale_selection", "mse") + if self.scale_selection not in ("mse", "mae", "abs_max"): + raise ValueError( + f"scale_selection must be 'mse', 'mae', or 'abs_max', " + f"got '{self.scale_selection}'" + ) + self._observed_blocks: torch.Tensor | None = None + self._token_args = self.args.model_copy( + update={"strategy": QuantizationStrategy.TOKEN} + ) + + def update_statistics_from_observed(self, observed: torch.Tensor) -> None: + self.min_vals = torch.amin(observed, dim=(0, -1)) + self.max_vals = torch.amax(observed, dim=(0, -1)) + self._observed_blocks = observed.detach().clone() + + @torch.no_grad + def get_qparams(self) -> QParamsDict: + assert ( + self.has_statistics + ), "No statistics available. Call observer(value) first." + + global_scale = None + + if self.args.strategy == QuantizationStrategy.TENSOR_GROUP: + all_stats = self.fusion_handler.get_fused_statistics() + global_absmax = all_stats[0]["max_vals"].max() + for stats in all_stats: + global_absmax = torch.max( + global_absmax, -stats["min_vals"].min() + ) + global_absmax = torch.max( + global_absmax, stats["max_vals"].max() + ) + + global_scale = generate_gparam( + -global_absmax.reshape(1), + global_absmax.reshape(1), + scale_data=_FP8ScaleData256, + ) + + scale_6, zero_point = calculate_qparams( + min_vals=self.min_vals, + max_vals=self.max_vals, + quantization_args=self.args, + global_scale=global_scale, + ) + + if self._observed_blocks is not None: + scale = self._select_block_scales( + self._observed_blocks, scale_6, zero_point, global_scale + ) + self._observed_blocks = None + else: + scale = scale_6 + + self.delete_statistics() + + return {"scale": scale, "zero_point": zero_point, "global_scale": global_scale} + + def _select_block_scales( + self, + observed: torch.Tensor, + scale_6: torch.Tensor, + zero_point: torch.Tensor, + global_scale: torch.Tensor | None, + ) -> torch.Tensor: + """Choose M=6 or M=4 per block to minimize quantization error. + + 1. Expand M=6 scales by 1.5x -> M=4 scales, round to FP8. + 2. Fake-quantize the observed blocks with each scale set. + 3. Pick the scale with lower reconstruction error per block. + """ + scale_4_raw = scale_6.float() * self.SCALE_EXPANSION + if self.args.scale_dtype is not None: + scale_4 = round_to_quantized_type_dtype( + scale_4_raw, dtype=self.args.scale_dtype + ) + else: + scale_4 = scale_4_raw + + q_6 = fake_quantize( + observed, + scale_6.unsqueeze(-1), + zero_point.unsqueeze(-1), + self._token_args, + global_scale=global_scale, + ) + q_4 = fake_quantize( + observed, + scale_4.unsqueeze(-1), + zero_point.unsqueeze(-1), + self._token_args, + global_scale=global_scale, + ) + + obs_f = observed.float() + diff_6 = (q_6.float() - obs_f).abs_() + diff_4 = (q_4.float() - obs_f).abs_() + + if self.scale_selection == "mse": + err_6 = diff_6.pow_(2).sum(dim=(0, -1)) + err_4 = diff_4.pow_(2).sum(dim=(0, -1)) + elif self.scale_selection == "mae": + err_6 = diff_6.sum(dim=(0, -1)) + err_4 = diff_4.sum(dim=(0, -1)) + else: # abs_max + err_6 = diff_6.amax(dim=(0, -1)) + err_4 = diff_4.amax(dim=(0, -1)) + + select_4 = err_4 < err_6 + + scale_f = torch.where(select_4, scale_4.float(), scale_6.float()) + return scale_f.to(scale_6.dtype) diff --git a/src/llmcompressor/observers/mse.py b/src/llmcompressor/observers/mse.py index 319770824b..8ad23ba2c4 100644 --- a/src/llmcompressor/observers/mse.py +++ b/src/llmcompressor/observers/mse.py @@ -1,5 +1,6 @@ import torch from compressed_tensors.quantization import QuantizationStrategy +from compressed_tensors.quantization.utils import generate_gparam from torch import distributed as dist from llmcompressor.observers.base import Observer @@ -26,6 +27,10 @@ def __init__(self, *args, **kwargs): self.grid = observer_kwargs.get("grid", 100.0) self.norm = observer_kwargs.get("norm", 2.4) self.chunk_size = observer_kwargs.get("chunk_size", 5) + self.expand = observer_kwargs.get("expand", 1.0) + self.use_global_scale_prior = observer_kwargs.get( + "use_global_scale_prior", False + ) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") @@ -35,7 +40,20 @@ def __init__(self, *args, **kwargs): update={"strategy": QuantizationStrategy.TOKEN} ) + def _compute_approx_global_scale( + self, observed: torch.Tensor + ) -> torch.Tensor | None: + if ( + not self.use_global_scale_prior + or self.args.strategy != QuantizationStrategy.TENSOR_GROUP + ): + return None + absmax = observed.abs().max() + absmax = torch.clamp(absmax, min=torch.finfo(absmax.dtype).tiny) + return generate_gparam(-absmax.reshape(1), absmax.reshape(1)) + def update_statistics_from_observed(self, observed: torch.Tensor) -> None: + gs_prior = self._compute_approx_global_scale(observed) self.min_vals, self.max_vals = _grid_search_mse( observed, self.args, @@ -45,6 +63,8 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.grid, self.norm, self.chunk_size, + self.expand, + global_scale=gs_prior, ) @@ -69,6 +89,7 @@ def __init__(self, *args, **kwargs): self.grid = observer_kwargs.get("grid", 100.0) self.norm = observer_kwargs.get("norm", 2.4) self.chunk_size = observer_kwargs.get("chunk_size", 5) + self.expand = observer_kwargs.get("expand", 1.0) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") @@ -88,6 +109,7 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.grid, self.norm, self.chunk_size, + self.expand, ) if hasattr(self, "min_vals") and self.avg_constant != 1.0: @@ -96,3 +118,30 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.min_vals = min_vals self.max_vals = max_vals + + +@Observer.register("nvfp4_expanded_mse") +class NVFP4ExpandedMSEObserver(MemorylessMSEObserver): + """ + MSE observer with defaults tuned for NVFP4 range expansion. + + Searches from ``expand`` times the observed range down to + ``(1 - maxshrink) * expand`` times the observed range. + With the defaults below, this covers 1.8x down to ~0.8x of + the original per-group range in 112 search steps. + + Usage:: + + QuantizationArgs( + ... + observer="nvfp4_expanded_mse", + ) + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + observer_kwargs = self.args.observer_kwargs + self.expand = observer_kwargs.get("expand", 1.8) + self.maxshrink = observer_kwargs.get("maxshrink", 1 - 0.8 / 1.8) + self.grid = observer_kwargs.get("grid", 200.0) + self.patience = observer_kwargs.get("patience", 1000) diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index 088220cb43..b90adbcdfb 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -25,6 +25,8 @@ def _grid_search_mse( grid: float, norm: float, chunk_size: int, + expand: float = 1.0, + global_scale: torch.Tensor | None = None, ) -> MinMaxTuple: """Find per-channel min/max ranges that minimize quantization error. @@ -44,9 +46,14 @@ def _grid_search_mse( in shrink factors :param norm: exponent used when computing the error. norm = 2 approximates MSE :param chunk_size: number of grid steps per compiled call + :param expand: factor to scale the initial min/max range before searching. + Values > 1.0 let the search explore ranges wider than the observed + values (e.g. expand=2.0 starts at 2x the observed range). """ - min_val = torch.amin(observed, dim=(0, -1)) - max_val = torch.amax(observed, dim=(0, -1)) + original_min = torch.amin(observed, dim=(0, -1)) + original_max = torch.amax(observed, dim=(0, -1)) + min_val = original_min * expand + max_val = original_max * expand best_error = torch.full_like(min_val, torch.finfo(min_val.dtype).max) best_min_val = min_val.clone() best_max_val = max_val.clone() @@ -68,6 +75,7 @@ def _grid_search_mse( grid, norm, chunk_size, + global_scale, ) return _grid_search_eager( observed, @@ -82,6 +90,7 @@ def _grid_search_mse( patience, grid, norm, + global_scale, ) @@ -98,6 +107,7 @@ def _grid_search_eager( patience: int, grid: float, norm: float, + global_scale: torch.Tensor | None = None, ) -> MinMaxTuple: """Per-step grid search with boolean-indexing updates and early stopping.""" no_improve_count = 0 @@ -114,6 +124,7 @@ def _grid_search_eager( shrinked_min_val, shrinked_max_val, norm, + global_scale, ) improved = err < best_error @@ -144,6 +155,7 @@ def _grid_search_compiled( grid: float, norm: float, chunk_size: int, + global_scale: torch.Tensor | None = None, ) -> MinMaxTuple: """Chunked grid search using torch.compiled inner loop. @@ -183,6 +195,7 @@ def _grid_search_compiled( best_error, best_min_val, best_max_val, + global_scale, ) if torch.equal(prev_best, best_error): @@ -209,6 +222,7 @@ def _compute_chunk( best_error: torch.Tensor, best_min_val: torch.Tensor, best_max_val: torch.Tensor, + global_scale: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Evaluate ``chunk_size`` shrink-factor candidates in one call. @@ -231,6 +245,7 @@ def _compute_chunk( shrinked_min, shrinked_max, norm, + global_scale, ) improved = err < best_error @@ -248,6 +263,7 @@ def _calculate_error( shrinked_min: torch.Tensor, shrinked_max: torch.Tensor, norm: float, + global_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Fake-quantize ``observed`` using the given shrinked min/max range and return the per-channel error. @@ -258,7 +274,7 @@ def _calculate_error( min_vals=shrinked_min, max_vals=shrinked_max, quantization_args=args, - global_scale=None, + global_scale=global_scale, ) q = fake_quantize( @@ -266,6 +282,7 @@ def _calculate_error( candidate_scales.unsqueeze(-1), candidate_zero_points.unsqueeze(-1), token_args, + global_scale=global_scale, ).to(observed.dtype) err = torch.sum((q - observed).abs().pow(norm), dim=(0, -1))