FakeQuantize Callback

Train through a simulated width, and bake it in

Overview

FakeQuantizeCallback trains a model through [FakeQuantizer](https://FasterAI-Labs.github.io/fasterai/quantize/fake_quantizer.html#fakequantizer)’s rounding. On every forward the weights are rounded onto the grid of the width you asked for, the loss is computed on the rounded model, and the gradient reaches an untouched floating-point copy — the master — which the optimizer keeps training. When the fit ends the rounding is baked in, leaving exactly the model FakeQuantizer.quantize_model() leaves: ordinary modules, floating-point dtype, rounded weights.

That is the one question FakeQuantizer cannot answer. Rounding a trained model onto a narrow grid costs accuracy; letting the weights adapt to the rounding they will be subject to may win some of it back. How much, on your model, is a number only your fit can produce.

Everything the post-training page says still holds: nothing is packed, no kernel is swapped, and biases stay in floating point. Two more things this callback deliberately does not do:

  • BatchNorm is untouched — not folded, not frozen. A deployed flow folds it before quantizing, which changes the per-channel weight ranges the rounding sees; fold it yourself first with [BN_Folder](https://FasterAI-Labs.github.io/fasterai/misc/bn_folding.html#bn_folder) if you want the model a deployed flow would produce.
  • The activation rounding sits on the module’s own output, before whatever activation function follows it. No integer runtime does that — it fuses conv + relu and quantizes once, after the fusion — so an activation width measured here is a proxy, not the arithmetic a backend will run.

The mechanism

The weight rounding is a torch.nn.utils.parametrize parametrization: the trained parameter becomes m.parametrizations.weight.original, and m.weight is that master rounded, on every access. Three consequences worth knowing:

  • The optimizer needs no rebinding. register_parametrization keeps the same Parameter object as the master, and remove_parametrizations hands the same object back, so the optimizer fit built before the callback ran holds exactly the tensors being trained.
  • The gradient is straight-through. The forward uses the value _fake_quantize writes — the same function the post-training path calls, so both round identically — and the backward passes the gradient to the master unchanged. Nothing is ever clipped on the weight side: the grid is recomputed from the master at every forward, so it always covers it.
  • Anything reading a weight must read the master. m.weight is a rounded snapshot, and a write to it is silently discarded. [Sparsifier](https://FasterAI-Labs.github.io/fasterai/sparse/sparsifier.html#sparsifier) and [Criteria](https://FasterAI-Labs.github.io/fasterai/core/criteria.html#criteria) go through _master, which is why sparsification and QAT compose.

Activations are observed the way torch’s MovingAverageMinMaxObserver observes them: the range moves towards each batch by act_averaging (0.01 by default), rather than being the absolute min and max of everything ever seen — one outlier batch would otherwise widen the grid permanently, with nothing to narrow it again. Pass act_averaging=None for the absolute min and max. Once a range has been observed, only training batches move it — the module has to be in training mode — and activations outside the frozen grid are clipped, which means they get no gradient at all.

Moment What happens
before_fit builds the FakeQuantizer, parametrizes every weight it rounds, hooks every activation
every training batch the forward rounds, the gradient reaches the master, the optimizer moves it
pct_train >= freeze_act_pct the activation scales stop moving, and the fit finishes on the grid it will keep
after_fit bake(): the trained master is rounded into an ordinary weight, and the spec is attached

pct_train is incremented after each batch, so the largest value a fit of N steps ever shows is (N-1)/N: a 4-step fit tops out at 0.75 and never reaches the default freeze_act_pct=0.9. The scales are therefore also frozen unconditionally by bake(), so a model always leaves a fit on frozen scales. With act_bits=None, or observer='dynamic', there is no scale to freeze and freeze_act_pct does nothing.

The arguments are [FakeQuantizer](https://FasterAI-Labs.github.io/fasterai/quantize/fake_quantizer.html#fakequantizer)’s, spelled the same way and validated by the same code — the callback builds one at the start of the fit and exposes it as self.fake_quantizer, so a width, an axis or a group_size this grammar cannot honor is refused there, with the same sentence, and before a single weight is touched. Three arguments are the callback’s own, and are checked at construction:

  • freeze_act_pct: how far into the fit the activation scales stop moving. A no-op when there are no static activation scales, i.e. with act_bits=None or observer='dynamic'.
  • act_averaging: how far the observed range moves towards each batch, 0.01 by default. None takes the absolute min and max instead, which one outlier batch widens for good.
  • model: the model to round, when it is not learn.model.

There is no schedule: the rounding is applied at full width from the first batch, the way prepare_qat_fx does.

Usage

from fasterai.quantize.fake_quantize_callback import FakeQuantizeCallback

cb = FakeQuantizeCallback(weight_bits=4)            # weight-only QAT at 4 bits, per channel
learn.fit(5, cbs=[cb])

cb = FakeQuantizeCallback(8, 8, freeze_act_pct=0.9)  # weights and activations, scales frozen at 90%
learn.fit(5, cbs=[cb])

Bind the callback: fit takes it off the learner when it returns, and everything below is called on it.

After the fit learn.model is an ordinary model holding its rounded weights, tagged with the precision it trained through:

from fasterai.core.precision import fake_quant_spec

fake_quant_spec(learn.model)          # FakeQuantSpec(..., trained=True)
learn.validate()                      # what the trained width costs on your metric
cb.fake_quantizer.remove()            # the TRAINED floating-point weights back

trained=True is the one thing a spec from this callback says that a post-training one does not: the model was fitted through the rounding rather than rounded after the fact.

A fit that raises stops before after_fit, so the parametrizations are still in place and torch.save(learn.model) refuses to serialize the model. Either finish the job or undo it:

cb.bake()    # keep the rounding: an ordinary model at the width it was training through
cb.strip()   # drop it: the floating-point weights, as the interrupted fit left them

What it costs

Rounding on every forward is not free, and nothing here claims otherwise. One architecture, one machine, one shape:

ResNet-18, 10 classes, batch 32 of 3x64x64, CPU step time peak resident memory
plain fit 109 ms 1.15 GiB
FakeQuantizeCallback(8, 8) 128 ms 1.27 GiB

Measured on one machine: an Intel i9-14900KS with torch.set_num_threads(4), torch 2.9.1, 12 training steps per run, six runs per arm in fresh processes, alternating between the arms; each figure is the median of the per-run medians (step time) and of the per-run high-water marks (ru_maxrss). That is 1.17x the step time and +120 MiB here. Those are the numbers this fixture produced, not a target: both move with the architecture, the batch size and the machine, so measure your own if the budget matters.

The master doubles nothing that was not already there — it is the parameter, and the rounded weight is the extra tensor — while the activation hooks add one rounded copy of every feature map they round.

With Sparsifier

The two compose in one fit, in either order: the mask is applied to the master, and a weight that is exactly 0 rounds to exactly 0 on a symmetric grid and on an affine one alike, so sparsity survives the rounding.

learn.fit(5, cbs=[SparsifyCallback(0.5, 'weight', 'local', large_final, one_cycle),
                  FakeQuantizeCallback(weight_bits=8)])

Sparsifier scores, masks, snapshots and rewinds the master, so print_sparsity() reports the mask rather than the rounding, and a winning ticket is saved as floating-point weights. Structured pruning is the exception: [Pruner](https://FasterAI-Labs.github.io/fasterai/prune/pruner.html#pruner) refuses a parametrized model, because torch-pruning rewrites the modules it traces. Bake or strip before pruning.

See Also

Tests live in nbs/tests/test_fake_quantize_callback.ipynb.