From b84ec506523b708da65082b195ea626cf3b8aa59 Mon Sep 17 00:00:00 2001 From: Benjamin Luo Date: Sun, 12 Jul 2026 13:58:07 -0700 Subject: [PATCH 1/3] Fix quantize_tensor logic - quantize_tensor now uses scale quantization and is symmetric in its range. - PGD now uses 10 steps rather than 15 --- src/QuantAdv.py | 18 ++++++++---------- src/config.py | 2 +- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/src/QuantAdv.py b/src/QuantAdv.py index 36fd025..98c0ba9 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -147,16 +147,14 @@ def backward(ctx, grad_output): return grad_output -def quantize_tensor(t, bits, use_ste): - if bits is None: - return t - qmax = 2 ** (bits - 1) - 1 - scale = torch.clamp(t.detach().abs().max() / qmax, min=QUANT_SCALE_MIN) - t_scaled = t / scale - t_round = FakeQuantSTE.apply(t_scaled) if use_ste else torch.round(t_scaled) - t_round = torch.clamp(t_round, -qmax - 1, qmax) - return t_round * scale - +def quantize_tensor(t: torch.Tensor, bits: int, alpha: int): + """ + Quantizes a tensor t to a given number of bits through + scale quantization. See https://arxiv.org/pdf/2004.09602 + """ + max: int = 2 ** (bits - 1) - 1 + scale: float = max / alpha + return torch.clamp(torch.round(scale * t), min=-max, max=max) def chaotic_sequence_like(t, seed=CHAOTIC_QUANT_SEED, map_name=CHAOTIC_QUANT_MAP): n = t.numel() diff --git a/src/config.py b/src/config.py index 8b4aed9..0b5c2fc 100644 --- a/src/config.py +++ b/src/config.py @@ -118,7 +118,7 @@ DEFAULT_EPS = 8 / 255 PGD_ALPHA = 2 / 255 -PGD_STEPS = 15 +PGD_STEPS = 10 PGD_RANDOM_START = True QAT_BITS = 8 From bea8b7f213992ddb1b827f117ce690d5ffa57e2e Mon Sep 17 00:00:00 2001 From: Benjamin Luo Date: Sun, 12 Jul 2026 14:04:40 -0700 Subject: [PATCH 2/3] Re-add quantize_tensor --- src/QuantAdv.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/QuantAdv.py b/src/QuantAdv.py index c38c067..a4230a1 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -163,14 +163,6 @@ def _torchao_int_dtype(bits): return dtype -def quantize_tensor(t: torch.Tensor, bits: int, alpha: int): - """ - Quantizes a tensor t to a given number of bits through - scale quantization. See https://arxiv.org/pdf/2004.09602 - """ - max: int = 2 ** (bits - 1) - 1 - scale: float = max / alpha - return torch.clamp(torch.round(scale * t), min=-max, max=max) def _make_torchao_fake_quantizer(bits, *, role): """Build TorchAO fake quantization for a weight or activation tensor. """ @@ -200,6 +192,14 @@ def _hard_or_ste(fake_quantizer, tensor, use_ste): quantized = fake_quantizer(tensor) return quantized if use_ste else quantized.detach() +def quantize_tensor(t: torch.Tensor, bits: int, alpha: int): + """ + Quantizes a tensor t to a given number of bits through + scale quantization. See https://arxiv.org/pdf/2004.09602 + """ + max: int = 2 ** (bits - 1) - 1 + scale: float = max / alpha + return torch.clamp(torch.round(scale * t), min=-max, max=max) def chaotic_sequence_like(t, seed=CHAOTIC_QUANT_SEED, map_name=CHAOTIC_QUANT_MAP): """Create a deterministic chaotic sequence with the same shape as ``t``.""" From ebaa484a3100031a6a024f2906a67238317fe811 Mon Sep 17 00:00:00 2001 From: Benjamin Luo Date: Sun, 12 Jul 2026 14:10:05 -0700 Subject: [PATCH 3/3] Make specifying alpha in quantize_tensor optional --- src/QuantAdv.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/QuantAdv.py b/src/QuantAdv.py index a4230a1..79eb650 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -192,12 +192,17 @@ def _hard_or_ste(fake_quantizer, tensor, use_ste): quantized = fake_quantizer(tensor) return quantized if use_ste else quantized.detach() -def quantize_tensor(t: torch.Tensor, bits: int, alpha: int): +def quantize_tensor(t: torch.Tensor, bits: int, alpha: int | None = None): """ Quantizes a tensor t to a given number of bits through scale quantization. See https://arxiv.org/pdf/2004.09602 """ max: int = 2 ** (bits - 1) - 1 + + # Rescale to 0-1 range by default if no alpha is specified. + if alpha is None: + alpha = max + scale: float = max / alpha return torch.clamp(torch.round(scale * t), min=-max, max=max)