feat(quantization): support targeted weight QAT for LoRA

[hasuoshenyun]

What does this PR do ?

Add targeted LoRA QAT for DeepSeek-V4-Flash, GLM-5.3-Flash and Qwen3, with actual training validation and measured improvements in train-to-quantized-reference inference consistency in our tested cases. Existing TorchAO INT4 configuration remains supported.

At a glance: supported examples and measured benefit

Model	Targeted weight format	Recorded validation
DeepSeek-V4-Flash	Routed-expert MXFP4, group32, E8M0 scales	Original full 43-layer backbone; both QAT and ordinary LoRA completed training and reference reload
GLM-5.3-Flash (glm5_next)	Routed-expert FP8, block128, FP32 scales	Original full 45-layer backbone; both QAT and ordinary LoRA completed training and reference reload
Qwen3-0.6B	Q/V-projection FP8, block32, E8M0 scales	Single-GPU pretrained-model QAT training smoke; no full-model paired benefit claim for Qwen

Measured in the matched V4/GLM cases: training-eval versus quantized-reference-reload logit RMSE fell from 0.2418 to 0.0000 for V4 and 0.8230 to 0.0000 for GLM. Top-1 agreement between those two execution paths increased from 91.0156% to 100.0000% and 67.5781% to 100.0000%, respectively. This is an improvement in behavioral consistency after merge-and-quantize, not a measured increase in task accuracy or inference throughput.

The full-model experiments used two training steps, synthetic data and FSDP2/EP32 on 4 nodes × 8 H100 GPUs, preserving backbone depth, width, vocabulary and expert count; MTP was disabled and GLM evaluation was text-only. They used local checkpoint/data overrides and forced eager execution, rather than running the published HellaSwag YAMLs unchanged. These historical runs precede subsequent refactoring/rebasing; current-SHA CI and sanitized full-run artifacts remain outstanding below. Native serving-engine parity is not claimed.

Why this matters: fine-tuning gains may not survive deployment quantization

A fine-tuning run can look successful in floating point, yet lose part of its learned behavior when its weights are merged and quantized for inference. This PR makes LoRA training account for that final quantization step instead of discovering the mismatch only after deployment.

When a model is fine-tuned in floating point but served with lower-precision weights, training and inference can execute materially different weight computations. This matters for workflows such as DeepSeek and GLM fine-tuning followed by merged FP4/FP8 inference:

1. Training can improve the floating-point model. The optimizer and floating-point validation see updates that have not yet undergone the intended deployment quantization.
2. Merge-and-quantize can weaken or erase those updates. For a fixed quantization grid, an update that stays inside the same rounding bin can produce no change in the encoded weight. Changes to block scales can also alter neighboring values. This is not a claim that every small update disappears, but it explains why apparently useful floating-point changes may not survive in the quantized model.
3. The deployed model may therefore fail to retain the behavior observed during training. Good training loss or floating-point validation alone cannot establish quality after quantization; the quantized model must be evaluated too.

The issue is floating-point training followed by a different, quantized inference computation—not an inherent inconsistency in all floating-point training. It can affect full-parameter fine-tuning as well as adapters; this PR's new FP8/MXFP4 path specifically addresses merged LoRA weights.

Workflow	Weights seen by the training forward	Weights used by the quantized reference inference
Ordinary LoRA, then merge and quantize	W_eff = W_base + delta_W	D(Q(W_eff))
This PR: QAT-LoRA	D(Q(W_base + delta_W))	D(Q(W_base + delta_W))

Here delta_W = (alpha / rank) * B @ A for a dense linear layer; grouped experts use the corresponding transposed storage layout. Q encodes the weights and D decodes them for the reference computation. This PR fake-quantizes the combined effective weight, not only the frozen base, and uses a straight-through estimator to update the adapters. The base remains frozen and is not permanently overwritten during training.

Why we report train–inference consistency: the measurements quantify how much a trained model's output changes when it is converted to its quantized reference form. They are intended to expose the gap that floating-point training metrics can hide, and to verify that QAT models the same weight transformation used by the reference inference path.

In the recorded matched full-backbone experiments, ordinary-LoRA versus QAT train-to-reference-reload RMSE was 0.2418 versus 0.0000 for V4-Flash, and 0.8230 versus 0.0000 for GLM-5.3-Flash. The protocol and KL/top-1 results are below. These measurements demonstrate the output mismatch and its removal under the tested QDQ contract; they do not establish that all observed drift is erased learning, or that QAT improves task accuracy. Ordinary-LoRA drift also includes merge rounding and operation-order changes.

