lnn-quant-training

LNN Quantization-Aware Training

Core Workflow

Treat quantization-aware training as an engineering closed loop with acceptance criteria, rather than calling it once at the end of a floating-point training script linger.init. Unless the user requests only one of the stages, proceed in the order below.

1. Freeze acceptance conditions: Record task metrics, floating-point baseline, allowed accuracy drop, input shape/range, preprocessing, target platform, bit width, latency, and memory constraints. When the target platform is missing, generic 8/8-bit QAT can still be done, but no deployment compatibility conclusion may be given.
2. Audit the structure: Inspect operators, branches, normalization, activations, dynamic shapes, and output forms. Before designing or modifying the network, read references/model-design.md; hand target-platform deployment constraints to the peer $lnn-model-design skill for review.
3. Establish a floating-point baseline: Save the best float checkpoint, a fixed validation set, and the evaluation script, and keep a set of reproducible golden inputs/outputs. Without a float baseline using the same data, same preprocessing, and same metric, do not judge quantization loss.
4. Quick graph-export precheck: Before long training, use random weights or a float checkpoint to complete one trace_layers -> init -> forward -> linger.onnx.export, confirming that the structure can be exported. The precheck only proves that the graph can be generated; it does not prove quantization effectiveness.
5. Constrained training: Starting from the float checkpoint, perform one linger.trace_layers(...)as needed, call linger.constrain(...), rebuild the optimizer, and then do a short constrained fine-tuning cycle. First confirm the constrained-vs-float metric gap, then proceed to QAT.
6. Quantization initialization: Under the same structure transformation order, call linger.init(...), then load the matching checkpoint and rebuild the optimizer again. MOM causes the training forward to initialize running_data/scale; TQT tasks must first read ../doc/tutorial/tqt_training.md, and execute according to its backend selection, single calibration, optimizer override, and first backward gradient audit.
7. QAT fine-tuning: Usually start from about 0.1x of the floating-point learning rate, subject to the actual curve; at the same time record float, constrained, and QAT metrics and quantizer states, and save the complete QAT state_dict. For the complete implementation order, read references/training-workflow.md。
8. Effect diagnosis: First rule out checkpoint, preprocessing, BN/eval, uninitialized quantizer, and sample inconsistency issues, then locate the first high-error layer. When quantization effectiveness does not meet the target, you must read references/accuracy-recovery.md; when mixed precision, local clamp,const_module or quant_module is involved, you must also read references/mixed-precision.md, and iterate in order from low risk to high risk.
9. Graph export and validation: Strictly rebuild the structure and configuration used during training, load the QAT checkpoint, perform warm-up with real inputs, switch eval(), use torch.no_grad() to wrap linger.onnx.export(...). For the complete graph export and validation process, read references/export-and-validation.md。
10. Acceptance: Use references/report-template.md to output stage metrics, differences, configuration, commands, artifacts, open risks, and next steps. Only when the ONNX checker, quantized graph audit, and target runtime consistency all pass may it be called a deployment closed loop.

Project Fact Priority

Prioritize the live source and tests of the current checkout, then refer to tutorials; do not directly treat enums or capability descriptions in old documentation as current implementation facts. For entry points, see references/project-sources.md; for configuration fields and current limitations, see references/configuration.md。

Mandatory Rules

• Put linger.trace_layers、linger.constrain and linger.init are all regarded as operations that change the model structure; rebuild the optimizer/scheduler after each structural change.
• Always keep three independent checkpoints—float, constrained, and QAT—and do not overwrite the upstream baseline.
• Training, resuming, evaluation, and graph export must use the same model code, fusion order, config,disable_module/disable_submodel and input contract.
• After loading a checkpoint, check missing_keys and unexpected_keys; unless the differences have been explained item by item, do not use strict=False to mask problems.
• disable_module modifies the process-level quantization registry; use it only for isolated experiments, and prefer starting a new process for new controlled experiments.
• In the current code, linger.calibration() 's built-in calibration function depends on TQT's learning_data/is_calibrate; MOM does not go through that context and directly updates statistics from the training forward pass.
• TQT training uses quant_method: NATIVE or CUDA_GS; currently CUDA backward does not return learning_data gradients. Before validation and graph export, refresh the final scale according to ../doc/tutorial/tqt_training.md .
• For the final mixed precision solution, prefer writing the complete YAML partN; API overrides are only for localization or dynamic tensor ops, and must be replayed during resume/graph export. First localize sensitive layers, then confirm target hardware support; effective high-bit-width training does not mean it is deployable on Thinker/the target chip.
• Do not use whether ordinary ONNX Runtime can execute linger custom-domain ops as the criterion for successful graph export; prioritize ONNX checker, graph attribute auditing, and Thinker/runtime consistency.
• Change only one major variable at a time, and save the same-set metrics and output differences before and after the change.

Executable tools

Use under the toolchain root:

python3 lnn-quant-training/scripts/generate_qat_config.py --output configs/qat.yaml --platform venusA --device cuda --qat-method MOM
python3 lnn-quant-training/scripts/inspect_quant_checkpoint.py checkpoints/best_qat.pt --format markdown
python3 lnn-quant-training/scripts/compare_outputs.py artifacts/float_outputs.npz artifacts/qat_outputs.npz --format markdown
python3 lnn-quant-training/scripts/inspect_quant_onnx.py artifacts/model_quant.onnx --format markdown --fail-on-issues

The scripts' detailed purposes, input conventions, and return codes are written in their respective --help . For any checkpoint, load only from trusted sources; inspection scripts should preferentially use PyTorch's weights_only=True。

Deliverables

At minimum keep:

• float, constrained, and best-QAT checkpoints and their metrics.
• Final YAML, training parameters, random seeds, code version, and input contract.
• golden input, float output, QAT output, and output comparison report.
• Exported ONNX, ONNX audit report, export commands, and logs.
• When Thinker is available, the pkg, memory report,tvalidator/runtime results; when not available, clearly mark as unverified.