fix(quantization): Inherit Fields should not copy into empty provisional qspec - #96
Merged
Merged
Conversation
guru-desh
marked this pull request as ready for review
September 12, 2026 04:50
InheritFields should not copy into empty provisional qspec
InheritFields should not copy into empty provisional qspec
guru-desh
requested review from
dengqiaoyu
and removed request for
dengqiaoyu
September 12, 2026 04:57
guru-desh
force-pushed
the
fix-annotation-bug
branch
2 times, most recently
from
September 16, 2026 17:48
16acd06 to
a12cd06
Compare
anotheranshu
requested changes
Sep 16, 2026
anotheranshu
left a comment
Contributor
There was a problem hiding this comment.
The proposed change to the setter is optional, but can you add a test which runs prepare on a Linear -> Reshape -> Permute model?
anotheranshu
approved these changes
Sep 16, 2026
guru-desh
force-pushed
the
fix-annotation-bug
branch
from
September 17, 2026 01:12
0caf9f1 to
41c6048
Compare
guru-desh
enabled auto-merge (squash)
September 17, 2026 01:14
InheritFields guarded on the target's ProvisionalQSpec merely existing,
while its own comment promised "only slots that already hold fields ...
can inherit". ShareObserverInstance materializes an empty, non-declined
spec for any slot group no config spoke for, and inserts it into the
map. Such a slot passed the guard and received exactly _DATA_FACT_FIELDS
(QSCHEME, FLOAT_RANGE), leaving a partial spec.
_build_concrete_spec deliberately refuses to default a partial group --
"every observed slot is seeded from a whole QuantizationSpec, so a
partial group is a generation bug" -- so annotation then died with:
ReconciliationError: Cannot rebuild a QuantizationSpec: reconciled
group is missing ['DTYPE', 'QFORMULATION', 'GRANULARITY',
'FAKE_QUANTIZE_CLS', 'QPARAM_CALCULATOR_CLS', 'RANGE_CALCULATOR_CLS',
'SCALE_DTYPE']. Present: ['FLOAT_RANGE', 'QSCHEME'].
This reproduces on any graph-mode activation-quantized transformer, not
just MoE: the trigger is the ordinary QKV chain linear -> reshape ->
permute -> narrow, where the rank-changing passthrough sits beside slots
left unseeded by module_type_configs entries set to None.
Make the guard match the documented intent by also requiring the target
spec to hold fields. An unseeded slot stays empty, so
_build_concrete_spec returns None for it and the op is left in float --
which is the correct reading of "no config spoke for this slot".
Introduced in f226301 (apple#68).
Both the constraints module and the types module need the complete field set: ShareObserverInstance to reconcile across a group, and the next commit to decide whether a write leaves a qspec whole. Defining it next to the enum it is derived from leaves one definition instead of two.
A ProvisionalQSpec's field map is either empty, meaning nothing has spoken for the slot, or whole, carrying every FieldName. _build_concrete_spec rejects anything between the two, but the map was a public dict, so a constraint could write its own subset in and the failure only surfaced later, at resolution, naming the poisoned group rather than the constraint that poisoned it. Store the map privately, expose it as a read-only view, and route every write through merge_fields. Overwriting a field already present stays free: reconciliation does it on every pass and cannot change the key set. Introducing one is allowed only if it leaves the map whole, which is the only operation that can break the invariant. The constructor is not checked, because ShareObserverInstance builds from an already-reconciled group and so yields empty or whole either way, and _build_concrete_spec still backstops a map built by hand. ShareFields now raises where it previously broadcast one field into a slot no config addressed, which produced a DTYPE-only spec that _build_concrete_spec rejected later anyway. That is a deliberate behaviour change and it is not covered: nothing fails if ShareFields is changed to skip such a slot instead of raising. Reaching it needs a cat whose output is per-channel along the concat axis with an input no config addressed, and that path has no end-to-end coverage either. Dropping the dataclass, so callers keep a public fields= keyword alongside private storage, also swaps value equality and unhashability for identity equality and hashability. Nothing compares or hashes a ProvisionalQSpec today, since sharing is expressed with id(), and identity is the truer semantics for this type, but the change is silent.
The check reasoned about which keys an update introduced and then rebuilt the merged key set inline, describing the merged map twice without ever building it. Build it once, validate it, assign it. Behaviour is unchanged. The surviving clause, comparing the merged key set against the current one, is what limits the check to writes that add a key. Dropping it would also reject a map that was already partial before the write, which fails seven existing tests that build one- to three-field specs to exercise a single field's policy and then reconcile them. Assigning rather than updating in place is safe: the fields property builds a fresh view per read, and sharing is expressed by ProvisionalQSpec identity rather than by the dict object.
Nothing exercised this failure through Quantizer.prepare(). The reconciler unit tests build constraints by hand, so they pin the type's contract but not the pipeline that produced the crash, and the QKV chain that triggers it is three ops. linear -> reshape -> permute is the whole trigger: permute is rank-preserving so it emits ShareObserverInstance, which enters the shape ops into the map holding nothing, and reshape is rank-changing so it emits InheritFields, which writes two data facts into that spec. Reverting the InheritFields guard fails both tests; removing the type guard as well reproduces the original report, missing eight fields with FLOAT_RANGE and QSCHEME present. The first test asserts both shape ops are in the graph before asserting only linear is annotated, so export eliminating one cannot make it pass vacuously. The second pins three observers rather than two because the shared fixture quantizes weights; op_state_spec weight None gives two, and both configs reproduce the bug. The ShareFields half of the invariant is still uncovered, and reaching it needs a cat whose output is per-channel along the concat axis.
Co-authored-by: anotheranshu <anotheranshu@gmail.com>
41c6048 made merge_fields reject any write whose result is partial, not only one that adds a key. Seven tests built specs holding one to three fields on purpose, to exercise a single field's policy in isolation, so every overwrite in them left the map partial and raised before reaching an assertion. None of those tests is about how many fields a spec holds. Two cover drain-loop convergence, one covers ShareFields not merging spec objects, one covers a relaxed float range not being re-pinned. The partial map was incidental scaffolding, so the fix is in the two helpers that build it rather than in the tests, and no assertion changes. The placeholder values are identical across specs, which is what keeps this inert: every field policy reconciles them to the same value at the same priority, so no test sees a spurious changed slot. Passing no keyword still yields the empty map, because several tests need a slot no config spoke for; test_missing_field_returns_none fails if that branch goes. Mutation coverage is unchanged, including that reverting the InheritFields guard still fails test_skips_target_whose_spec_is_empty.
guru-desh
force-pushed
the
fix-annotation-bug
branch
from
September 17, 2026 16:09
9e1aa48 to
59f1564
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
While I was quantizing a model with graph mode set to True, I ran into this issue:
(shortened traceback to only focus on important parts)
This was occuring in a model that had an attention layer and did the following operations as part of that attention module:
In the first part of annotation reconciliation, we traverse the entire graph. Here, we make two provisional qspecs:
Because permute touches reshape, this part of the annotation reconciliation loop contains the reshape node. This then creates an InheritFields constraint. InheritFields apply then triggers which copies linear's qscheme and float range into permute's provisional qspec, which is the bug. It should not have copied linear's qscheme + float range as the provisional qspec was empty, so it creates a malformed provisional qspec triggering the error.
The fix is one line which is to add
target_qspec.fieldsinto the if statement so.apply()does nothing