Remove nested directory: BitTransformerLM/bit_transformer/quantization.py

BitTransformerLM/bit_transformer/quantization.py

@@ -1,89 +0,0 @@
-import torch
-import torch.nn as nn
-
-from torch.ao.quantization.fake_quantize import FakeQuantize
-from torch.ao.quantization.observer import MinMaxObserver
-from torch.ao.quantization.qconfig import QConfig
-from torch.ao.quantization import convert
-
-from .model import BitTransformerLM
-
-
-def quantize_dynamic(model: BitTransformerLM, dtype: torch.dtype = torch.qint8) -> BitTransformerLM:
-    """Return a dynamically quantized copy of the model for inference."""
-    quantized = torch.quantization.quantize_dynamic(
-        model, {nn.Linear}, dtype=dtype
-    )
-    return quantized
-
-
-class FourBitObserver(MinMaxObserver):
-    """Min-max observer configured for 4-bit quantization."""
-
-    def __init__(self, **kwargs):
-        super().__init__(
-            quant_min=0,
-            quant_max=15,
-            dtype=torch.quint8,
-            qscheme=torch.per_tensor_affine,
-            **kwargs,
-        )
-
-
-FourBitFakeQuantize = FakeQuantize.with_args(observer=FourBitObserver)
-
-four_bit_qconfig = QConfig(activation=FourBitFakeQuantize, weight=FourBitFakeQuantize)
-
-
-class QATLinear(nn.Linear):
-    """Linear layer with fake quantization for QAT."""
-
-    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
-        super().__init__(in_features, out_features, bias)
-        self.weight_fake_quant = FourBitFakeQuantize()
-        self.activation_post_process = FourBitFakeQuantize()
-
-    @classmethod
-    def from_float(cls, mod: nn.Linear) -> "QATLinear":
-        qat = cls(mod.in_features, mod.out_features, mod.bias is not None)
-        qat.weight = mod.weight
-        qat.bias = mod.bias
-        return qat
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.activation_post_process(x)
-        w = self.weight_fake_quant(self.weight)
-        return nn.functional.linear(x, w, self.bias)
-
-
-def prepare_qat_fx(model: BitTransformerLM) -> BitTransformerLM:
-    """Prepare BitTransformerLM for quantization-aware training."""
-
-    for name, module in model.named_children():
-        if isinstance(module, nn.Linear):
-            setattr(model, name, QATLinear.from_float(module))
-        else:
-            prepare_qat_fx(module)
-    return model
-
-
-def convert_qat_fx(model: BitTransformerLM) -> BitTransformerLM:
-    """Convert a QAT-prepared model to a quantized version."""
-
-    for name, module in model.named_children():
-        if isinstance(module, QATLinear):
-            w = module.weight.data
-            qmin, qmax = 0, 15
-            min_w = w.min()
-            max_w = w.max()
-            scale = (max_w - min_w) / (qmax - qmin + 1e-8)
-            zero_point = qmin - torch.round(min_w / scale)
-            q_w = torch.clamp(torch.round(w / scale + zero_point), qmin, qmax)
-            new_mod = nn.Linear(module.in_features, module.out_features, module.bias is not None)
-            new_mod.weight = nn.Parameter((q_w - zero_point) * scale)
-            if module.bias is not None:
-                new_mod.bias = nn.Parameter(module.bias.data)
-            setattr(model, name, new_mod)
-        else:
-            convert_qat_fx(module)
-    return model