QATFactory

A Versatile, Deployment-Aligned Framework for Quantization-aware Training and Distillation of LLMs

Weili Xu1,2,*, Jisen Li1,*, Yuqing Jian1,*, Chenxi Li1, Zhizhou Sha1,3, Yifan Yu1,2, Qingyang Wu1, Chenfeng Xu1,3, Zhongzhu Zhou1, Tianyi Zhang1, Ben Athiwaratkun1

1Together AI2University of Illinois Urbana-Champaign3The University of Texas at Austin

*Equal contribution

Paper arXiv Code Models & data Jump to findings
Two paths from a high-precision model. Top: post-training quantization calibrates and rounds the model, deploys it on vLLM, SGLang or llama.cpp, and loses accuracy. Bottom: QATFactory trains quantization-aware with self-distillation or RL, exports MXFP4, NVFP4 or Q4_K checkpoints, and deploys them on the same engines. Inset, Qwen3.5-9B NVFP4: LiveCodeBench 0.823 BF16, 0.637 RTN, 0.660 GPTQ, 0.830 QAD; BigCodeBench 0.317, 0.250, 0.233, 0.297.
Post-training quantization (PTQ) quantizes a trained model directly and can degrade accuracy. QATFactory trains with quantization in the loop, by self-distillation or reinforcement learning, then exports a deployment-ready checkpoint. On Qwen3.5-9B, its NVFP4 model beats round-to-nearest (RTN) and GPTQ on LiveCodeBench and BigCodeBench.

Overview

Recover the accuracy that PTQ loses

Inference keeps moving to lower precision to use the throughput of new accelerators, and aggressive PTQ can degrade model quality. QATFactory is an open-source framework for quantization-aware distillation (QAD) and quantization-aware reinforcement learning (QARL).

Training uses simulated quantization. Weights, and optionally activations, are rounded onto the target format's grid, while every matrix multiplication runs in BF16. A model can therefore train for a format its hardware cannot run: NVFP4 models train on H100 GPUs, which have no FP4 Tensor Cores.

The framework supports NVFP4, MXFP4 and llama.cpp's Q4_K. It trains dense and mixture-of-experts (MoE) models from 8B to 230B parameters, with full-parameter updates or low-rank adaptation (LoRA). Exported checkpoints load in vLLM and llama.cpp with no extra lossy conversion and no added inference cost.

Method

Framework overview

QATFactory supports two training workflows: quantization-aware distillation (QAD) and quantization-aware reinforcement learning (QARL). Both share the same quantization-aware layers, format implementations and export pipeline. In the forward pass, these layers apply a lossy quantize-dequantize transformation that reproduces inference-time numerics. In the backward pass, a straight-through estimator (STE) updates the continuous latent weights.

(a) Quantization-aware distillation: training data feeds a frozen BF16 teacher and a quantization-aware student; a distillation loss between them updates the student. (b) Quantization-aware reinforcement learning: a low-precision vLLM rollout engine samples rollouts for each prompt, a reward or verifier scores them, and the quantization-aware learner sends its weights back to the rollout engine. Right: inside a quantization-aware layer, latent weights and activations pass through quantize-dequantize into simulated-quantized tensors that feed a BF16 matmul; the straight-through estimator approximates the gradient with respect to the latent weights by the gradient with respect to the simulated-quantized weights.
QAD distills from the frozen high-precision model. QARL samples rollouts from the packed low-precision model in vLLM and syncs the learner's weights back after every step. Both use the quantization-aware layer on the right.
Numerics

Deployment-aligned simulated quantization

The inference engine acts as the numerical specification: each layer follows the target format's values, block sizes and scaling rules. Under NVFP4, projections that vLLM and SGLang fuse, such as QKV and gate/up, share one tensor scale in training as they do in serving. MXFP4 and Q4_K need no such tying, because neither uses tensor-level scales.

QAD

Quantization-aware distillation

The quantized student learns to match its frozen original model by minimizing forward KL divergence, on responses the original model generated for Open Perfect Blend prompts. Full-parameter QAD updates every quantized weight. LoRA-QAD shares one frozen base between teacher and student, trains only the adapters and quantizes the merged weight. For MoE models, it merges and quantizes one expert at a time to bound peak memory.

QARL

Quantization-aware reinforcement learning

A vLLM engine samples GRPO rollouts from the packed NVFP4 policy while the learner trains under simulated quantization. After each step, the learner packs its weights in the engine's layout and syncs them, so the engine needs no requantization.

Export

Engine-native checkpoints

Export packs weights and scales into each engine's native layout: vLLM compressed-tensors checkpoints for NVFP4 and MXFP4, and a llama.cpp GGUF file for Q4_K. NVFP4 activation scales are tracked during training and saved, so no calibration pass runs.

