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(base, *, bits=4, clip_percentile=None, enabled=True)[source]

Bases: Module

Wrap an nn.Linear so its weight is straight-through fake-quantized on every forward call: the module computes F.linear(x, fake_quantize(weight), bias) instead of F.linear(x, weight, bias). Bias stays fp32, matching PTQ’s scheme (mixle.task.quantize.quantize_mlp()) where only weights are quantized.

Drop-in composition: QATWrapper(linear) has the same forward(x) -> Tensor contract as the Linear it wraps, so it slots into any module tree (see apply_qat()) without the surrounding model or training loop changing at all.

Parameters:

  • base (Any)

  • bits (int)

  • clip_percentile (float | None)

  • enabled (bool)

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Any)

Return type:

Any

extra_repr()[source]

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

apply_qat(model, *, bits=4, clip_percentile=None)[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:

  • model (Any)

  • bits (int)

  • clip_percentile (float | None)

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