mixle.models.qat module

Quantization-aware training (QAT): straight-through fake-quant, composed by wrapping.

Post-training quantization (mixle.task.quantize.quantize_mlp()) quantizes an already-trained float model’s weights to int4/int8 after the fact – the model never saw its own quantization error during training, so gradient descent had no chance to route around it. QAT simulates that error DURING training instead: every forward pass runs the weight through a real quantize-then-dequantize round trip (the exact int4 math already in mixle.task.quantize – this module does not reimplement it), so the optimizer sees a loss landscape shaped by the precision it will actually run at. The backward pass uses the straight-through estimator (Bengio et al.): the quantize/dequantize round trip is a step function with zero gradient almost everywhere, so STE simply copies the incoming gradient through as if the round trip were the identity – the standard trick that makes QAT trainable at all.

Composition follows this codebase’s established pattern (see mixle.experimental.program.lora, the peft/LoRA-style wrapper): QAT is a module-level wrapper around nn.Linear, not a change to the training loop. GradLeaf’s M-step only ever calls module.log_density(x) and back-propagates through whatever module is – it has no idea some of that module’s Linear layers fake-quantize their weights on every forward call, so QAT drops into GradLeaf.fit/estimate (and, when it exists, the F1 distributed trainer and J4 distillation students that also route through the same log_density contract) with no changes to ``grad_leaf.py``.

model = build_causal_lm(vocab=…, d_model=…, n_layer=…, n_head=…, block=…) apply_qat(model) # every nn.Linear now fake-quantizes its weight to int4 leaf = GradLeaf(SomeWrapperWithLogDensity(model)) fitted = leaf.estimator().estimate(None, suff_stat) # trains QAT-aware, unmodified M-step

F1 (a real distributed trainer skeleton) and J4 (distillation students) are separate, not-yet-built roadmap items; this module does not depend on either. What is real and tested here: the STE fake-quant op itself, and that wrapping a real transformer’s Linear layers with it and training end-to-end beats post-hoc PTQ at matched int4 size (see mixle/tests/qat_test.py).

fake_quantize(x, *, bits=4, clip_percentile=None)[source]

Straight-through fake-quantize x to bits (int4 or int8, per mixle.task.quantize._QMAX): forward returns the real quantize->dequantize round trip, backward passes the gradient through unchanged (STE).

Parameters:
Return type:

Any

fake_quantize_int4(x, *, clip_percentile=None)[source]

fake_quantize(x, bits=4) – the int4 case this roadmap item targets.

Parameters:

  • x (Any)

  • clip_percentile (float | None)

Return type:

Any

class QATWrapper(*args, **kwargs)[source]

Bases: object

Torch-absent stand-in: keeps QATWrapper always importable (mixle.models imports it unconditionally, matching every other torch-optional model in this package – see GaussianProcessRegressor), raising only when actually constructed, since it must subclass nn.Module to be real and nn does not exist here.

Parameters:

  • args (Any)

  • kwargs (Any)

apply_qat(model, *, bits=4, clip_percentile=None, reconfigure=False)[source]

Replace every nn.Linear under model in place with a QATWrapper, so the whole model trains quantization-aware (straight-through int4 fake-quant on every Linear weight, every forward call). Mirrors mixle.experimental.program.lora’s wrapping pattern: walk named_children, swap Linear leaves, recurse into everything else. Weight-tied layers (e.g. CausalLM.head sharing tok.weight) wrap cleanly – only the wrapped module’s forward changes, the underlying nn.Parameter (and anything else pointing at it) is untouched. Returns model (mutated in place) for chaining.

Parameters:
Return type:

Any

set_fake_quant_enabled(model, enabled)[source]

Toggle every QATWrapper under model on/off in place. enabled=False runs the model at its real fp32 weights (e.g. to check that QAT training didn’t wreck full-precision quality); enabled=True (the default after apply_qat()) restores the fake-quant forward. Returns model for chaining.

Parameters:
Return type:

Any