From a7025c92e7baf68cf7d48d00139e2e30c6d99042 Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Thu, 23 Jul 2026 11:50:03 -0400 Subject: [PATCH 01/12] MSE observer: global_scale fix and expand parameter - Thread global_scale through _calculate_error for correct NVFP4 two-level quantization error evaluation - Add expand parameter for wider-than-observed grid search ranges - Add global_scale_max observer_kwargs for custom FP8 scale cap - Add _CustomFP8ScaleData and global_scale_max support in get_global_scale Co-Authored-By: Claude Opus 4.6 --- src/llmcompressor/observers/base.py | 24 +++++++++++++- src/llmcompressor/observers/mse.py | 8 +++++ src/llmcompressor/observers/mse_quant.py | 42 ++++++++++++++++++++---- 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/src/llmcompressor/observers/base.py b/src/llmcompressor/observers/base.py index a06105f5bf..16d6f60fdb 100644 --- a/src/llmcompressor/observers/base.py +++ b/src/llmcompressor/observers/base.py @@ -5,6 +5,7 @@ from compressed_tensors import InternalModule from compressed_tensors.offload.dist_utils import as_broadcastable from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy +from compressed_tensors.quantization.quant_args import FloatArgs from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam from compressed_tensors.registry.registry import RegistryMixin from torch import distributed as dist @@ -26,6 +27,18 @@ class QParamsDict(TypedDict, total=False): global_scale: Optional[torch.Tensor] +class _CustomFP8ScaleData(FloatArgs): + exponent = 4 + mantissa = 3 + bits = 8 + max = 0 + min = 0 + + def __init__(self, max: float, min: float): + self.max = max + self.min = min + + class Observer(InternalModule, RegistryMixin): """ Base class for observers which compute quantization parameters given @@ -93,7 +106,16 @@ def get_global_scale(self) -> Optional[torch.Tensor]: global_absmax = torch.max(global_absmax, -stats["min_vals"].min()) global_absmax = torch.max(global_absmax, stats["max_vals"].max()) - return generate_gparam(-global_absmax.reshape(1), global_absmax.reshape(1)) + gparam_kwargs = {} + gs_max = (self.args.observer_kwargs or {}).get("global_scale_max") + if gs_max is not None: + gparam_kwargs["scale_data"] = _CustomFP8ScaleData( + max=gs_max, min=-gs_max + ) + + return generate_gparam( + -global_absmax.reshape(1), global_absmax.reshape(1), **gparam_kwargs + ) @torch.no_grad def get_qparams(self) -> QParamsDict: diff --git a/src/llmcompressor/observers/mse.py b/src/llmcompressor/observers/mse.py index 319770824b..e379fe425f 100644 --- a/src/llmcompressor/observers/mse.py +++ b/src/llmcompressor/observers/mse.py @@ -26,6 +26,8 @@ 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.global_scale_max = observer_kwargs.get("global_scale_max", None) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") @@ -45,6 +47,8 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.grid, self.norm, self.chunk_size, + self.expand, + self.global_scale_max, ) @@ -69,6 +73,8 @@ 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.global_scale_max = observer_kwargs.get("global_scale_max", None) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") @@ -88,6 +94,8 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.grid, self.norm, self.chunk_size, + self.expand, + self.global_scale_max, ) if hasattr(self, "min_vals") and self.avg_constant != 1.0: diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index 088220cb43..fad56cbd12 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -3,12 +3,12 @@ import torch import torch._dynamo.config import torch._dynamo.decorators -from compressed_tensors.quantization import QuantizationArgs +from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy from compressed_tensors.quantization.lifecycle import fake_quantize -from compressed_tensors.quantization.utils import calculate_qparams +from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam from llmcompressor.core import active_session -from llmcompressor.observers.base import MinMaxTuple +from llmcompressor.observers.base import MinMaxTuple, _CustomFP8ScaleData # Allow torch.compile to handle scalar conversions inside # compressed_tensors' calculate_qparams (float(bit_range)). @@ -25,6 +25,8 @@ def _grid_search_mse( grid: float, norm: float, chunk_size: int, + expand: float = 1.0, + global_scale_max: float | None = None, ) -> MinMaxTuple: """Find per-channel min/max ranges that minimize quantization error. @@ -44,13 +46,31 @@ 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). + :param global_scale_max: optional cap for the FP8 global scale 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() + global_scale = None + if args.strategy == QuantizationStrategy.TENSOR_GROUP: + global_absmax = torch.max(-original_min.min(), original_max.max()) + gparam_kwargs = {} + if global_scale_max is not None: + gparam_kwargs["scale_data"] = _CustomFP8ScaleData( + max=global_scale_max, min=-global_scale_max + ) + global_scale = generate_gparam( + -global_absmax.reshape(1), global_absmax.reshape(1), **gparam_kwargs + ) + total_steps = int(maxshrink * grid) if active_session().state.enable_compile: @@ -68,6 +88,7 @@ def _grid_search_mse( grid, norm, chunk_size, + global_scale, ) return _grid_search_eager( observed, @@ -82,6 +103,7 @@ def _grid_search_mse( patience, grid, norm, + global_scale, ) @@ -98,6 +120,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 +137,7 @@ def _grid_search_eager( shrinked_min_val, shrinked_max_val, norm, + global_scale, ) improved = err < best_error @@ -144,6 +168,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 +208,7 @@ def _grid_search_compiled( best_error, best_min_val, best_max_val, + global_scale, ) if torch.equal(prev_best, best_error): @@ -209,6 +235,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 +258,7 @@ def _compute_chunk( shrinked_min, shrinked_max, norm, + global_scale, ) improved = err < best_error @@ -248,6 +276,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 +287,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 +295,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)) From feb1291f77504957b69f17da9a999a729c1e06ea Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Fri, 24 Jul 2026 09:39:43 -0400 Subject: [PATCH 02/12] lint Signed-off-by: Roderick Wu --- src/llmcompressor/observers/base.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/llmcompressor/observers/base.py b/src/llmcompressor/observers/base.py index 16d6f60fdb..e3926b9dd9 100644 --- a/src/llmcompressor/observers/base.py +++ b/src/llmcompressor/observers/base.py @@ -109,9 +109,7 @@ def get_global_scale(self) -> Optional[torch.Tensor]: gparam_kwargs = {} gs_max = (self.args.observer_kwargs or {}).get("global_scale_max") if gs_max is not None: - gparam_kwargs["scale_data"] = _CustomFP8ScaleData( - max=gs_max, min=-gs_max - ) + gparam_kwargs["scale_data"] = _CustomFP8ScaleData(max=gs_max, min=-gs_max) return generate_gparam( -global_absmax.reshape(1), global_absmax.reshape(1), **gparam_kwargs From 9cab2aca4485b80fb54e6e13f4b01c7889d2853d Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Fri, 24 Jul 2026 16:42:59 -0400 Subject: [PATCH 03/12] Update src/llmcompressor/observers/mse_quant.py Co-authored-by: Brian Dellabetta Signed-off-by: Roderick Wu --- src/llmcompressor/observers/mse_quant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index fad56cbd12..16c4c7dbae 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -8,7 +8,7 @@ from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam from llmcompressor.core import active_session -from llmcompressor.observers.base import MinMaxTuple, _CustomFP8ScaleData +from llmcompressor.observers.base import MinMaxTuple, CustomFP8ScaleData # Allow torch.compile to handle scalar conversions inside # compressed_tensors' calculate_qparams (float(bit_range)). From 4ca6147d87ea52b90f7613282418829dc263a0fe Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Fri, 24 Jul 2026 16:46:43 -0400 Subject: [PATCH 04/12] lint Signed-off-by: Roderick Wu --- src/llmcompressor/observers/base.py | 4 ++-- src/llmcompressor/observers/mse_quant.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/llmcompressor/observers/base.py b/src/llmcompressor/observers/base.py index e3926b9dd9..4c3817d588 100644 --- a/src/llmcompressor/observers/base.py +++ b/src/llmcompressor/observers/base.py @@ -27,7 +27,7 @@ class QParamsDict(TypedDict, total=False): global_scale: Optional[torch.Tensor] -class _CustomFP8ScaleData(FloatArgs): +class CustomFP8ScaleData(FloatArgs): exponent = 4 mantissa = 3 bits = 8 @@ -109,7 +109,7 @@ def get_global_scale(self) -> Optional[torch.Tensor]: gparam_kwargs = {} gs_max = (self.args.observer_kwargs or {}).get("global_scale_max") if gs_max is not None: - gparam_kwargs["scale_data"] = _CustomFP8ScaleData(max=gs_max, min=-gs_max) + gparam_kwargs["scale_data"] = CustomFP8ScaleData(max=gs_max, min=-gs_max) return generate_gparam( -global_absmax.reshape(1), global_absmax.reshape(1), **gparam_kwargs diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index 16c4c7dbae..295cc17ac0 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -8,7 +8,7 @@ from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam from llmcompressor.core import active_session -from llmcompressor.observers.base import MinMaxTuple, CustomFP8ScaleData +from llmcompressor.observers.base import CustomFP8ScaleData, MinMaxTuple # Allow torch.compile to handle scalar conversions inside # compressed_tensors' calculate_qparams (float(bit_range)). @@ -64,7 +64,7 @@ def _grid_search_mse( global_absmax = torch.max(-original_min.min(), original_max.max()) gparam_kwargs = {} if global_scale_max is not None: - gparam_kwargs["scale_data"] = _CustomFP8ScaleData( + gparam_kwargs["scale_data"] = CustomFP8ScaleData( max=global_scale_max, min=-global_scale_max ) global_scale = generate_gparam( From fd61a44f9dc16304f2c933e89d2c8625bab33ff9 Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Tue, 4 Aug 2026 19:15:40 +0000 Subject: [PATCH 05/12] Remove fp8 scale apis Signed-off-by: Roderick Wu --- src/llmcompressor/observers/base.py | 21 ++------------------- src/llmcompressor/observers/mse.py | 4 ---- src/llmcompressor/observers/mse_quant.py | 11 ++--------- 3 files changed, 4 insertions(+), 32 deletions(-) diff --git a/src/llmcompressor/observers/base.py b/src/llmcompressor/observers/base.py index 4c3817d588..df2c455454 100644 --- a/src/llmcompressor/observers/base.py +++ b/src/llmcompressor/observers/base.py @@ -5,7 +5,6 @@ from compressed_tensors import InternalModule from compressed_tensors.offload.dist_utils import as_broadcastable from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy -from compressed_tensors.quantization.quant_args import FloatArgs from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam from compressed_tensors.registry.registry import RegistryMixin from torch import distributed as dist @@ -27,18 +26,6 @@ class QParamsDict(TypedDict, total=False): global_scale: Optional[torch.Tensor] -class CustomFP8ScaleData(FloatArgs): - exponent = 4 - mantissa = 3 - bits = 8 - max = 0 - min = 0 - - def __init__(self, max: float, min: float): - self.max = max - self.min = min - - class Observer(InternalModule, RegistryMixin): """ Base class for observers which compute quantization parameters given @@ -101,18 +88,14 @@ def get_global_scale(self) -> Optional[torch.Tensor]: return None 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()) - gparam_kwargs = {} - gs_max = (self.args.observer_kwargs or {}).get("global_scale_max") - if gs_max is not None: - gparam_kwargs["scale_data"] = CustomFP8ScaleData(max=gs_max, min=-gs_max) - return generate_gparam( - -global_absmax.reshape(1), global_absmax.reshape(1), **gparam_kwargs + -global_absmax.reshape(1), global_absmax.reshape(1) ) @torch.no_grad diff --git a/src/llmcompressor/observers/mse.py b/src/llmcompressor/observers/mse.py index e379fe425f..8c5a961bfc 100644 --- a/src/llmcompressor/observers/mse.py +++ b/src/llmcompressor/observers/mse.py @@ -27,7 +27,6 @@ def __init__(self, *args, **kwargs): 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.global_scale_max = observer_kwargs.get("global_scale_max", None) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") @@ -48,7 +47,6 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.norm, self.chunk_size, self.expand, - self.global_scale_max, ) @@ -74,7 +72,6 @@ def __init__(self, *args, **kwargs): 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.global_scale_max = observer_kwargs.get("global_scale_max", None) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") @@ -95,7 +92,6 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.norm, self.chunk_size, self.expand, - self.global_scale_max, ) if hasattr(self, "min_vals") and self.avg_constant != 1.0: diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index 295cc17ac0..c1e25cdedc 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -8,7 +8,7 @@ from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam from llmcompressor.core import active_session -from llmcompressor.observers.base import CustomFP8ScaleData, MinMaxTuple +from llmcompressor.observers.base import MinMaxTuple # Allow torch.compile to handle scalar conversions inside # compressed_tensors' calculate_qparams (float(bit_range)). @@ -26,7 +26,6 @@ def _grid_search_mse( norm: float, chunk_size: int, expand: float = 1.0, - global_scale_max: float | None = None, ) -> MinMaxTuple: """Find per-channel min/max ranges that minimize quantization error. @@ -49,7 +48,6 @@ def _grid_search_mse( :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). - :param global_scale_max: optional cap for the FP8 global scale range """ original_min = torch.amin(observed, dim=(0, -1)) original_max = torch.amax(observed, dim=(0, -1)) @@ -62,13 +60,8 @@ def _grid_search_mse( global_scale = None if args.strategy == QuantizationStrategy.TENSOR_GROUP: global_absmax = torch.max(-original_min.min(), original_max.max()) - gparam_kwargs = {} - if global_scale_max is not None: - gparam_kwargs["scale_data"] = CustomFP8ScaleData( - max=global_scale_max, min=-global_scale_max - ) global_scale = generate_gparam( - -global_absmax.reshape(1), global_absmax.reshape(1), **gparam_kwargs + -global_absmax.reshape(1), global_absmax.reshape(1) ) total_steps = int(maxshrink * grid) From 32632f8297b30c4a476bec5dea398d72de7196f2 Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Tue, 4 Aug 2026 19:23:14 +0000 Subject: [PATCH 06/12] just expands? Signed-off-by: Roderick Wu --- src/llmcompressor/observers/mse_quant.py | 22 ++-------------------- 1 file changed, 2 insertions(+), 20 deletions(-) diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index c1e25cdedc..df3de4daf0 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -3,9 +3,9 @@ import torch import torch._dynamo.config import torch._dynamo.decorators -from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy +from compressed_tensors.quantization import QuantizationArgs from compressed_tensors.quantization.lifecycle import fake_quantize -from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam +from compressed_tensors.quantization.utils import calculate_qparams from llmcompressor.core import active_session from llmcompressor.observers.base import MinMaxTuple @@ -57,13 +57,6 @@ def _grid_search_mse( best_min_val = min_val.clone() best_max_val = max_val.clone() - global_scale = None - if args.strategy == QuantizationStrategy.TENSOR_GROUP: - global_absmax = torch.max(-original_min.min(), original_max.max()) - global_scale = generate_gparam( - -global_absmax.reshape(1), global_absmax.reshape(1) - ) - total_steps = int(maxshrink * grid) if active_session().state.enable_compile: @@ -81,7 +74,6 @@ def _grid_search_mse( grid, norm, chunk_size, - global_scale, ) return _grid_search_eager( observed, @@ -96,7 +88,6 @@ def _grid_search_mse( patience, grid, norm, - global_scale, ) @@ -113,7 +104,6 @@ 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 @@ -130,7 +120,6 @@ def _grid_search_eager( shrinked_min_val, shrinked_max_val, norm, - global_scale, ) improved = err < best_error @@ -161,7 +150,6 @@ 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. @@ -201,7 +189,6 @@ def _grid_search_compiled( best_error, best_min_val, best_max_val, - global_scale, ) if torch.equal(prev_best, best_error): @@ -228,7 +215,6 @@ 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. @@ -251,7 +237,6 @@ def _compute_chunk( shrinked_min, shrinked_max, norm, - global_scale, ) improved = err < best_error @@ -269,7 +254,6 @@ 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. @@ -280,7 +264,6 @@ def _calculate_error( min_vals=shrinked_min, max_vals=shrinked_max, quantization_args=args, - global_scale=global_scale, ) q = fake_quantize( @@ -288,7 +271,6 @@ 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)) From 67dd2a218b042edd0733a9fb57174a119752062b Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Tue, 4 Aug 2026 19:31:13 +0000 Subject: [PATCH 07/12] asldfjkl Signed-off-by: Roderick Wu --- src/llmcompressor/observers/mse_quant.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index df3de4daf0..06815de415 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -264,6 +264,7 @@ def _calculate_error( min_vals=shrinked_min, max_vals=shrinked_max, quantization_args=args, + global_scale=None, ) q = fake_quantize( From 14a63614acce0bc4f98589efbd57b7d868333021 Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Tue, 4 Aug 2026 19:50:40 +0000 Subject: [PATCH 08/12] modify expand/maxshrink defaults Signed-off-by: Roderick Wu --- src/llmcompressor/observers/base.py | 5 +---- src/llmcompressor/observers/mse.py | 8 ++++---- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/src/llmcompressor/observers/base.py b/src/llmcompressor/observers/base.py index df2c455454..a06105f5bf 100644 --- a/src/llmcompressor/observers/base.py +++ b/src/llmcompressor/observers/base.py @@ -88,15 +88,12 @@ def get_global_scale(self) -> Optional[torch.Tensor]: return None 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()) - return generate_gparam( - -global_absmax.reshape(1), global_absmax.reshape(1) - ) + return generate_gparam(-global_absmax.reshape(1), global_absmax.reshape(1)) @torch.no_grad def get_qparams(self) -> QParamsDict: diff --git a/src/llmcompressor/observers/mse.py b/src/llmcompressor/observers/mse.py index 8c5a961bfc..d8e2ef84f0 100644 --- a/src/llmcompressor/observers/mse.py +++ b/src/llmcompressor/observers/mse.py @@ -21,12 +21,12 @@ class MemorylessMSEObserver(Observer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) observer_kwargs = self.args.observer_kwargs - self.maxshrink = observer_kwargs.get("maxshrink", 0.20) + self.maxshrink = observer_kwargs.get("maxshrink", 0.6) self.patience = observer_kwargs.get("patience", 5) 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.expand = observer_kwargs.get("expand", 2.0) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") @@ -66,12 +66,12 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.avg_constant = self.args.observer_kwargs.get("averaging_constant", 0.01) observer_kwargs = self.args.observer_kwargs - self.maxshrink = observer_kwargs.get("maxshrink", 0.20) + self.maxshrink = observer_kwargs.get("maxshrink", 0.6) self.patience = observer_kwargs.get("patience", 5) 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.expand = observer_kwargs.get("expand", 2.0) if self.chunk_size <= 0: raise ValueError(f"chunk_size must be positive, got {self.chunk_size}") From 4a6ec4fe75cf07e4ecdabed1c41944aeac97f921 Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Wed, 5 Aug 2026 20:30:42 +0000 Subject: [PATCH 09/12] dslkfjdls --- src/llmcompressor/observers/__init__.py | 1 + src/llmcompressor/observers/fouroversix.py | 178 +++++++++++++++++++++ src/llmcompressor/observers/mse.py | 33 +++- 3 files changed, 208 insertions(+), 4 deletions(-) create mode 100644 src/llmcompressor/observers/fouroversix.py 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 d8e2ef84f0..14130e8319 100644 --- a/src/llmcompressor/observers/mse.py +++ b/src/llmcompressor/observers/mse.py @@ -21,12 +21,12 @@ class MemorylessMSEObserver(Observer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) observer_kwargs = self.args.observer_kwargs - self.maxshrink = observer_kwargs.get("maxshrink", 0.6) + self.maxshrink = observer_kwargs.get("maxshrink", 0.20) self.patience = observer_kwargs.get("patience", 5) 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", 2.0) + 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}") @@ -66,12 +66,12 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.avg_constant = self.args.observer_kwargs.get("averaging_constant", 0.01) observer_kwargs = self.args.observer_kwargs - self.maxshrink = observer_kwargs.get("maxshrink", 0.6) + self.maxshrink = observer_kwargs.get("maxshrink", 0.20) self.patience = observer_kwargs.get("patience", 5) 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", 2.0) + 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}") @@ -100,3 +100,28 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.min_vals = min_vals self.max_vals = max_vals + + +@Observer.register("nvfp4_mse") +class NVFP4MSEObserver(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 (expand=1.8, maxshrink=0.56), this covers + 1.8x down to ~0.8x of the original per-group range. + + Usage:: + + QuantizationArgs( + ... + observer="nvfp4_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", 0.56) From 48af6ddc05e30491b1a89d7ca7e02ed9d9470acd Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Fri, 7 Aug 2026 17:23:06 +0000 Subject: [PATCH 10/12] =?UTF-8?q?Rename=20nvfp4=5Fmse=20=E2=86=92=20nvfp4?= =?UTF-8?q?=5Fexpanded=5Fmse,=20update=20defaults=20and=20docstring?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Register name: nvfp4_mse → nvfp4_expanded_mse - maxshrink: 0.56 → 1 - 0.8/1.8 (sweeps 1.8x→0.8x) - grid: 100 → 200 (finer search resolution) - patience: 5 → 1000 (full range sweep) - Fix docstring to reference correct observer name Co-Authored-By: Claude Opus 4.6 --- src/llmcompressor/observers/mse.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/llmcompressor/observers/mse.py b/src/llmcompressor/observers/mse.py index 14130e8319..84df791795 100644 --- a/src/llmcompressor/observers/mse.py +++ b/src/llmcompressor/observers/mse.py @@ -102,21 +102,21 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.max_vals = max_vals -@Observer.register("nvfp4_mse") -class NVFP4MSEObserver(MemorylessMSEObserver): +@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 (expand=1.8, maxshrink=0.56), this covers - 1.8x down to ~0.8x of the original per-group 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_mse", + observer="nvfp4_expanded_mse", ) """ @@ -124,4 +124,6 @@ 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", 0.56) + 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) From 6c35d08bb5632e6094b7dd719acf50df12ffc78f Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Mon, 10 Aug 2026 00:06:59 -0400 Subject: [PATCH 11/12] test changes --- src/llmcompressor/observers/mse.py | 18 ++++++++++++++++++ src/llmcompressor/observers/mse_quant.py | 13 ++++++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/src/llmcompressor/observers/mse.py b/src/llmcompressor/observers/mse.py index 84df791795..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 @@ -27,6 +28,9 @@ def __init__(self, *args, **kwargs): 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}") @@ -36,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, @@ -47,6 +64,7 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None: self.norm, self.chunk_size, self.expand, + global_scale=gs_prior, ) diff --git a/src/llmcompressor/observers/mse_quant.py b/src/llmcompressor/observers/mse_quant.py index 06815de415..b90adbcdfb 100644 --- a/src/llmcompressor/observers/mse_quant.py +++ b/src/llmcompressor/observers/mse_quant.py @@ -26,6 +26,7 @@ def _grid_search_mse( 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. @@ -74,6 +75,7 @@ def _grid_search_mse( grid, norm, chunk_size, + global_scale, ) return _grid_search_eager( observed, @@ -88,6 +90,7 @@ def _grid_search_mse( patience, grid, norm, + global_scale, ) @@ -104,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 @@ -120,6 +124,7 @@ def _grid_search_eager( shrinked_min_val, shrinked_max_val, norm, + global_scale, ) improved = err < best_error @@ -150,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. @@ -189,6 +195,7 @@ def _grid_search_compiled( best_error, best_min_val, best_max_val, + global_scale, ) if torch.equal(prev_best, best_error): @@ -215,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. @@ -237,6 +245,7 @@ def _compute_chunk( shrinked_min, shrinked_max, norm, + global_scale, ) improved = err < best_error @@ -254,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. @@ -264,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( @@ -272,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)) From 64d27c4a81127fb3116d2a1f33e2a20c08753f5d Mon Sep 17 00:00:00 2001 From: Roderick Wu Date: Mon, 10 Aug 2026 11:39:43 -0400 Subject: [PATCH 12/12] eval scripts --- compress_all.py | 235 ++++++++++++++++++++++++++++++++++++++++++++++++ eval_all.py | 158 ++++++++++++++++++++++++++++++++ 2 files changed, 393 insertions(+) create mode 100644 compress_all.py create mode 100644 eval_all.py 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()