The practical objective is to retain useful fine-tuning behavior after quantization. Exact reference alignment is an implementation requirement toward that objective, not a substitute for task-quality evaluation. It requires matched encoding, scales, rounding, dtype and operation order; arbitrary native serving kernels, activation quantizers and KV-cache formats are outside the demonstrated contract.

Changelog

• Use the existing QATConfig entry point with independent QAT target selection and typed numerical rules.
• Add FP8 block32/block128 with E8M0 or FP32 scales, and MXFP4 group32 with E8M0 scales.
• Apply fake quantization to the effective base-plus-adapter weight with identity-gradient STE.
• Gate unsupported precision, aliasing and distributed combinations explicitly.
• Isolate local packed-reference checkpoint persistence from training preparation.
• Add model-family examples for Qwen, DeepSeek-V4-Flash and GLM-5.3-Flash.
• Add scalar-oracle, forward/backward, corruption, fresh-process reload, and distributed numerical tests.

Before your PR is "Ready for review"

• Read the contributor guidelines and signed off commits.
• Added tests and repository documentation.
• Attach sanitized full-model experiment driver and original run artifacts alongside the recorded measurement table.
• Run upstream CI and documentation checks on the submitted SHA.

Additional Information

Depends on #3902 (ordinary expert LoRA correctness) and #3903 (opt-in typed configuration). Please review those independently; rebase this branch after they land. The QAT-specific commit is 11c262eff. Do not interpret the stacked diff as a request to silently change ordinary training semantics within QAT.

Validation

• On the rebased integrated branch: 1454 passed, 53 conditional skips on CPU (PyTorch 2.10 CPU, Transformers 5.15.1). Skips are not passes.
• Earlier GPU evidence includes real two-rank FSDP2/EP numerical tests, composed tiny-model tests, and full-depth V4/GLM expert-LoRA experiments. Those runs precede subsequent refactoring and rebasing; they are historical experimental evidence, not a GPU rerun of this exact PR SHA.
• Full-model paired measurements retained original backbone dimensions and expert counts, used two steps, fixed synthetic input, and matching initial adapters/source/config/topology. V4 ordinary/QAT RMSE was 0.2418/0.0000 and GLM 0.8230/0.0000 on 256 evaluation positions.

Recorded train-forward versus quantized-reference-reload measurements

Original backbone	Training arm	Logit RMSE	Mean KL/token	Top-1 agreement
V4-Flash, 43 layers	Ordinary LoRA	0.2418	0.0211	233/256
V4-Flash, 43 layers	QAT-LoRA	0.0000	0.0000	256/256
GLM-5.3-Flash, 45 layers	Ordinary LoRA	0.8230	0.4843	173/256
GLM-5.3-Flash, 45 layers	QAT-LoRA	0.0000	0.0000	256/256

Errors are displayed to four decimal places; QAT errors were measured as exact zero, not merely rounded to zero. Raw experimental values are retained separately.

These are bounded two-step experiments using expert-only rank-4/alpha-8 LoRA, BF16, fixed synthetic input and FSDP2/EP32. Each arm is compared against its own decoded reference output, not against the other arm's learned weights. The full-model probe reuses the same distributed model with disk-decoded local expert weights, zeroed adapters and identity quantizers to preserve the reference projection/routing path. It is not a fresh standalone inference engine. Native serving and task-quality gains are unverified, and sanitized reproduction artifacts remain an outstanding draft checklist item.

Important limitations

• The full-model comparison is same-distributed-model decoded local-weight reference inference, not native low-bit serving or a fresh standalone inference engine.
• QAT's exact reference match is expected under the matched QDQ contract. It is not a task-quality or convergence claim; baseline drift also includes merge/arithmetic-order differences.
• General packed export currently requires a complete local model. The full-model experimental shard probe is not a distributed QAT.export API.
• Supported training is restricted to the verified FSDP2/EP reference mechanism; unsupported TP/CP/PP/HSDP, additional expert-DP and optimized expert backends fail closed.
• FP32 scale representation support does not guarantee bitwise equivalence to every native FP8 encoder.

Please keep this PR as draft until reproducible evidence attachments and upstream CI review are complete.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-preview-11c262eff115.docs.buildwithfern.com/nemo/automodel