Skip to content

fix(quantization): Inherit Fields should not copy into empty provisional qspec - #96

Merged
guru-desh merged 7 commits into
apple:mainfrom
guru-desh:fix-annotation-bug
Sep 17, 2026
Merged

guru-desh merged 7 commits into
apple:mainfrom
guru-desh:fix-annotation-bug

Conversation

@guru-desh

@guru-desh guru-desh commented Sep 12, 2026

Copy link
Copy Markdown
Contributor

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)

    def prepare(
        self,
        example_inputs: tuple[Any, ...],
        dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
        export_with_no_grad: bool = True,
    ) -> torch.fx.GraphModule:
	...
        try:
>           prepared_model = prepare_qat_pt2e(exported_model, quantizer)
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
...
.venv/lib/python3.11/site-packages/coreai_opt/quantization/_graph/_qspec_resolution.py:145: in _assign_group_specs
    concrete = _build_concrete_spec(group.qspec)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
...

    def _build_concrete_spec(qspec: ProvisionalQSpec) -> TorchAOQuantizationSpec | None:
...
        missing = [
            field_name.name for field_name in _SPEC_KWARG_FROM_FIELD if field_name not in qspec.fields
        ]
        if missing or FieldName.QUANTIZATION_TARGET not in qspec.fields:
            # Every observed slot is seeded from a whole QuantizationSpec, so a
            # partial group is a generation bug. Don't paper over it with defaults.
>           raise ReconciliationError(
                f"Cannot rebuild a QuantizationSpec: reconciled group is missing "
                f"{missing or ['QUANTIZATION_TARGET']}. Present: "
                f"{sorted(f.name for f in qspec.fields)}."
            )
E           coreai_opt.quantization._graph._qspec_types.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'].
.venv/lib/python3.11/site-packages/coreai_opt/quantization/_graph/_qspec_resolution.py:204: ReconciliationError

This was occuring in a model that had an attention layer and did the following operations as part of that attention module:

linear -> reshape -> permute

In the first part of annotation reconciliation, we traverse the entire graph. Here, we make two provisional qspecs:

  1. Linear gets a filled provisional qspec as it was targeted in the user config
  2. permute gets a provisional qspec as ShareObserverInstance creates one as its rank preserving. This provisional qspec covers reshape's output + permute's input + output

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.fields into the if statement so .apply() does nothing

@guru-desh guru-desh added the bug Something isn't working label Sep 12, 2026
@guru-desh
guru-desh marked this pull request as ready for review September 12, 2026 04:50
@guru-desh guru-desh changed the title [WIP] fix(quantization): don't inherit data facts into an unseeded qspec [WIP] fix(annotation): inherit fields should not copy into empty provisional qspec Sep 12, 2026
@guru-desh guru-desh changed the title [WIP] fix(annotation): inherit fields should not copy into empty provisional qspec fix(annotation): inherit fields should not copy into empty provisional qspec Sep 12, 2026
@guru-desh guru-desh changed the title fix(annotation): inherit fields should not copy into empty provisional qspec fix(quantization): inherit fields should not copy into empty provisional qspec Sep 12, 2026
@guru-desh guru-desh changed the title fix(quantization): inherit fields should not copy into empty provisional qspec fix(quantization): InheritFields should not copy into empty provisional qspec Sep 12, 2026
@guru-desh guru-desh changed the title fix(quantization): InheritFields should not copy into empty provisional qspec fix(quantization): Inherit Fields should not copy into empty provisional qspec Sep 12, 2026
@guru-desh
guru-desh requested review from dengqiaoyu and removed request for dengqiaoyu September 12, 2026 04:57
@guru-desh
guru-desh force-pushed the fix-annotation-bug branch 2 times, most recently from 16acd06 to a12cd06 Compare September 16, 2026 17:48
Comment thread src/coreai_opt/quantization/_graph/_qspec_types.py

@anotheranshu anotheranshu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The proposed change to the setter is optional, but can you add a test which runs prepare on a Linear -> Reshape -> Permute model?

Comment thread src/coreai_opt/quantization/_graph/_qspec_types.py Outdated
Comment thread src/coreai_opt/quantization/_graph/_qspec_types.py Outdated
@guru-desh
guru-desh enabled auto-merge (squash) September 17, 2026 01:14
guru-desh and others added 7 commits September 17, 2026 09:09
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
guru-desh merged commit 0388751 into apple:main Sep 17, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants