Source-linked AI summary
QATFactory: A Versatile, Deployment-Aligned Framework for Quantization-aware Training and Distillation of LLMs
Weili Xu, Jisen Li, Yuqing Jian, Chenxi Li, Zhizhou Sha, Yifan Yu, Qingyang Wu, Chenfeng Xu, Zhongzhu Zhou, Tianyi Zhang, Ben Athiwaratkun
TL;DR
Aggressive low-precision quantization can degrade LLM quality, while existing training pipelines may not match production inference numerics. QATFactory simulates deployment quantization during training using BF16 matrix multiplications and supports distillation and reinforcement learning. Across models and formats, QAD improves deployed-model quality over strong PTQ baselines, and the experiments find that the best training quantization choices vary by format.
Problem
Aggressive low-precision deployment can perturb model computations, while existing QAT pipelines may mismatch production-engine numerics and support for newer formats remains underexplored.
Source: §1 · pp. 2–3Method
Results
Across models and formats, QAD consistently outperforms strong PTQ baselines; NVFP4 QARL improves mean reasoning accuracy by 2.7 percentage points over BF16 RL followed by PTQ.
Source: §1 · p. 4, §5 · p. 22Takeaways & Limitations
The best training strategy depends on format: MXFP4 benefits from weight-and-activation quantization, whereas NVFP4 generally performs better with weight-only quantization.
Source: §1 · p. 3
Abstract
from arXiv · showhide
Large language model (LLM) inference is increasingly moving toward lower precision to realize the throughput of hardware accelerators, but aggressive post-training quantization (PTQ) can degrade model quality. We present QATFactory, an open-source framework for deployment-aligned quantization-aware distillation (QAD) and reinforcement learning (QARL). QATFactory simulates deployment-time quantization while performing matrix multiplications in BF16, allowing models to adapt to quantization noise without requiring training hardware that natively supports the target format; for example, it supports NVFP4 training on H100 GPUs, which lack FP4 Tensor Cores. The framework supports NVFP4, MXFP4, and $\text{llama}.\text{cpp}$'s Q4_K format; dense and mixture-of-experts models; and both full-parameter and LoRA-based training. It exports checkpoints directly to vLLM and $\text{llama}.\text{cpp}$ without an additional lossy conversion step or added inference overhead. With QATFactory, we conduct extensive experiments on models ranging from 8B to 230B parameters and evaluate exported checkpoints in production inference engines. Across models and formats, QAD consistently improves deployed-model quality over strong PTQ baselines. On Qwen3.5-9B, QAD achieves average benchmark accuracies of 68.9% under NVFP4 and 66.0% under MXFP4, outperforming the best PTQ results of 65.4% and 56.4%, respectively. Through our experiments, we found that although both FP4 formats quantize weights and activations at deployment, the best training strategy is format-dependent: NVFP4 generally performs better when only weights are quantized during training, whereas MXFP4 benefits from quantizing both weights and activations. At a fixed training token budget, training on fewer 32K sequences improves average accuracy by 1.9 points over training on more 4K sequences.
1. Introduction
Aggressive low-precision inference promises greater hardware throughput but can make model quality more vulnerable to quantization error. QATFactory trains with deployment-aligned quantization and reports quality gains over PTQ alongside format- and workload-dependent training findings.
1. Introduction
- The paper frames the central challenge as training models that withstand deployment numerics, since low-precision errors affect both operands in W4A4 and can compound in reasoning models.Existing QAT pipelines may mismatch production-engine semantics, and support for NVFP4, MXFP4, and Q4_K remains underexplored.Source: §1 · pp. 2–3
- QATFactory emulates each deployment format’s numerical behavior while retaining high-precision linear algebra, enabling NVFP4 QAT on H100 GPUs without native NVFP4 Tensor Cores.It supports direct checkpoint consumption by vLLM, SGLang, and llama.cpp without custom operators or runtime overhead.Source: §1 · pp. 3–4
- MXFP4 benefits from quantizing weights and activations during training, whereas NVFP4 generally produces higher-quality checkpoints with weight-only quantization.The finding exposes a format-dependent tradeoff between deployment fidelity and optimization stability.Source: §1 · p. 3
- NVFP4 QARL accelerates end-to-end RL training by 1.23× and improves mean reasoning accuracy by 2.7 percentage points over BF16 RL followed by PTQ.Source: §1 · p. 4
- At a fixed token budget, fewer long sequences generalize better overall than more short sequences, while QAD improvements transfer across domains, especially from code to mathematical reasoning.LoRA-based QAD reduces training memory but underperforms full-parameter QAD in these experiments; increasing adapter rank alone does not close the quality gap.Source: §1 · p. 4
- Across model families, architectural classes, and scales, QAD-trained checkpoints consistently outperform strong PTQ baselines in production inference engines.The framework supports NVFP4, MXFP4, and Q4_K, with dense and MoE models and both full-parameter and LoRA-based training.Source: §1 · p. 4
2. Related Work
Related work spans post-training quantization, quantization-aware training and distillation, efficient low-precision adaptation, FP4 training, and low-precision reinforcement learning. The cited studies motivate adaptation to deployment numerics and more efficient training workflows.
Post-training quantization for LLMs
- PTQ methods reduce LLM memory and inference costs without retraining, using approaches such as second-order weight reconstruction, redistribution of quantization difficulty, and activation-informed channel preservation.The section notes that aggressive weight-and-activation quantization remains challenging.Source: §2 · p. 4
Quantization-aware training and distillation
- Classical QAT uses fake quantization and straight-through estimators, while later work explores clipping and step-size learning, distillation, and self-distillation for low-bit language models.Prior LLM work also studies joint quantization of weights, activations, and KV caches.Source: §2 · pp. 4–5
Efficient and scalable low-precision training
- EfficientQAT, QA-LoRA, LR-QAT, and L4Q pursue lower-cost QAT or parameter-efficient adaptation, while SketchTune optimizes shared lookup-table values.Other studies examine how QAT varies with scale, data, and quantization configuration.Source: §2 · p. 5
FP4 and low-precision reinforcement learning
- Prior work studies native FP4 training and low-precision RL, including rollout-aligned methods that address numerical mismatch between rollout generation and policy optimization.Source: §2 · p. 5
3. QATFactory
QATFactory adapts pretrained LLMs to deployment-native formats through QAD or QARL. Both workflows share quantization layers, format implementations, and an export pipeline.
3. QATFactory
- QATFactory simulates quantization and dequantization during training while performing matrix multiplications in BF16, then exports packed checkpoints for target inference engines without extra PTQ.Supported formats include NVFP4, MXFP4, and Q4_K, without requiring specialized hardware for those formats.Source: §3 · p. 6
- QAD trains a quantized student to match its frozen high-precision counterpart, while QARL uses reward or verifier feedback to train the quantized model.Source: §3 · p. 6
- QARL uses low-precision inference engines for rollout generation, and both workflows use the same quantization layers, format implementations, and export pipeline.Source: §3 · p. 6
3.1. Framework overview and design goals
The framework design addresses deployment-format mismatch and the hardware and memory demands of low-bit training. Its shared quantization-aware layers and training options target fidelity, extensibility, accessibility, and scale.
3.1. Framework overview and design goals
- Existing QAT frameworks often target generic formats unsupported by modern inference engines, while native low-bit training can require specialized hardware.Source: §3.1 · p. 6
- Full-parameter distillation has high memory overhead because teacher weights, student latent weights, and optimizer states must coexist in GPU memory.Source: §3.1 · p. 7
- QATFactory’s deployment-aligned layers apply lossy quantize–dequantize transformations in the forward pass and use a straight-through estimator to update continuous latent parameters backward.Source: §3.1 · p. 7
- The design aims for deployment fidelity and direct export, extensibility across formats and engines, and accessibility through BF16 matrix multiplications, FSDP, and LoRA.Source: §3.1 · p. 7
3.2. Deployment-aligned simulated quantization
QATFactory simulates target-format quantization in high-precision training so models adapt to deployment numerics while matrix multiplications remain in BF16. Its unified quantizer interface supports several formats and engine-native layouts, with format-specific scaling and grouping rules.
3.2.1. Quantization-aware layer
- Quantization-aware layers simulate deployment-time quantization of weights and activations while performing matrix multiplications in BF16.The layers keep high-precision tensors during training and reconstruct values restricted to the target format’s representable grid.Source: §3.2.1 · p. 7, §3.2 · p. 7
- Each forward pass recomputes quantized representations from changing latent weights, using on-the-fly scales and GPU kernels before multiplication.The fake-quantized tensors are transient and discarded after the matrix multiplication.Source: §3.2.1 · p. 8
- A straight-through estimator treats quantize–dequantize operations as identity mappings during backpropagation, passing gradients to high-precision latent weights.The continuous latent weights and optimizer states are updated in high precision.Source: §3.2.1 · p. 8
3.2.2. Quantizer interface
- The unified quantizer interface separates format-specific quantization and engine-native weight packing from training logic, making new formats straightforward to support.During training it reconstructs BF16 operands; during export it packs low-precision weights and metadata in the target engine’s layout.Source: §3.2.2 · p. 8
3.2.3. Supported quantization formats
- NVFP4 uses blocks of 16 weights with E4M3 block scales and a global FP32 tensor scale; activation scales are tracked from running layer-input maxima.The tracked activation scales are exported with the checkpoint.Source: §3.2.3 · p. 9
- MXFP4 groups E2M1 codes into 32-value blocks with exact power-of-two scales and no tensor-level scale.Its coarser scales and larger groups can make model quality more difficult to preserve than NVFP4.Source: §3.2.3 · p. 9
- Q4_K is an asymmetric weight-only format with 256-weight super-blocks divided into eight 32-value sub-blocks using local scales and offsets.It double-quantizes the sub-block scales and offsets to reduce metadata overhead.Source: §3.2.3 · p. 9
- For NVFP4, QATFactory ties weight and activation scales across fused projections; this scale sharing is unnecessary for MXFP4 and Q4_K.The shared NVFP4 weight scale is computed from concatenated projection weights, with one activation scale for their common input.Source: §3.2.3 · p. 10
3.3. Quantization-aware distillation
QATFactory uses a frozen high-precision teacher to supervise a quantized student, minimizing token-averaged forward KL divergence. It supports full-parameter and LoRA adaptation for dense and MoE models, with different memory and adaptation trade-offs.
3.3. Quantization-aware distillation
- Self-distillation trains a low-precision student initialized from the pretrained weights to match the frozen teacher’s token-level output distributions using forward KL divergence.The loss is averaged across sequence tokens over the shared vocabulary.Source: §3.3 · p. 10
- The teacher remains frozen in its original precision, while the student quantizes linear projections and optimizes their continuous latent weights.Embeddings, normalization layers, routers, and the language-modeling head remain frozen in their original precision.Source: §3.3 · p. 10
- Full-parameter QAD typically provides the highest quality but requires substantially more memory than LoRA-QAD, which trains adapters while sharing frozen base weights.LoRA parameters are merged with the base weights before quantization during the student forward pass.Source: §3.3 · p. 10
- For MoE models, QATFactory supports full expert-weight training and unshared or hybrid shared LoRA layouts, processing and releasing one merged expert weight at a time.The hybrid layout shares A for gate and up projections and B for down projections across experts, respectively.Source: §3.3 · pp. 10–11
3.4. Quantization-aware reinforcement learning
QARL extends deployment-aligned training to reinforcement learning with verifiable rewards, using separate rollout and policy-training servers. Native low-bit rollout inference reduces generation time and improves end-to-end throughput in the evaluated workload.
3.4. Quantization-aware reinforcement learning
- QARL extends deployment-aligned training to RLVR and supports algorithms including GRPO and DAPO.Source: §3.4 · p. 11
- A rollout server generates responses with native low-bit inference, while a training server simulates quantization and computes policy updates in BF16.Verifiers score rollout responses and send trajectories and rewards to the training server.Source: §3.4 · p. 11
- Because rollout generation dominates step latency in the evaluated RLVR workload, native low-precision inference reduces generation time and improves end-to-end training throughput.Updated weights are quantized and packed in the inference engine’s native layout before synchronization to the rollout server.Source: §3.4 · p. 11
3.5. Engine-native checkpoint export
QAD and QARL share an export pipeline that produces low-precision checkpoints in the target inference engine’s native format. The pipeline avoids additional post-training quantization or calibration before direct loading by supported engines.
3.5. Engine-native checkpoint export
- QAD and QARL share an export pipeline that produces low-precision checkpoints in the target inference engine’s native format.Source: §3.5 · p. 11
- LoRA-QAD merges trained adapters into frozen base weights before quantization, while full-parameter QAD and QARL quantize final high-precision latent weights directly.The LoRA-QAD checkpoint therefore comes from the same merged weights used during training.Source: §3.5 · p. 11
- The exporter packs low-bit values and metadata in the target engine’s layout using the training quantization configuration, then checkpoints load directly without additional quantization or calibration.For formats requiring tracked state, such as NVFP4 input-activation scales, it saves values collected during training.Source: §3.5 · p. 11
4. Experiments and Findings
The experiments evaluate exported QAD checkpoints across models, quantization formats, and inference engines, alongside training-strategy variations in activation precision, adapter rank, sequence length, and reinforcement learning. Across these settings, QAD improves on PTQ, while the best training choices depend on the format and configuration.
4.1. Experimental Setup
- Experiments cover dense and MoE models from 8B to 230B parameters, three quantization formats, and exported-checkpoint evaluations against PTQ baselines.QAD runs use teacher-generated data and forward-KL distillation; evaluations include downstream benchmarks and KL divergence to the unquantized reference.Source: §4.1 · pp. 12–13
4.2. Generality Across Formats, Scales, and Architectures
- QAD achieves higher average benchmark accuracy and lower mean KL divergence than PTQ baselines across formats and model families, including large MoEs and Q4_K models.For Qwen3.5 9B, NVFP4 QAD reaches 83.0% on LiveCodeBench versus 63.7% for RTN and 66.0% for GPTQ; MXFP4 QAD reaches 66.0% average accuracy versus 56.4% for the best PTQ baseline.Source: Table 2 · p. 13, §4.2 · p. 14, Table 3 · p. 15
4.3. Choosing Training-Time Activation Precision for W4A4 Deployment
- For W4A4 deployment, NVFP4 matches or favors W4A16 training, whereas MXFP4 benefits from W4A4 training.On Qwen3.5 9B, MXFP4 W4A4 training raises average accuracy from 57.8% to 66.0% and lowers average KL divergence from 0.1510 to 0.1438.Source: Table 2 · p. 13, §4.3 · p. 15
4.4. Parameter-Efficient QAT via LoRA: Memory Footprint and Scaling Limits
- Rank-16 LoRA reduces Qwen3.5 9B training memory by 2.91×, but LoRA remains behind full-parameter QAD and higher rank does not reliably improve quality.For Muse-Glimmer 30B, rank 32 exceeds rank 64 in average accuracy (78.4% vs. 76.5%) and average KL divergence (0.0260 vs. 0.0262).Source: Table 4 · p. 16, §4.4 · pp. 17–18, Table 6 · p. 18
4.5. Sequence-Length and Cross-Domain Generalization
- At the same training-token budget, 32K sequences raise pooled average accuracy from 50.1% to 51.9% (1.9 percentage points), improving three of four comparisons and tying one.AIME25 and LCB-v6 improve in every comparison; BCB declines in both W4A4 comparisons.Source: §4.5 · p. 20
4.6. Quantization-Aware Reinforcement Learning (QARL)
- NVFP4 QARL improves mean reasoning accuracy by 2.7 percentage points over BF16 RL followed by PTQ and increases training throughput by 1.23×.Mean accuracy is 50.8% with QARL versus 48.1% after PTQ; total step time falls from 41.83 s to 33.96 s.Source: §4.6 · pp. 20–21, Table 8 · p. 21, Table 9 · p. 21
4.7. Guidelines for Quantization-Aware Training
- The experiments recommend choosing activation precision per format, using LoRA for memory savings without assuming higher ranks help, and favoring longer sequences at fixed token budgets.When native low-precision rollouts are available, quantization-aware RL can improve both training speed and quantized-model quality.Source: §4.7 · p. 21
5. Conclusion and Future Directions
QATFactory is an open, deployment-aligned framework that reproduces target quantization numerics while computing in BF16, and QAD outperforms strong PTQ baselines across the reported settings. The conclusion highlights format- and workload-specific training choices and outlines extensions to additional formats, engines, architectures, and methods.
5. Conclusion and Future Directions
- Across reported settings, QAD outperforms strong post-training quantization baselines, while effective training choices vary by format and workload.The experiments find that activation quantization helps MXFP4 but not NVFP4, LoRA reduces memory use but does not match full-parameter QAD, longer sequences improve generalization at a fixed token budget, and NVFP4 QARL improves quality and throughput over BF16 reinforcement learning followed by quantization.Source: §5 · p. 22
- QATFactory reproduces NVFP4, MXFP4, and Q4_K quantization behavior in BF16 and supports dense and MoE models from 8B to 230B parameters.It exports checkpoints directly to vLLM, SGLang, and llama.cpp.Source: §5 · p. 22
- Future work will extend QATFactory to additional quantization formats, inference engines, model architectures, and QAT methods, including lookup-table formats planned for NVIDIA Rubin GPUs.The stated broader goal is a common, deployment-aligned platform for developing and evaluating quantization-aware training techniques.Source: §5 · p. 22