[4/N] Add QAT (quantization-aware training) model converter

[mori360]

Summary

Add a QATConverter that uses torchao's QATConfig to insert fake quantization into nn.Linear modules during training. This enables quantization-aware training as a first-class model converter in torchtitan. QAT composes with LoRA by converter ordering: QAT first replaces nn.Linear → FakeQuantizedLinear, then LoRA inherits from FakeQuantizedLinear, giving fake-quantized base weights with full-precision adapters.

• torchtitan/components/quantization/qat.py — New QATConverter with support for 7 torchao quantization schemes: int4_weight_only, intx_weight_only, int8_dynamic_act_intx_weight, float8_dynamic_act_float8_weight, float8_dynamic_act_int4_weight, nvfp4, and mx. A shared _build_base_config() helper maps scheme names to torchao PTQ base configs.
• torchtitan/components/lora.py — Add adapter_qat_scheme and adapter_qat_group_size config fields to LoRAConverter, enabling fake quantization of LoRA adapter weights (lora_a/lora_b) independently of base model quantization. Validates that rank % group_size == 0 for schemes that use per-group granularity.
• torchtitan/protocols/model_converter.py — Extend converter ordering validation to enforce that QATConverter (like quantization converters) must come before LoRAConverter, since QAT replaces nn.Linear with FakeQuantizedLinear and LoRA must wrap the final linear class.
• torchtitan/models/llama3/config_registry.py — Add llama3_debugmodel_qat and llama3_debugmodel_qat_lora debug configs for testing QAT alone and QAT+LoRA composition.
• tests/unit_tests/test_model_converter.py — Add tests for dtype preservation, all 7 QAT schemes, unknown scheme error, and group_size warning for unsupported schemes.

Test Plan

• test_qat_preserves_weight_dtype: verifies all parameter dtypes remain unchanged after QAT conversion

Stack from ghstack (oldest at bottom):

• [5/N] Support lora moe #2569
• -> [4/N] Add QAT (quantization-aware training) model converter #2488
• [2/N] Support lora checkpoint on partial save and multi-source load #2485
• [1/N] Add lora adapters for finetune #2484

[tianyu-l]

similar to the comment in [1/N]: order matters

[andrewor14]

Looks great!

[andrewor14]

This was only for torch 2.5 and before (which didn't have torch.int4), I think now you can just use torch.int4 directly here

[andrewor14]

Should we support activation quantization too? There are pre-set QAT configs in torchao that do this, e.g. Int8DynamicActivationIntxWeightConfig (used primarily for edge) can be passed to QATConfig as a base config (docs here, a bit outdated).

Another example that I think will be important is NVFP4DynamicActivationNVFP4WeightConfig. I don't think a single dtype field will be sufficient in capturing that

[andrewor14]

Just for future reference but there's also NVFP4FakeQuantizedLinear that's not a subclass of this. We can add support for this later but just wanted to bring it up early so we're not baked into this specific linear class

[andrewor14]

By the way not all quantization configs have group_size, you can see the list of supported QAT configs here

[andrewor14]

Does the LoRAConverter already handle the case when the inner linear is FakeQuantizedLinear? I think additionally we will need to make the lora adapters also fake quantized (according to the same configs, e.g. int4). Does this part exist yet?

[mori360]

Currently we follow: Linear -> QAT(Linear) -> LoRA(QAT) that only adds adapters to the QATLinear, so the currently design has lora adapters in dtype same as the origin linear.