FormatNVFP4MXFP4Q4_K
Deployment precisionW4A4W4A4W4A16, weight-only
ElementE2M1 (FP4)E2M1 (FP4)Unsigned INT4
Block size163232, in a 256 super-block
ScalesBlock E4M3, tensor FP32Block E8M0Block 6-bit, super-block FP16
SymmetrySymmetricSymmetricAsymmetric
Bits per weight4.504.254.50
Training precisionW4A16 or W4A4W4A16 or W4A4W4A16
EnginesvLLM, SGLangvLLM, SGLangllama.cpp

Supported formats. Bits per weight include element codes and scale metadata.

Key findings

What we learned training in low precision

Every quantized checkpoint below is exported to its engine and evaluated there: vLLM for NVFP4 and MXFP4, llama.cpp for Q4_K. KL divergence to the unquantized model is measured on UltraChat, the Pile and WikiText.

01Does QAD beat post-training quantization across formats, scales and architectures?

Yes. In every model and format tested, exported QAD checkpoints reach higher average accuracy and lower mean KL divergence than the post-training baselines run for that setting.

Training stays stable on every model, including the 230B-scale MiniMax-M2.7. The gap is widest on coding. On Qwen3.5-9B under NVFP4, QAD scores 83.0 on LiveCodeBench, where the best baseline scores 66.0.

Training loss, training top-1 agreement, held-out KL and held-out top-1 agreement over QAD training for Qwen3.5-9B and DeepSeek-R1-Distill-Llama-8B under NVFP4 and MXFP4.
Dense models. Training loss and held-out KL fall while top-1 agreement with the BF16 model rises, for Qwen3.5-9B and DeepSeek-R1-Distill-Llama-8B under NVFP4 and MXFP4.
Training and held-out evaluation loss for Qwen3-8B and Qwen3.5-9B under Q4_K, and for Qwen3-30B-A3B and MiniMax-M2.7 under NVFP4.
Q4_K and MoE. Training and held-out loss for Qwen3-8B and Qwen3.5-9B under Q4_K, and for Qwen3-30B-A3B and MiniMax-M2.7 under NVFP4.
02W4A4 deployment: should training quantize activations too?

It depends on the format. For NVFP4, weight-only W4A16 training matches or beats W4A4 training. MXFP4 benefits substantially from W4A4 training.

On Qwen3.5-9B under MXFP4, W4A4 training raises average accuracy from 57.8 to 66.0. Under NVFP4 on the same model, the two averages tie and W4A16 training gives the lower KL.

Training precisionQwen3.5-9B, NVFP4Qwen3.5-9B, MXFP4DS-R1-Distill-Llama-8B, NVFP4DS-R1-Distill-Llama-8B, MXFP4
Avg ↑KL ↓Avg ↑KL ↓Avg ↑KL ↓Avg ↑KL ↓
W4A16, weights only68.90.068257.80.151047.70.111343.70.2305
W4A4, weights and activations68.90.074666.00.143846.60.114844.00.2090

Both rows are QAD checkpoints served W4A4; they differ only in whether training quantized activations.

03Can LoRA make QAD cheaper without losing quality?

It saves memory, not quality. Rank-16 LoRA cuts GPU memory 2.9× on Qwen3.5-9B at an 8K sequence length. On aggregate metrics every rank still trails full-parameter QAD, and a larger rank does not close the gap.

LoRA drops the separate teacher copy and keeps gradients and optimizer states only for the adapters. Quality does not rise with rank. Rank 4 gives the best LoRA average accuracy, and rank 256 the lowest.

GiBFull QADLoRA, r=16
Teacher weights, BF1616.680.00
Student weights, frozen BF163.7916.68
Student weights, trainable FP3225.770.16
Gradients, FP3225.770.16
Optimizer states, FP3251.550.32
Activations32.9735.37
Temporary10.004.50
GPU memory in use166.5357.19

Training memory for Qwen3.5-9B at an 8K sequence length.

CheckpointAvg ↑KL ↓
BF16 reference71.7—
Full QAD68.90.0682
LoRA, r=468.00.0722
LoRA, r=1667.80.0735
LoRA, r=6467.90.0714
LoRA, r=25667.60.0777

Qwen3.5-9B NVFP4, same eight benchmarks and three KL corpora as Finding 1.

Training loss, training top-1 agreement, held-out evaluation loss and held-out top-1 agreement for LoRA-QAD at ranks 4, 16, 64 and 256 and for full-parameter QAD on Qwen3.5-9B.
LoRA ranks. LoRA-QAD at ranks 4, 16, 64 and 256 recovers quality during training, yet stays behind full-parameter QAD in held-out loss and agreement. Red is full-parameter QAD. The blues are LoRA ranks, darker for larger ranks.
04At a fixed training token budget, should QAD train on fewer, longer samples or more, shorter samples?

Fewer, longer samples. In the paper's test, long samples are cut at 32K tokens and short ones at 4K. At the same training token budget, the long-sample runs score 1.9 points higher in average accuracy, pooled over four comparisons.

