huggingface / huggingface/trl

Explore QAT support in TRL

Open
#6,567 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
19.3k
Forks
3k
Avg merge
1d 20h
Merged PRs (30d)
194

Description

TRL only does *post-training* quantization today: `get_quantization_config` returns a `BitsAndBytesConfig`, so the model is frozen-quantized then adapted with LoRA (QLoRA). No quantization-aware training anywhere.

Question: is QAT worth supporting; fake-quant during the forward pass so the model learns to absorb quantization error, then convert to a real quantized checkpoint?

Why now?: the big releases ship QAT, not PTQ. [Kimi K3](https://www.kimi.com/blog/kimi-k3) applies it from SFT onward (MXFP4 weights, MXFP8 activations) and K2-Thinking shipped INT4 W4A16. [Gemma 4](https://blog.google/innovation-and-ai/technology/developers-tools/quantization-aware-training-gemma-4/) released QAT checkpoints for every variant, including *unquantized* ones meant to be fine-tuned then re-quantize. And [gpt-oss](https://huggingface.co/docs/transformers/quantization/mxfp4) is trained natively in MXFP4, yet `examples/scripts/sft_gpt_oss.py` must pass `dequantize=True` to fine-tune it: we hand back a bf16 model that needs PTQ to redeploy.

To look into:
- Backend: `torchao.quantization.qat` — `quantize_(model, QATConfig(base_config, step="prepare"))`, train, then the same call with `step="convert"`. Axolotl and Unsloth both integrate it; torchao is not yet a TRL dependency.
- Where it lives: `TorchAoConfig` is load-time PTQ only, but `HfQuantizer` already exposes `is_qat_trainable`, honoured by `compressed_tensors`, `bitnet` and `fp_quant`; not torchao. Upstreaming the prepare/convert plumbing there and leaving TRL a config flag looks better than building it locally. Worth asking the transformers quantization folks first.
- Convert is a post-training step we don't have today (Axolotl ships a separate `quantize` command), and `fake_quant_after_n_steps` (delay fake-quant for N steps) needs a training-loop hook -> a callback?
- Which trainers: SFT first, where K3 starts and no rollout story is needed. For GRPO, fake-quant hits only the training forward while vLLM generates from unquantized weights; [slime/SGLang](https://www.lmsys.org/blog/2026-01-26-int4-qat/) quantize the rollout too and report train–infer logprob gaps overlapping BF16. A design question, not a blocker, and it lands on the train–infer mismatch we already care about.
- PEFT: QAT + LoRA vs full-parameter QAT.
- Export: does a converted model round-trip through `save_pretrained` / Hub / vLLM? torchao safetensors serialization exists (>= 0.15) but is documented for the `TorchAoConfig` path, not `quantize_`.

Outcome wanted: a yes/no on feasibility plus a sketch of the smallest useful version, not an implementation.

Related: https://github.com/huggingface/trl/issues/3399, https://github.com/huggingface/transformers/issues/10639

Contributor guide

Open the contributing guide

Research direction

Start by reading TorchAoConfig, HfQuantizer, and examples/scripts/sft_gpt_oss.py, then compare the related Transformers and TRL issues. Done means a feasibility decision and a sketch of the smallest useful QAT version, including prepare/convert flow, trainer scope, PEFT choice, delayed fake-quant handling, and checkpoint export.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.