FakeQuantize

Observer + straight-through fake-quantization, toggleable for QAT.

The building block of quantization-aware training (QAT). A FakeQuantize sits on a weight or activation and, on every forward, does two things: it lets its wrapped observer refresh (scale, zero_point) from the live statistics, then it applies lucid.quantization.fake_quantize — rounding the tensor onto the integer grid and back to float — so the downstream graph sees the quantization error while training still runs in float. This is what lets a network learn weights that are robust to the eventual int8 inference numerics: the forward simulates quantization, and a straight-through estimator (STE) keeps the operation differentiable so the rounding does not block gradients.

Observer vs fake-quant, decoupled. The two behaviours are independently gated by enable_observer / enable_fake_quant, which drives the standard QAT schedule: observe + fake-quant early so statistics settle and the net adapts, then disable_observer to freeze the grid once ranges converge (and, in fused stacks, freeze BN stats). With the observer off but fake-quant on, the layer keeps quantizing against the last-learned qparams; with fake-quant off it is a pure passthrough (the observer may still watch).

The straight-through fake-quant. Writing $S, Z$ for the current (scale, zero_point) and [q_\min, q_\max] for the grid, the forward computes the quantize→dequantize round-trip in the float domain:

The $\operatorname{round}$ step has a zero (a.e.) derivative, which would kill training, so the backward substitutes the straight-through estimator — gradient passes through unchanged inside the grid and is zeroed where the code saturates:

so $\nabla_x = \nabla_{\hat{x}} \cdot \mathbb{1}[\,\text{code in range}\,]$. The qparams $S, Z$ are treated as constants of the step (they are refreshed by the observer, not learned through this gradient).