[Bug]: EagerQuantizer maps root module weight fake-quantizers to "parametrizations", silently dropping QATSchedule and breaking training transitions

[rohith500]

Describe the Bug

In EagerQuantizer._get_fake_quantize_modules() (src/coreai_opt/quantization/_eager/quantizer.py:363-374), weight fake-quantize modules are mapped back to their owning module from PyTorch's ParametrizationList using a right-split on dots:

mapping: dict[str, list[FakeQuantizeImplBase]] = defaultdict(list)
for name, module in self._model.named_modules():
    fq_list = [c for c in module.children() if isinstance(c, FakeQuantizeImplBase)]
    if fq_list:
        # Weight FQs live inside ParametrizationList; map back
        # to the owning module (strip ".parametrizations.<param>")
        if isinstance(module, ParametrizationList):
            key = name.rsplit(".", 2)[0]
        else:
            key = name
        mapping[key] += fq_list
return dict(mapping)

The string decomposition logic name.rsplit(".", 2)[0] assumes that the qualified module path always contains at least two dots (i.e. <owning_module>.parametrizations.<tensor_name>).

When quantizing a standalone root module (e.g., model = nn.Linear(4, 4), single-layer wrappers, or custom modules acting as the root), the parametrization module's qualified name in named_modules() is "parametrizations.weight":

• "parametrizations.weight" contains only one dot.
• Calling "parametrizations.weight".rsplit(".", 2) returns ["parametrizations", "weight"].
• Index [0] evaluates to "parametrizations" instead of "" (the root module identifier).

---

Downstream Architectural Failure Analysis

The incorrect mapping key triggers a cascade of silent failures during Quantization-Aware Training (QAT):

1. Config Hierarchy Mismatch in _build_fq_to_schedule:
   In Quantizer.prepare(), after fake-quantize modules are constructed, _build_fq_to_schedule() maps each module's fake-quantizers to their corresponding QATSchedule:

   for module_name, fq_list in self._module_to_fqs.items():
       schedule = self._resolve_schedule(module_name)
       if schedule is not None:
           for fq_mod in fq_list:
               self._fq_to_schedule[fq_mod] = schedule

   _resolve_schedule(module_name) inspects self._module_config_dict across MODULE_NAME, MODULE_TYPE, and GLOBAL priority levels. The root module's configuration is stored under key "". When it queries for "parametrizations", no matching module exists, and _resolve_schedule("parametrizations") returns None.

2. Silent Drop from _fq_to_schedule:
   Because _resolve_schedule returns None, the weight fake-quantizer is never registered in self._fq_to_schedule.

3. Broken QAT Training Dynamics:
   During QAT, quantizer.step() is invoked inside with quantizer.training_mode(): to advance the schedule:

   • Weight-only quantization: self._fq_to_schedule remains completely empty. Calling quantizer.step() outputs:

     UserWarning: step() called but no qat_schedule is configured on any module. Observer and fake-quant states will not be updated.
     

     even though the user explicitly configured a valid QATSchedule.
   • Joint weight and activation quantization: Activation fake-quantizers (which are child modules directly attached to the root module, mapped under "") receive their schedule, but the weight fake-quantizer does not. As training steps advance:
     • Activation observers freeze and activate fake quantization on schedule.
     • Weight observers never freeze and fake quantization never transitions.
     • The training loop silently runs with corrupted quantization dynamics, resulting in degraded post-training numerical accuracy and inference mismatch.

---

Step-by-Step Reproduction

import torch
import torch.nn as nn
from coreai_opt.quantization.config import ExecutionMode, ModuleQuantizerConfig, QATSchedule
from coreai_opt.quantization.quantizer import Quantizer, QuantizerConfig
from coreai_opt.quantization.spec import default_weight_quantization_spec

# 1. Define a standalone root module
model = nn.Linear(4, 4)
sched = QATSchedule(enable_observer=0, enable_fake_quant=5)
config = QuantizerConfig(
    global_config=ModuleQuantizerConfig(
        op_state_spec={"weight": default_weight_quantization_spec()},
        qat_schedule=sched,
    ),
    execution_mode=ExecutionMode.EAGER,
)

# 2. Prepare the model
quantizer = Quantizer(model, config)
prepared_model = quantizer.prepare(torch.randn(1, 4))

# 3. Inspect internal FQ mapping
fq_map = quantizer._get_fake_quantize_modules()
print("Keys in fq_map:", list(fq_map.keys()))

# 4. Check schedule registration
print("Registered in _fq_to_schedule:", len(quantizer._fq_to_schedule))
for fq, schedule in quantizer._fq_to_schedule.items():
    print(f" - Scheduled FQ type: {type(fq.qparams_calculator).__name__}")

Actual Output:

Keys in fq_map: ['', 'parametrizations']
Registered in _fq_to_schedule: 2
 - Scheduled FQ type: MovingAverageQParamsCalculator (input activation)
 - Scheduled FQ type: MovingAverageQParamsCalculator (output activation)
# The weight fake quantizer (StaticQParamsCalculator) is completely missing!

Expected Output:

Keys in fq_map: ['']
Registered in _fq_to_schedule: 3
 - Scheduled FQ type: MovingAverageQParamsCalculator (input activation)
 - Scheduled FQ type: MovingAverageQParamsCalculator (output activation)
 - Scheduled FQ type: StaticQParamsCalculator (weight)

---

Proposed Permanent Fix

In src/coreai_opt/quantization/_eager/quantizer.py:369-373, replace the dot-counting heuristic with exact namespace boundary extraction:

# Before
if isinstance(module, ParametrizationList):
    key = name.rsplit(".", 2)[0]
else:
    key = name

# After
if isinstance(module, ParametrizationList):
    key = name.rsplit(".parametrizations.", 1)[0] if ".parametrizations." in name else ""
else:
    key = name

Why This Fix Is Canonical:

• Nested submodules ("sub.linear.parametrizations.weight"):
  ".parametrizations." in name is True $\rightarrow$ splits on the exact parametrization namespace, returning "sub.linear".
• Root modules ("parametrizations.weight"):
  ".parametrizations." in name is False $\rightarrow$ evaluates directly to "" (the canonical root module identifier).
• Deeply nested hierarchies ("a.b.c.parametrizations.bias"):
  Splits at the rightmost .parametrizations., returning "a.b.c".

---

Verification Plan

1. Unit Test: Add test_root_module_fake_quantize_mapping_eager in tests/quantization/test_qat_schedule.py asserting:
   • _get_fake_quantize_modules() produces only key "" for root modules.
   • All fake-quantizers (weights + activations) are present in _fq_to_schedule.
   • quantizer.step() transitions weight fake-quantization states as specified by QATSchedule.
2. Quality Gates: Ensure make check and the full quantization test suite (pytest tests/quantization/ -m "not slow" -n auto) pass cleanly.