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 6 commits into
Conversation
spalne
changed the title
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| from .qwen3_modeling import ( | ||
| WinMLQwen3Attention, | ||
| WinMLQwen3DecoderLayer, | ||
| WinMLQwen3MLP, | ||
| WinMLQwen3Model, | ||
| WinMLQwen3RMSNorm, | ||
| ) |
| from .qwen3_export_ops import ( | ||
| GroupQueryAttentionOnnxExport, | ||
| LpNormOnnxExport, | ||
| TransposeConv2d1x1Transpose, | ||
| ) |
|
|
||
| COMPOSITE_MODEL_REGISTRY[("qwen3", "text-generation")] = WinMLQwen3TransformerOnlyModel | ||
|
|
||
| _INSTALLED = True |
|
@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.
| @@ -0,0 +1,235 @@ | |||
| """E2E test for the transformer-only Qwen3 export path. | |||
There was a problem hiding this comment.
This is a standalone runner at the repo root that drives the build via subprocess and judges success by "did a fresh artifact file appear". Repo convention (and tests/CLAUDE.md) is pytest under tests/ with code-generated expectations - there are no other root-level test_*.py scripts. Could this move under tests/e2e/ as a real pytest (marked e2e/npu/qnn), or under examples/ if it's really a demo rather than a test? As-is it'll get picked up by name but isn't a pytest, and it lives outside the tree the suite runs from.
| @@ -0,0 +1,230 @@ | |||
| """Transformer-only w8a16 quantization for Qwen3. | |||
There was a problem hiding this comment.
Quantization in winml-cli is normally config-driven through WinMLBuildConfig.quant and runs as part of the build pipeline. This adds a parallel standalone quant entry point at the repo root that reaches into sub_models[*]._onnx_path directly and is "run via test_qwen.py". Could the transformer-only calibration reader be wired into the standard quant flow so it's reachable from winml build / the config instead of a bespoke script? Also minor: Qwen3TransformerOnlyCalibReader structurally satisfies winml.modelkit.quant.config.CalibrationDataReader but doesn't declare it - worth importing/typing against the protocol so it stays in sync.
| samples=num_samples, | ||
| weight_type=weight_type, # type: ignore[arg-type] | ||
| activation_type=activation_type, # type: ignore[arg-type] | ||
| calibration_method="minmax", |
There was a problem hiding this comment.
Accuracy concern worth resolving before this lands. I ran the produced w8a16 graphs against the FP optimized graphs on the same GSM8K-style input (ORT CPU EP): the FP export matches HF eager exactly (top-1 next token identical), but the w8a16 output flips the top-1 token on both prefill and decode - top-5 overlap 0-1/5, KL(FP||quant) 0.66 / 2.75, output_hidden_states cosine 0.64-0.72. The present-KV path is ~0.999, so the damage is concentrated in the residual stream.
Likely causes: minmax calibration over a residual stream with large outliers (+/-76), calibrating with an all-zero KV cache, and only 30 samples. Suggest trying calibration_method="percentile" (or entropy), feeding a realistic non-zero KV during calibration, and reporting an actual task metric (e.g. GSM8K logits/top-1 agreement) so we can see the quant is acceptable, not just that QDQ nodes were inserted.
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.
|
|
||
| @staticmethod | ||
| def forward(ctx, input, axis, p): # noqa: ARG004 | ||
| return input # placeholder — real compute happens in symbolic |
There was a problem hiding this comment.
Warning: Eager-mode forward returns incorrect (un-normalized) results
LpNormOnnxExport.forward returns input unchanged (identity). This is only correct during ONNX tracing where symbolic runs instead. Any eager execution (unit tests, calibration debug runs) silently gets un-normalized values. Consider computing the real norm for eager mode or raising NotImplementedError to make misuse obvious.
| kv_num_heads, | ||
| num_heads, | ||
| ): # noqa: ARG004 | ||
| return query, past_key, past_value # placeholder shapes |
There was a problem hiding this comment.
Warning: Stale KV cache in eager mode
GroupQueryAttentionOnnxExport.forward returns (query, past_key, past_value) — the present_keys/present_values are the old un-updated tensors. Eager execution silently produces a KV cache that never advances. A NotImplementedError here would be safer than a silently-wrong placeholder.
| @@ -0,0 +1,229 @@ | |||
| """E2E test for the transformer-only Qwen3 export path. | |||
There was a problem hiding this comment.
Info: Not a pytest test
Despite the test_ prefix, this file uses __main__, sys.path mutations at import time, subprocess orchestration, and os._exit. It lives in the repo root and won't be collected by uv run pytest tests/. Consider renaming to scripts/run_qwen3_quant.py to avoid accidental pytest collection, or convert to a proper pytest integration test with hardware skip markers.
| feed: dict[str, np.ndarray] = { | ||
| "input_hidden_states": embeds.astype(np.float32), | ||
| # seqlens_k for GQA = (valid context length - 1), i.e. | ||
| # ``embeddings.shape[1] - 1``. We pad to seq_len, so the query | ||
| # has seq_len valid positions → past_seq_len = seq_len - 1. | ||
| # (Using 0 here declares only 1 valid token while feeding a | ||
| # seq_len-token query, which makes the GQA prefill kernel read | ||
| # out of bounds → native access violation.) | ||
| "past_seq_len": np.array([[self.seq_len - 1]], dtype=np.int32), | ||
| "total_seq_len": np.array([self.max_cache_len], dtype=np.int32), | ||
| } | ||
| kv_shape = (1, self.num_kv_heads, self.max_cache_len, self.head_dim) | ||
| zeros = np.zeros(kv_shape, dtype=np.float16) | ||
| for i in range(self.num_layers): | ||
| feed[f"past_keys_{i}"] = zeros | ||
| feed[f"past_values_{i}"] = zeros |
There was a problem hiding this comment.
Follow-up from the accuracy verification: the symmetric-int8 + GQA-exclusion fix correctly aligns the graph with the reference, but the produced model still degenerates in generation — it emits a correct first token and then collapses into repeated garbage (e.g. Paris → Parisammedammed...).
I traced it to this calibration reader, specifically the decode (seq_len=1, decoder_gen) path. Stage-isolation (mixing FP/quant ctx & iter) shows the prefill quant is fine; only the quantized iter breaks decode, and since ctx/iter share identical int8 weights and activations are uint16 (plenty of precision), the culprit is the iter activation calibration ranges, i.e. the data fed here:
ids = ids[:, : self.seq_len]withseq_len=1→ afterapply_chat_template, every sample is the same first token<|im_start|>(151644). I verified across prompts: the decode reader sees one identical token repeated, so MinMax effectively calibrates on a single activation value.past_seq_len = seq_len - 1 = 0→ declares an empty context.past_keys_/past_values_ = np.zeros(...)→ attention is calibrated with an empty KV cache.
So the iter activation ranges are derived from "one unrepresentative token + empty KV + length 0", which is far narrower than real decode. Real decode activations then saturate → collapse. (Prefill survives because seq_len=64 gives ~1920 diverse real states.)
Verified fix (single variable changed): I re-quantized the FP iter changing only calibration_data — keeping weight_type=int8/weight_symmetric=True, activation_type=uint16, calibration_method=minmax, and the same nodes_to_exclude=gqa_nodes — and fed a real decode trajectory instead: run a short FP prefill + N decode loop and capture each step's true feed (current token embedding, accumulated KV, growing past_seq_len), 10 prompts × 16 steps = 160 samples. Result (e2e greedy, same ctx/emb/head):
| prompt | current iter calib | trajectory calib | fp16 reference |
|---|---|---|---|
| The capital of France is | Parisammedammed... |
Paris, and the capital of Italy is Rome... |
Paris, and the capital of Italy is Rome... (token-identical) |
| 2 + 2 = | garbage | 4, 2 + 2 = 4... |
4, 3 + 3 = 6... |
| The opposite of hot is | garbage | cold, and the opposite of cold is hot... |
cold, and the opposite of cold is hot... |
QDQ count is unchanged (2279), so this is purely a calibration-data issue, not graph structure or bit-width.
Suggested change: for the decoder_gen (seq_len=1) sub-model only, replace the single-token + zeroed-KV feed with real decode-step states (a brief FP prefill+decode trajectory). The decoder_prefill (seq_len=64) path is fine as-is. Happy to share the prototype script that produces the trajectory calibration if useful.
There was a problem hiding this comment.
Non-blocking follow-up / request for the reference calibration recipe.
I compared this path's w8a16 transformer initializers against the reference bundle (qwen3_gqa_fp16_ctx.onnx / qwen3_gqa_fp16_iter.onnx) by hashing every initializer's bytes and classifying each FLOAT scale as a weight scale vs an activation scale (by whether the (De)QuantizeLinear it feeds quantizes an INT8 weight init or a runtime tensor):
| INT8 weights | weight scales | activation scales | |
|---|---|---|---|
| prefill / ctx | 392/392 identical | 196/196 identical | 122/704 identical |
| decode / iter | 392/392 identical | 196/196 identical | 121/704 identical |
So the quantization scheme + the deterministic RTN weights + the weight scales are already byte-identical to the reference — the only thing that diverges is ~83% of the activation scales, which are purely calibration-derived. That's expected (activation ranges are data-dependent), and I don't think byte-matching them is the right goal — functional task-metric parity is. But it does mean the only free variable left between the two pipelines is the calibration.
To make the accuracy comparison apples-to-apples (and to rule out "we calibrated on GSM8K, the reference used something else" as a confound), could you share the reference calibration recipe — not just the dataset, but the whole thing:
- dataset (and how many samples / selection order / tokenization),
- calibration method (minmax / percentile / entropy, and any clipping),
- how the activations + KV are fed during calibration for the iter (decode, seq_len=1) sub-model — i.e. the equivalent of the reference's
CalibrationDataReader. This is the important one: as documented above, the decode collapse here comes from feeding a single token + zeroed KV, and I expect the reference feeds real decode-trajectory states. The dataset alone won't fix that — the feed mechanics will.
Note this isn't a blocker for landing: the decode fix (real-trajectory calibration) is independent of the reference data, and we can validate via a task metric. The recipe would just let us do a clean side-by-side on the same distribution and potentially close the activation-scale gap.
Evidence scripts: temp/weight_overlap.py (initializer byte-overlap by dtype) and temp/scale_classify.py (weight-scale vs activation-scale classification + overlap).
| if final_path.name.endswith("_model.onnx"): | ||
| stem = final_path.name[: -len("_model.onnx")] | ||
| optimized = final_path.with_name(f"{stem}_optimized.onnx") | ||
| if optimized.exists(): | ||
| sub_paths[name] = optimized |
There was a problem hiding this comment.
Non-blocking / out of scope for the quant fix — just recording a graph-hygiene difference for later.
The *_optimized.onnx consumed here carries a bloated set of opset imports compared with the reference graph (measured on the produced ctx model):
- This PR (MINE): main domain opset 17 plus 8 extra domains —
com.microsoft,com.microsoft.nchwc,ai.onnx.ml,ai.onnx.training,ai.onnx.preview.training,com.microsoft.experimental,org.pytorch.aten,com.microsoft.dml. - Reference: main domain opset 18 +
com.microsoftonly (2 domains total).
The extra nchwc/dml/training/aten domains are artifacts of the ORT optimize pass that produces _optimized.onnx; they're not used by the actual nodes. This doesn't change quant numerics, but it does inflate the graph metadata and can confuse downstream EP partitioning/loaders. Worth cleaning up the optimize pass (or stripping unused opset imports) at some point — not required for this 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.
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.