src/agiws_neural_quant/training/ste.py

"""STE (Straight-Through Estimator) for quantization-aware training.

Forward: quantize (round) — non-differentiable
Backward: identity — gradient passes through as if round() was identity()

This allows training with fake quantization: the model "sees" quantized weights
in forward, but gradients flow to the latent fp32/fp16 weights in backward.
"""

import torch


class STEQuantize(torch.autograd.Function):
    """Straight-Through Estimator for symmetric quantization.

    Forward:  weight_int = clamp(round(weight / scale), -n_levels, n_levels)
              For asymmetric (int8/int4): clamp(min=-n_levels, max=n_levels-1)
              For symmetric (ternary {-1,0,+1}): clamp(min=-n_levels, max=n_levels)
    Backward: d(weight_int)/d(weight) = 1 (identity, STE)

    The clamp range depends on whether the quantization is symmetric (includes
    both -n_levels and +n_levels, like ternary {-1,0,+1}) or asymmetric (int8
    [-127,127], int4 [-7,7]). Use symmetric=True for ternary, False for int8/int4.
    """

    @staticmethod
    def forward(ctx, weight: torch.Tensor, scale: torch.Tensor, n_levels: int,
                symmetric: bool = False) -> torch.Tensor:
        """Quantize weight and dequantize back (fake quantization).

        Args:
            weight: float tensor (latent weight)
            scale: float tensor (per-channel or per-group scale)
            n_levels: max quantized value (127 for int8, 7 for int4, 1 for ternary)
            symmetric: if True, clamp to [-n_levels, n_levels] (ternary {-1,0,+1}).
                      if False, clamp to [-n_levels, n_levels-1] (int8 [-127,127]).

        Returns:
            Dequantized weight (fake quantized) — same shape as input.
        """
        # Quantize
        if symmetric:
            q = torch.clamp(
                torch.round(weight / scale),
                min=-n_levels,
                max=n_levels,
            )
        else:
            q = torch.clamp(
                torch.round(weight / scale),
                min=-n_levels,
                max=n_levels - 1,
            )
        # Dequantize back (fake quantization)
        return q * scale

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        # STE: gradient passes through as identity
        # grad for weight = grad_output, grad for scale = None, grad for n_levels = None
        return grad_output, None, None, None


def fake_quantize(weight: torch.Tensor, scale: torch.Tensor, n_levels: int,
                  symmetric: bool = False) -> torch.Tensor:
    """Apply fake quantization with STE for backward.

    In forward: simulates quantization error.
    In backward: gradient flows to weight (latent) as identity.

    Args:
        weight: latent weight (fp32 or fp16) — nn.Parameter
        scale: quantization scale (per-channel or per-group)
        n_levels: 127 for int8, 7 for int4, 1 for ternary
        symmetric: True for ternary {-1,0,+1} (clamp to [-n, n]).
                  False for int8/int4 (clamp to [-n, n-1]).

    Returns:
        Fake-quantized weight (same dtype as input).
    """
    return STEQuantize.apply(weight, scale, n_levels, symmetric)


class STECodebook(torch.autograd.Function):
    """STE for codebook quantization (learnable codebooks).

    Forward:  indices = argmin(|x - codebook|); deq = codebook[indices]
    Backward: gradient flows to x (latent weight) as identity (STE),
              AND to codebook (learnable) as identity for the selected entries.

    This allows training both the latent weight AND the codebook via gradient
    descent. The argmin (non-differentiable) is bypassed by STE.
    """

    @staticmethod
    def forward(ctx, weight: torch.Tensor, codebook: torch.Tensor,
                indices: torch.Tensor) -> torch.Tensor:
        """Codebook lookup with precomputed indices.

        Args:
            weight: latent weight (for STE identity backward).
            codebook: codebook tensor [K] or [K, dim].
            indices: precomputed argmin indices (same shape as weight, or [..., dim]).

        Returns:
            Dequantized weight: codebook[indices], same shape as weight.
        """
        ctx.save_for_backward(weight, codebook, indices)
        if codebook.dim() == 1:
            return codebook.to(weight.device)[indices.long()].to(torch.float32)
        else:
            # VQ: codebook [K, dim], indices [...]
            return codebook.to(weight.device)[indices.long()].to(torch.float32)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        weight, codebook, indices = ctx.saved_tensors
        # STE: grad to weight = grad_output (identity).
        grad_weight = grad_output
        # Grad to codebook: scatter-add gradient to selected entries.
        grad_codebook = torch.zeros_like(codebook)
        if codebook.dim() == 1:
            grad_codebook.scatter_add_(0, indices.long().flatten(),
                                       grad_output.flatten())
        else:
            # VQ: codebook [K, dim], indices [...]
            grad_codebook.scatter_add_(0, indices.long().reshape(-1, 1).expand(-1, codebook.shape[1]),
                                       grad_output.reshape(-1, codebook.shape[1]))
        return grad_weight, grad_codebook, None


def fake_codebook_quantize(weight: torch.Tensor, codebook: torch.Tensor,
                           indices: torch.Tensor) -> torch.Tensor:
    """Apply codebook fake quantization with STE for backward.

    Forward: codebook[indices] (precomputed argmin).
    Backward: gradient to weight (identity STE) + gradient to codebook (scatter-add).

    Args:
        weight: latent weight (nn.Parameter, for STE identity).
        codebook: codebook tensor [K] (for learnable codebook gradient).
        indices: precomputed argmin indices.

    Returns:
        Fake-quantized weight (codebook[indices]), differentiable to both
        weight and codebook.
    """
    return STECodebook.apply(weight, codebook, indices)