dew.training.quantization

Quantized training through Qwix, applied to the model before it trains.

Qwix (google/qwix, Apache 2.0) expresses quantization as rules over module paths and applies them without editing the model. One call wraps the module, and the matmuls in the wrapped methods’ extent run quantized.

Dew’s version of that call is apply_quantization. A caller builds its model from the registry as always, then wraps it before the objective ever sees it. A run that names --trainer.quantization instead hands RunConfig.train the objective, and quantize wraps the model it holds before anything initialises it.

What trains is fake-quantized. The parameter tree keeps fp32 master weights with the same structure, so the checkpoint layout, the sharding derivation, the Muon parameter split and Hugging Face loading are unchanged. The quantization lives in the forward and backward matmuls, with a straight-through estimator on the backward pass.

The vocabulary head stays fp32 with the rest of Dew’s fp32 zones. Its einsum lives in the objective’s chunked cross entropy, outside any model method Qwix wraps.

The value mirrors MaxText’s knob set (configs/base.yml:128-167) where Qwix has an equivalent. dtype is its quantization for the dynamic-range forms and patterns its quant_cfg_path, written inline as the regexes Qwix matches; the backward fields are Qwix’s finer-grained version of the same idea.

Three of its knobs have no equivalent and are refused with the reason. Static activation scaling (fp8_full) needs a calibration pass Dew has no seam for, nanoo_fp8 is AMD-only kernels, and KV-cache quantization has no reader here since the cache holds the compute dtype.

Qwix is not a dependency. The import sits inside apply_quantization, and without the package the call raises naming it, the way the tokamax branch of dew.nn.moe behaves.

Name	Summary
QuantizedDtype	The gemm dtypes a run trains with: int8 on any backend, fp8 where the backend lowers it (measured in docs/performance.md).
Rounding	How a quantized gradient rounds: Qwix’s two stochastic modes.
CALIBRATIONS	The weight calibration methods Qwix parses, before an optional ,args suffix (qwix/_src/qconfig.py, QuantizationRule).
METHODS
Quantization	Says how a run quantizes its trunk matmuls, for Qwix’s provider.
apply_quantization	Wrap model so its trunk matmuls train in spec’s dtype.
ModelObjective	Trains one module, which is the shape quantize can wrap.
quantize	Quantize the trunk matmuls of the module objective trains.

QuantizedDtype

attribute source

QuantizedDtype = Literal['int8', 'fp8']

The gemm dtypes a run trains with: int8 on any backend, fp8 where the backend lowers it (measured in docs/performance.md).

Rounding

attribute source

Rounding = Literal['uniform', 'low_bit_uniform']

How a quantized gradient rounds: Qwix’s two stochastic modes.

CALIBRATIONS

attribute source

CALIBRATIONS = ('absmax', 'minmax', 'rms', 'fixed')

The weight calibration methods Qwix parses, before an optional ,args suffix (qwix/_src/qconfig.py, QuantizationRule).

METHODS

attribute source

METHODS = ('__call__', 'hidden_states', 'mtp_hidden_states')

Quantization

dataclass source

class Quantization(
    dtype: QuantizedDtype = 'int8',
    patterns: tuple[str, ...] = ('.*',),
    calibration: str = 'absmax',
    tile_size: int | None = None,
    bwd_qtype: QuantizedDtype | None = None,
    bwd_stochastic_rounding: Rounding | None = None,
)

Says how a run quantizes its trunk matmuls, for Qwix’s provider.

dtype: QuantizedDtype

The dtype weights and activations quantize to, in the forward pass.

patterns: tuple[str, ...]

Module-path regexes the rules apply to, in Qwix precedence order: the first rule whose regex full-matches a module’s /-joined scope path wins. '.*mlp.*' quantizes the feed-forward blocks and leaves attention in fp32; the default quantizes every matmul of the wrapped methods.

calibration: str

How weights calibrate, as Qwix parses it: a method with an optional ,args suffix, for example absmax,0.8.

tile_size: int | None

Sub-channel tiling of the contraction axis; unset keeps per-channel scales, the coarser and cheaper form.

bwd_qtype: QuantizedDtype | None

The dtype gradients quantize to in the backward pass; unset keeps them in the compute dtype.

bwd_stochastic_rounding: Rounding | None

Stochastic rounding on the quantized gradients. A run that sets this passes a stochastic_rounding RNG stream at apply time, which Qwix draws (qwix/_src/providers/qt.py:361); unset rounds deterministically.

apply_quantization

function source

def apply_quantization(model: nn.Module, spec: Quantization) -> nn.Module

Wrap model so its trunk matmuls train in spec’s dtype.

The returned module is a copy of the same class with the entry methods it defines of METHODS wrapped, so everything the registry, the objective and the checkpoint code read off the model still answers. Construction already refused what the value cannot ask for; without the package the call raises naming it.

ModelObjective

class source

class ModelObjective(Protocol)

Trains one module, which is the shape quantize can wrap.

The module is the objective’s model, and every trace it runs reads it there.

quantize

function source

def quantize(objective: object, spec: Quantization) -> None

Quantize the trunk matmuls of the module objective trains.

apply_quantization wraps a module before an objective is built, which is what a recipe that builds its own model does. A run that names --trainer.quantization has handed RunConfig.train the objective already, so the wrap lands on the objective’s own model instead, before anything has initialised or traced it; the wrapped module is a copy of the same class, so what the objective read off the model at construction still holds.

An objective that trains something other than one module has nothing to wrap and is refused by name.