Constraint-queue reconciler for quantization annotation#40
Draft
anotheranshu wants to merge 6 commits into
Draft
Constraint-queue reconciler for quantization annotation#40anotheranshu wants to merge 6 commits into
anotheranshu wants to merge 6 commits into
Conversation
added 6 commits
July 14, 2026 10:41
Rewrites _AnnotationHandler.annotate around an explicit constraint-queue reconciler that operates on a per-slot ProvisionalQSpec state. Slots sharing an observer at runtime reference the same ProvisionalQSpec object (identity IS the sharing relation); ShareFields and ShareObserverInstance constraints drive per-field reconciliation (dtype-priority-wins, qscheme-lattice-join, range-union) and observer- instance merging. Pattern-driven generators (adjacent-edge, shared-observer, shared-weight-tensor) emit constraints as slot state evolves. Convergence is guaranteed by per-field priority monotonicity. Fixes the YOLOX-shape cat-of-sigmoid case by construction: sigmoid's op-intrinsic per_tensor_affine qscheme participates in the cat group's reconciliation and propagates via lattice-join to override the wider input's symmetric qscheme without an order-dependent traversal. New files: - src/coreai_opt/quantization/_graph/_qspec_reconcile.py - external/tests/quantization/test_qspec_reconcile.py (22 tests) Existing helpers (adjust_output_qspec_for_qscheme_and_propagate, _get_output_qspec's child-inheritance branch, _propagate_output_qspec, _adjust_input_qspec_map_for_shared_observers) in _annotation_utils.py are no longer called from annotate() but are left in place for now. The 5-line SharedQuantizationSpec early-return stopgap at :267 is dead code under the new pipeline but harmless. Test status: test_qspec_reconcile.py 22/22 pass; test_annotation_pattern_registry.py 101/102 pass. The one failure (test_shared_param_precedence op_state_multiple_specs_with_None) needs "explicit None" represented as a first-class annotation signal; design discussion pending.
Break up the 960-line _qspec_reconcile.py into four phase-specific
modules plus a shared types module:
_qspec_types types + ReconciliationError
_qspec_constraints Constraint ABC + per-field policies
_provisional_qspec_generation Phase 2 (build initial state)
_qspec_resolution Phase 4 (write per-node annotations)
_qspec_reconcile Entry point + edge/state generators
Behavioral changes carried alongside the split:
* Raise ReconciliationError (was silent) when an op-intrinsic node
has all consumers in-pattern, or when a covered node has both
in-pattern and external consumers -- pattern-shape assumptions
that silently dropped intrinsic contributions before.
* Extend op-intrinsic override past qscheme: when the user's dtype
matches the intrinsic dtype (sigmoid/tanh/hardsigmoid), also
enforce range and scale/zp via FixedQParamsFakeQuantize.with_args;
log every field the override changes.
* Give OBSERVER_CLASS a lattice policy -- non-fixed observers beat
fixed ones instead of must-agree failing on any mismatch.
* Move shared-observer constraint generation into pattern classes
via a new abstract generate_qspec_sharing_constraints so each
pattern owns its sharing semantics (ConcatPattern's axis-aware
logic vs Flatten/MaxPool/AvgPool's whole-boundary sharing).
* Rename State to ProvisionalQSpecMap; introduce SlotOrderKey
dataclass in place of the bare tuple used for anchor selection.
_qspec_reconcile.py now contains only annotate_via_reconciliation (the drain-loop driver). Everything else that used to live there -- the _AnnotationContext bundle, _generate_constraints_for_node, _adjacent_edge_constraints, _shared_state_constraints, _find_input_slot_for_producer, _nodes_covered_by -- moves to a new _qspec_constraint_generation module. The two context/coverage symbols quantizer.py imports (_AnnotationContext, _nodes_covered_by) are re-exported from _qspec_reconcile so no call sites move.
The reconciler only ever assigns one of two values to a slot -- a concrete TorchAOQuantizationSpec (anchor / singleton group) or a SharedQuantizationSpec pointing at the anchor (non-anchor slot in a multi-slot group). Replace the Any placeholder with a _SlotSpec union alias, and thread it through _assign_group_specs, the two _bucket_per_node output dicts, _write_annotations, and _backfill_input_qspec_map (the one legitimate None entry point, whose return type is _SlotSpec | None).
The reconciler at annotate_via_reconciliation() has been the sole live annotation code path since it landed; the old per-node annotator callbacks, their dispatch machinery, and the Phase-6 qscheme post-pass were still defined but unreachable. Deleted: - _annotation_utils.annotate_weighted_mod_match, annotate_n_ary_act_match, annotate_shared_observer_match -- per-match annotator callables, no longer invoked - _annotation_utils.adjust_output_qspec_for_qscheme_and_propagate and its _propagate_qscheme_to_child_nodes helper -- Phase-6 post-pass replaced by op-intrinsic override in the reconciler - _annotation_utils._get_output_qspec, _propagate_output_qspec, _adjust_input_qspec_map_for_shared_observers -- reachable only from the deleted annotator callables - _annotation_utils.is_all_annotated, is_any_annotated, mark_nodes_as_annotated, find_consumers -- zero remaining callers - _annotation_pattern_registry.AnnotatorFunc type alias, BaseAnnotationPattern.get_annotator_func (+ all three concrete overrides), AnnotatorMatchInfo.annotator_func field -- populated but never read - _annotation_config.AnnotationContext.shared_observer_nodes field -- no longer read anywhere (kv-cache override only needs module_name_to_state_names_map) - test stub in test_graph_mode_quantizer.py::temp_pattern fixture -- get_annotator_func is no longer an abstract method to satisfy _get_input_qspec_map and its state-lookup helpers stay -- they are still called by _AnnotationHandler._override_cache_op_annotations to build kv-cache op annotations.
Covers the pattern that motivated the constraint-queue reconciler: a cat group whose branches carry disagreeing qscheme proposals (one symmetric conv branch + two affine-forced sigmoid branches). Asserts the load-bearing invariants of the full pipeline: - topo-first anchor selection (reg conv, not one of the sigmoids) - lattice-join of qscheme to per_tensor_affine - SharedQuantizationSpec on every non-anchor slot in the group - internal-edge suppression on the two convs feeding the sigmoid pattern (no output_qspec on them) - prepared model still runs forward
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
1 participant
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.
Rewrite of the graph-mode annotation handler. Concretely, the previous handler failed on a subgraph within YOLO: we concatenated the outputs of a Conv and two sigmoids:
concat(conv(...), sigmoid(...), sigmoid(...))Sigmoid's output has a fixed qspec with a (0,1) range, conv has a "floating" qspec. In the old model, the winning config would be order-dependent, leading to a crash in the annotation propagation logic where we expected a fixed qspec for the sigmoid node but instead received a shared qspec.
This PR changes the model from an implicit prioritization (where we would traverse the patterns in order, propagating as we go, directly writing the qspecs back to the graph) to an explict prioritization. Here's how that looks:
The main driver is
_qspec_reconcile.py, in particularannotate_via_reconciliation. This expects pattern matching to be done already, and the pattern matching is untouched by this PR.Once we have the patterns in priority order, we construct a "provisional qspec" for each Node, (see
_provisional_qspec_generation.py). This constructs a map fromNodeSlots toProvisionalQSpecs. A NodeSlot is basically a reference in the graph to a tensor that could be quantized. Either it's the input (including state) to a Node, or the output of a Node. A ProvisionalQSpec stores the current best QSpec for the tensor. Afterbuild_initial_provisional_qspecs, this is basically what each node would like to quantize its tensors to if there were no constraints imposed from other nodes.The next step is constraint generation. Here, we go through the nodes and generate constraints of the form "this set of NodeSlots must have the same value for this field" or "this set of NodeSlots must have the exact same quantizer". We enqueue these constraints into a queue.
We then begin processing the queue. We pop each constraint and modify the provisional qspecs it affects to satisfy the constraint. This allows us to "relax" a qspec to fit a constraint (for example, a sigmoid and hardtanh that get concated can be relaxed to a (-1, 1) range, more general than sigmoid's (0, 1)), force a winning field (by taking the highest priority qspec's value for that field), merge groups of qspecs into a shared qspec, or error out if we can't reconcile the set of qspecs. When we change a provisional qspec, we go back to the Node it belongs to and re-enqueue any constraints that Node originally created. This means that processing is order-independent, and we will continue relaxing QSpecs until no new constraints are created.
Convergence here is guaranteed because each time we change a qspec, we lower the priority value of some field in the graph (0 is the "highest priority"). We only enqueue new constraints if we actually change a qspec, so there's a bound on how many changes we can make.
Finally, once all of the provisional qspecs are set, we write the actual qspecs back as annotations on the graph, and proceed with the rest of quantization.
Tested on existing unit tests plus the repro script for the YOLO bug. Also added a unit test for the YOLO structure