Enable static quantization for Qwen3-0.6B decoder (transformer-only)#836
Enable static quantization for Qwen3-0.6B decoder (transformer-only)#836spalne wants to merge 12 commits into
Conversation
spalne
changed the title
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
@spalne please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
DingmaomaoBJTU
left a comment
DingmaomaoBJTU
left a comment
There was a problem hiding this comment.
Summary - structurally sound export, but registration/test/quant integration don't match repo conventions, and w8a16 accuracy regresses.
Nice work getting a fused GQA + LpNorm RMSNorm + 1x1-Conv transformer-only export running end-to-end on QNN, and the export itself is faithful - the FP optimized graph reproduces HF eager's next-token exactly. Three things to address before this is review-ready:
1. Registration is non-standard (highest priority). qwen_transformer_only.install() hot-patches the global registries at runtime and isn't imported by models/hf/__init__.py. Every other model registers declaratively at import time (@register_onnx_overwrite / @register_composite_model, merged in __init__.py). Please make this a first-class variant (distinct task/model_type or a build-config flag) instead of monkey-patching; it also removes the "must call install() before importing the composite machinery" ordering trap and the no-way-back override of the eager path.
2. Test & quant entry points violate repo layout. test_qwen.py and qwen3_transformer_only_quantize.py are standalone scripts at the repo root; test_qwen.py is a subprocess driver that judges success by artifact mtime and uses os._exit(0) to mask a native QNN/ORT teardown crash. Convention (tests/CLAUDE.md) is pytest under tests/. Move the runner to tests/e2e/ (or examples/), and wire the calibration reader into the config-driven quant flow (WinMLBuildConfig.quant) rather than a bespoke quantizer.
3. w8a16 accuracy is not yet acceptable. Measured against the FP graph on the same GSM8K-style input, the quantized model flips the top-1 next token on both prefill and decode (top-5 overlap 0-1/5, KL 0.66/2.75; hidden-state cosine 0.64-0.72), while present-KV stays ~0.999 - i.e. the residual stream is the casualty. Likely minmax + all-zero KV calibration + only 30 samples. Please try percentile/entropy calibration with a realistic non-zero KV feed and report an actual task metric, not just QDQ node count.
Naming and the custom-op export pattern look good and match the codebase.
DingmaomaoBJTU
left a comment
DingmaomaoBJTU
left a comment
There was a problem hiding this comment.
Code Review — PR #836 (Draft)
Well-structured PR. The transformer-only export topology (fused GQA, LpNorm RMSNorm, 1x1 Conv), GSM8K calibration pipeline, and model_type override mechanism are solid. A few correctness bugs and infrastructure concerns should be resolved before marking ready for merge.
Not approving since this is a draft PR.
| for i in range(num_layers): | ||
| result[f"past_keys_{i}"] = {2: kv_seq_axis} | ||
| result[f"past_values_{i}"] = {2: kv_seq_axis} |
There was a problem hiding this comment.
Non-blocking / out of scope for the quant fix — recording another export/pipeline difference vs the reference graph.
The KV time axis is declared symbolic here (kv_seq_axis="max_seq_len", applied to past_keys_{i}/past_values_{i} axis 2), which matches the reference. But the produced/optimized graph ends up with a static axis (measured on the ctx model):
- This PR (MINE):
past_keys_0axis2 = 256 (static). - Reference:
past_keys_0axis2 =max_seq_len(symbolic).
So the symbolic dim declared here is being frozen to the concrete max_cache_len (256) somewhere downstream — most likely the same ORT optimize pass that produces _optimized.onnx. A static 256 cache hard-codes the max sequence length into the graph (less flexible for longer contexts / different cache sizes) whereas the reference stays parametric. Doesn't affect quant numerics; flagging so the symbolic axis is preserved through the optimize/export step if that flexibility is wanted.
There was a problem hiding this comment.
Confirmed out of scope here. The symbolic KV axis being frozen to a static dimension is a property of the static-shape HTP export path (the exporter does not apply dynamic_axes for this model_type), not of the quantization change. Tracked as a future export-pipeline improvement.
Replace the standalone root-level quant driver and __main__/subprocess test runner with the regular build pipeline and pytest. - Move calibration logic into src/.../hf/qwen_transformer_only_quant.py; the decode wrapper exposes winml_finalize_quant_config, invoked generically from build/hf.py just before quantize_onnx. The build now quantizes via precision=w8a16 + config.quant instead of a separate script. - The hook reads seq_len / max_cache / GQA node names from the exported ONNX and selects the prefill vs decode-trajectory calibration reader, keeping the verified-good scheme (int8-symmetric weights, uint16 activations, minmax, GQA excluded from QDQ). - Delete root qwen3_transformer_only_quantize.py and test_qwen.py. - Add tests/unit/models/qwen_transformer_only (fast, offline) and tests/e2e/models/test_qwen3_transformer_only_quant.py (build+quant+decode-parity, QNN-gated NPU).
# Conflicts: # src/winml/modelkit/loader/config.py # src/winml/modelkit/models/auto.py
…c shapes - Add missing docstrings / return-type annotations and drop dead noqa directives across qwen3_export_ops.py, qwen3_modeling.py and the transformer-only registration so 'ruff check src/ tests/' (CI lint) passes. - build/hf.py: re-persist config.json after winml_finalize_quant_config runs, so the saved metadata reflects the actually-applied w8a16 scheme (int8/uint16/symmetry + GQA nodes_to_exclude) rather than the pre-finalize policy dtypes. - qwen_transformer_only_quant._graph_shapes: treat a non-positive dim_value (symbolic/dynamic axis) as a hard error instead of silently returning a zero-length shape.
…2e helper) - LpNormOnnxExport.forward now computes the real L2 normalization instead of a silent identity; export-invariant (node comes from symbolic) and correct in eager. - GroupQueryAttentionOnnxExport.forward keeps the non-raising placeholder, with a docstring explaining why raising is impossible (HTP hierarchy capture runs an eager forward outside trace/export). - Remove unused module-level logger in qwen_transformer_only.py (CodeQL). - Use a single onnx import form in test_quant_calibration.py (CodeQL). - Fix e2e _decoder_onnx_path helper to handle the single-model WinMLModelForGenericTask (.onnx_path) build, not just composite .sub_models.
…_type-override test - build_hf_model: look up winml_finalize_quant_config on type(pytorch_model) instead of the instance, and call it with explicit self. Fixes the mypy 'Tensor not callable' error (getattr yields Any) and stops the hook firing on raw HF models / MagicMock test doubles (whose attributes are instance-synthesized), which was serializing a MagicMock into config.json. - test_resolve_loader_config: replace the obsolete 'never mutated' test with one asserting the intended explicit-model_type override (needed for variants like qwen3_transformer_only).
Relocate the model-specific transformer-only calibration/quant logic out of
models/hf (an export-only package) into a new quant/calibration/ subpackage,
dispatched via a model_type-keyed registry that mirrors COMPOSITE_MODEL_REGISTRY.
- Add quant/calibration/{base,registry}.py: QuantConfigFinalizer protocol +
register_quant_finalizer / get_quant_finalizer (lazy, torch-free import).
- git mv qwen_transformer_only_quant.py -> quant/calibration/qwen3_transformer_only.py
and register Qwen3TransformerOnlyQuantFinalizer for 'qwen3_transformer_only'.
- build/hf.py: replace the winml_finalize_quant_config wrapper hook with explicit
registry dispatch keyed on config.model_type; unregistered types fall back to
the default DatasetCalibrationReader. Preserve the model_id/_name_or_path
fallback (now model-agnostic in the build layer).
- Remove the hook from the export wrapper (back to export-only).
- Relocate unit tests to tests/unit/quant/calibration/ and add test_registry.py.
w8a16 scheme unchanged; CPU e2e (quantized-graph + GQA-exclusion + FP-parity)
and 86 build/quant unit tests pass.
| __all__ = [ | ||
| "QuantizeResult", | ||
| "WinMLQuantizationConfig", | ||
| "get_quant_finalizer", |
| "WinMLQuantizationConfig", | ||
| "get_quant_finalizer", | ||
| "quantize_onnx", | ||
| "register_quant_finalizer", |
| model_id: str | None = None, | ||
| ) -> WinMLQuantizationConfig: | ||
| """Return ``quant`` populated with the graph-derived quant settings.""" | ||
| ... |
3 participants
Adds a transformer-only ONNX export path for Qwen3 that emits a fused (GQA) GroupQueryAttention op (with built-in rotary), LpNormalization RMSNorm, and 1×1 Conv projections, backed by an FP16 KV cache. The path is opt-in via install(), which hot-patches the build registries to produce two graphs (prefill seq=64, decode seq=1) without embeddings or lm_head. Quantization runs w8a16 static PTQ on these graphs using GSM8K calibration
Results
Produces two transformer-only ONNX files (prefill + decode) plus their w8a16-quantized variants.