Shared/tied weight gets its dtype and QAT schedule from different module configs (graph mode)

[eyupcanakman]

Description

When two named modules share a weight (weight tying), graph-mode Quantizer.prepare merges it into one fake-quantize module. That module picks its dtype and its QAT schedule through two different orderings. With different per-module configs the two can disagree. The weight quantizes with one module's dtype but follows the other module's training schedule, and nothing warns.

Reproduction

Two nn.Linear layers sharing one weight, l1 declared first (int8, fake-quant enabled at step 1), l2 declared last (int4, fake-quant enabled at step 5):

class Shared(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(10, 10, bias=False)
        self.l2 = nn.Linear(10, 10, bias=False)
        self.l2.weight = self.l1.weight          # tied
    def forward(self, x):
        return self.l2(self.l1(x))

cfg = QuantizerConfig(global_config=None, module_name_configs={
    "l1": ModuleQuantizerConfig(op_state_spec={"weight": wspec(torch.int8)},
                                qat_schedule=QATSchedule(enable_observer=0, enable_fake_quant=1)),
    "l2": ModuleQuantizerConfig(op_state_spec={"weight": wspec(torch.int4)},
                                qat_schedule=QATSchedule(enable_observer=0, enable_fake_quant=5)),
}, execution_mode="graph")
q = Quantizer(Shared().eval(), cfg)
q.prepare((torch.randn(1, 10),))

The single shared weight FQ ends up int4 (from l2) but enabled at step 1 (from l1):

distinct weight FQ objects: 1
weight FQ dtype quant_min/quant_max: -8 / 7  -> int4     # l2, declared last
weight FQ governing schedule enable_fake_quant step: 1   # l1, first in graph
step=1 fake_quant_enabled=1    # fires at step 1, not l2's step 5

Root cause

The two owners are resolved independently.

• Schedule owner comes from _get_fake_quantize_modules (_graph/quantizer.py), which walks the FX graph and maps each shared FQ to the first consumer with nn_module_stack, so graph order picks l1.
• Dtype owner comes from _get_state_node_shared_spec (_graph/_annotation_utils.py), which keeps the spec of the user annotated first in priority order, and priority follows config declaration order, so l2 wins.

Eager mode has the same split. It warns about the schedule conflict, but the warning only mentions the schedule, not the dtype. tests/quantization/test_qat_schedule.py::test_shared_weight_keeps_first_schedule covers a shared weight with per-module schedules but gives both modules the same dtype, so the conflict is never covered.

The draft tie-break changes in #40 touch the dtype side but not the schedule path, so this can remain after #40 lands.

Suggested fix

The schedule path was never updated when the dtype owner moved to the #17 priority mechanism. Aligning it to the same owner the spec already picks keeps the shared FQ consistent, matching how the dtype already keeps one winner. Or would you rather reject a conflicting tied-weight config?

Environment

• coreai-optimization at 3e01e4a
• torch 2.8.0, Python 3.12, macOS (Apple silicon)