Add support for QAT + LoRA

[andrewor14]

Summary: Quantization-aware training (QAT) helps mitigate quantization degradation by simulating quantization numerics in high precision during training (fake quantization). This PR combines QAT with LoRA by applying torchao's QAT support to the peft model.

See the following for more details:

• torchao QAT: https://github.com/pytorch/ao/blob/main/torchao/quantization/qat/README.md
• torchtune QAT + LoRA: https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700

Current QAT schemes supported are:

fp8-fp8, targeting the torch.ops.fbgemm.f8i4bf16_shuffled kernel
fp8-int4, targeting the torch.ops.fbgemm.f8f8bf16_rowwise kernel

Test Plan:

from unsloth import FastLanguageModel

lora_rank = 32

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Qwen3-4B-Base",
    max_seq_length = 2048,
    load_in_4bit = False,
    fast_inference = False,
    max_lora_rank = lora_rank,
)

model = FastLanguageModel.get_peft_model(
    model,
    r = lora_rank,
    target_modules = [ 
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],  
    lora_alpha = lora_rank*2,
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    qat_scheme = "fp8-fp8",
)

lora.Linear(
  (base_layer): FakeQuantizedLinear(
    in_features=2560, out_features=4096, bias=False
    (activation_fake_quantizer): FakeQuantizer(Float8FakeQuantizeConfig(dtype=torch.float8_e4m3fn, granularity=PerRow(), hp_value_lb=None, hp_value_ub=None))
    (weight_fake_quantizer): FakeQuantizer(Float8FakeQuantizeConfig(dtype=torch.float8_e4m3fn, granularity=PerRow(), hp_value_lb=None, hp_value_ub=None))
  )
  ... 
)

[jerryzh168]

we could skip this in the config handler itself

[jerryzh168]

we also plan to remove the group_size arg for now since there is no benefit of making it larger according to Josh

[danielhanchen]

No need for try except since pyproject.toml will add it