32K training wins three comparisons and ties the fourth. It improves AIME25 and LiveCodeBench-v6 every time, while BigCodeBench drops in both W4A4 comparisons. QAD also transfers across domains. Code-trained checkpoints match or beat math-trained ones on AIME25 in three of four settings.

Average of AIME25, BigCodeBench and LiveCodeBench-v6 for Qwen3.5-9B NVFP4 QAD. Each 4K and 32K pair sees the same number of training tokens. The axis starts at 44.

InferenceTraining dataStepsTokens (M)AIME25 ↑BCB ↑LCB-v6 ↑Avg ↑
BF16Teacher reference——54.6730.6786.3357.22
W4A16Math, 4K5,264336.747.3328.6779.0051.67
W4A16Math, 32K2,164336.851.3330.0081.6754.33
W4A16Code, 4K5,264307.946.6730.6779.6752.34
W4A16Code, 32K1,817307.952.0034.0082.6756.22
W4A4Math, 4K5,264336.742.0025.3372.0046.44
W4A4Math, 32K2,164336.844.0023.6774.6747.45
W4A4Code, 4K5,264307.945.3326.3377.6749.78
W4A4Code, 32K1,817307.946.6724.3378.3349.78

Qwen3.5-9B NVFP4 QAD. Each 4K and 32K pair sees the same number of training tokens. AIME25 averages five samples per problem.

Training loss and held-out fidelity over QAD training for the token-matched math and code runs at 4K and 32K sequence length on Qwen3.5-9B.
Token-matched runs. Training objective and held-out fidelity against supervised tokens seen. Blue is 4K and red is 32K context. Solid lines are math and dashed lines are code.
05Is quantization-aware RL better than RL followed by quantization?

Yes, on both speed and quality. NVFP4 QARL trains at 123.2% of BF16 RL's throughput and scores 2.7 points higher on five math benchmarks than BF16 RL followed by RTN.

Rollouts dominate each step, and native NVFP4 weights shorten them. That gain more than offsets the slower sync, forward and backward stages. QARL stays within 1.3 points of BF16 RL and slightly beats it on AIME24.

CheckpointTrainEvalMean ↑
Qwen3-8B-Base—BF1621.0
BF16 RLBF16BF1652.1
BF16 RL, then RTNBF16NVFP448.1
QARLNVFP4NVFP450.8

Mean accuracy over AIME24, AIME25, MATH500, OlympiadBench and GSM8K. Qwen3-8B-Base trains on DeepMath-103K for 1,024 steps, on four NVIDIA B200 GPUs, because native NVFP4 rollouts need Blackwell Tensor Cores.

Seconds per stepNVFP4 QARLBF16 GRPO
Weight sync0.990.73
vLLM rollout28.3137.46
Training forward1.520.98
Backward3.142.66
Total33.9641.83

Wall-clock time per optimizer step. Backward includes activation-checkpoint recomputation.

Training loss, rollout reward and completion length over 1,024 optimizer steps for BF16 GRPO and NVFP4 QARL on Qwen3-8B-Base.
RL training. Loss, rollout reward and completion length over 1,024 steps for BF16 GRPO and NVFP4 QARL on Qwen3-8B-Base.

Takeaways

Practical advice for low-precision training

Choose activation precision per format

Deployment precision alone should not set training precision. For NVFP4 W4A4 deployment, W4A16 training generally gives better aggregate accuracy and KL divergence. MXFP4 benefits from W4A4 training. Compare exported checkpoints rather than rely only on training-time metrics.

Evidence: Finding 02 →

Use LoRA for memory, not for quality

LoRA cuts training memory substantially. A larger adapter rank does not reliably improve accuracy or KL divergence, and it does not close the gap to full-parameter QAD.

Evidence: Finding 03 →

Prefer longer sequences at a fixed budget

Fewer 32K sequences generalize better overall than more 4K sequences with the same number of tokens. QAD updates also transfer across domains. In the paper's experiments, code-domain training improved math reasoning.

Evidence: Finding 04 →

Train RL in the deployment format

When native low-precision rollouts are available, combine quantized rollouts with quantization-aware updates. In the paper's NVFP4 experiment, this improved both speed and quality over BF16 RL followed by PTQ.

Evidence: Finding 05 →

Citation

BibTeX

@misc{xu2026qatfactory,
  title         = {QATFactory: A Versatile, Deployment-Aligned Framework for Quantization-aware Training and
                   Distillation of LLMs},
  author        = {Xu, Weili and Li, Jisen and Jian, Yuqing and Li, Chenxi and Sha, Zhizhou and Yu, Yifan and
                   Wu, Qingyang and Xu, Chenfeng and Zhou, Zhongzhu and Zhang, Tianyi and Athiwaratkun, Ben},
  year          = {2026},
  eprint        = {2609.39223},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2609.39223}
}