FakeQuantizeCallback trains a model throughFakeQuantizer’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 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 sameParameter 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 and Criteria go through _master(m), 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
a scheduled width steps down onto the rung its progress selects
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.
def FakeQuantizeCallback( weight_bits:int|None=8, # Weight width; None leaves them in floating point act_bits:int|None=None, # Activation width; None leaves them in floating point qscheme:str='per_channel', # Weight scale axis: 'per_tensor', 'per_channel' or 'per_group' symmetric:bool=True, # Zero-point 0, or an affine offset group_size:int|None=None, # Weights sharing one scale (qscheme='per_group') observer:str='static', # Activation scales: 'static' (observed, then frozen) or 'dynamic' layer_type:type|tuple=(<class'torch.nn.modules.conv.Conv2d'>, <class'torch.nn.modules.linear.Linear'>), # Module types that carry the rounding layer_bits:dict|None=None, # {layer_name: width} overriding `weight_bits` layer_act_bits:dict|None=None, # {layer_name: width} overriding `act_bits` freeze_act_pct:float=0.9, # Training fraction after which the activation scales stop moving act_averaging:float|None=0.01, # How far the observed range moves towards each batch; None takes its absolute min and max weight_schedule:Schedule |None=None, # Schedule stepping `weight_widths`; None rounds at `weight_bits` from the first batch weight_widths:Sequence[int] |None=None, # Weight widths to step down through, coarsest first, ending at `weight_bits` act_schedule:Schedule |None=None, # Schedule stepping `act_widths`; None rounds at `act_bits` from the first batch act_widths:Sequence[int] |None=None, # Activation widths to step down through, coarsest first, ending at `act_bits` model:nn.Module |None=None, # Model to round; None = learn.model):
Quantization-aware training on FakeQuantizer’s arithmetic: the fit sees rounded weights, the optimizer trains floating-point ones.
before_fit parametrizes every weight this rounds, so the forward computes it from a master the gradient reaches; after_fit bakes the rounding in and leaves the model quantize_model() leaves. self.fake_quantizer is that quantizer, and its remove() gives back the master as the fit left it — the trained weights, not the ones training started from.
A fit that raises never reaches after_fit and leaves the parametrizations in place, on a model torch.save refuses: strip() takes them off and keeps the floating-point weights, bake() keeps the rounding.
A schedule can lower the width during the fit rather than rounding at the target from the first batch (weight_schedule / weight_widths, and their activation twins). Whether that is worth doing is not measured here.
Nothing here folds or freezes BatchNorm, so the weight ranges rounded are not the ones a deployed flow would see. The activation rounding sits on each module’s own output, before whatever activation function follows it, where an integer runtime fuses the two and rounds once.
The arguments are 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. The others 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.
weight_schedule / weight_widths and act_schedule / act_widths: the width ladder below. Without them the rounding is applied at the width you asked for from the first batch, the way prepare_qat_fx does.
model: the model to round, when it is not learn.model.
Drop the rounding and keep the floating-point master — the way back from a fit that raised
Usage
from fasterai.quantize.fake_quantize_callback import FakeQuantizeCallbackcb = FakeQuantizeCallback(weight_bits=4) # weight-only QAT at 4 bits, per channellearn.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_specfake_quant_spec(learn.model) # FakeQuantSpec(..., trained=True)learn.validate() # what the trained width costs on your metriccb.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 throughcb.strip() # drop it: the floating-point weights, as the interrupted fit left them
The width ladder
The fit can also start wide and step down to the width you asked for, one axis at a time. A ladder is two arguments: a Schedule, whose progress in [0, 1] selects the rung, and at least two widths to step through — coarsest first, ending at the width the callback rounds to.
from fasterai.core.schedule import lin, one_shot# weights start 16-bit and reach the 4 bits `weight_bits` nameslearn.fit(5, cbs=[FakeQuantizeCallback(weight_bits=4, weight_schedule=lin, weight_widths=(16, 8, 4))])# the two axes are independent: 8-bit weights throughout, activations dropped to 4 halfwaylearn.fit(5, cbs=[FakeQuantizeCallback(8, 4, act_schedule=one_shot, act_widths=(8, 4))])
The rungs share the schedule’s progress equally: with three rungs, progress below 1/3 selects the first, below 2/3 the second, and the rest the last. cb.current_widths is the pair the model is rounding at right now. Three things the mechanism guarantees:
The last training batch is always on the final rung, whatever the arithmetic says. pct_train is incremented after each batch, so it never reaches 1 — a 4-step fit tops out at 0.75 — and a ladder that waited for it would bake a width the spec does not name. bake() sets the final rung too, so a fit cut short still leaves the model at weight_bits / act_bits.
Changing a width re-observes nothing. On the activation side the observed range is width-independent — only the scale = range / qmax it defines moves with the width, and that is rewritten from the range the observer already holds. On the weight side there is nothing to redo at all: the master stays in floating point and every forward rounds a copy of it, so a rung change is simply what the next forward reads.
The schedule is only read. The rung, not Schedule.changed, decides whether anything is written, and after_step() is never called — so handing the same schedule object to a SparsifyCallback in the same fit leaves that callback’s own reading of it intact.
A layer whose width was named by layer_bits or layer_act_bits keeps that width for the whole fit — the ladder moves the layers you did not name.
Each axis refuses on its own, at construction:
Request
Refusal
weight_schedule=lin with no weight_widths
a schedule has no ladder to step through
weight_widths=(16, 8, 4) with no weight_schedule
a ladder nothing steps through
weight_bits=4, weight_widths=(16, 8)
the last rung must be the width this rounds to
weight_widths=(8, 8, 4) or (4, 8)
a ladder goes down, strictly
weight_widths=(4,)
a single rung steps nothing
weight_widths=(17, 8, 4)
a rung outside the widths this engine rounds to ([2, 16])
weight_widths=8
a ladder is a sequence of widths, not one
weight_bits=None with a weight_schedule
an axis left in floating point has no width to lower
FakeQuantSpec records the ladder next to the width the model ended at, so weight_widths=(16, 8, 4) with weight_bits=4 reads back as FakeQuantSpec(weight_bits=4, ..., weight_widths=(16, 8, 4)) — a model that arrived at 4 bits through a ladder is not read as one trained at 4 bits throughout. The field is the ladder the fit was asked to step down: a fit short enough to skip a rung, or one that raised before reaching one, still records the whole ladder.
Whether stepping the width down is worth doing is not measured here. The ladder is a mechanism this callback offers; which schedule, which rungs, and how any of it compares to training at the target width from the first batch are questions for a separate experiment, on your model and your data.
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.
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.
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 refuses a parametrized model, because torch-pruning rewrites the modules it traces. Bake or strip before pruning.