fix(quantization): map root module weight fake quantizers to root key (#100)

[rohith500]

Summary of Changes

Closes #100.

This PR fixes a critical defect in EagerQuantizer where fake-quantize modules for root module weights are mapped to "parametrizations" instead of the canonical root module identifier "", causing weight fake-quantizers to silently lose their configured QATSchedule.

1. Exact Namespace Boundary Extraction in _get_fake_quantize_modules:
   • In src/coreai_opt/quantization/_eager/quantizer.py:

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

   • Replaced the brittle .rsplit(".", 2)[0] heuristic with exact .parametrizations. namespace detection.
   • For root modules (name == "parametrizations.weight"), .rsplit(".", 2)[0] previously returned "parametrizations" because there was only one dot. Now, ".parametrizations." in name is False, cleanly returning "".
   • For nested submodules (name == "layer.parametrizations.weight"), splits at .parametrizations., returning "layer".
   • For deeply nested modules (name == "encoder.layers.0.parametrizations.weight"), splits at the rightmost .parametrizations., returning "encoder.layers.0".

---

Motivation & Impact

In Eager QAT (ExecutionMode.EAGER):

• quantizer.prepare() registers schedules via:

  for mod_name, fq_list in self._get_fake_quantize_modules().items():
      schedule = self._module_config_dict.get(mod_name, ...).qat_schedule
      if schedule is not None:
          for fq in fq_list:
              self._fq_to_schedule[fq] = schedule

• For root modules, the module config is stored under "". Because weight FQs were mapped to "parametrizations", _module_config_dict.get("parametrizations") returned None, and weight fake-quantizers were never added to _fq_to_schedule.
• During QAT training (quantizer.step()):
  • Weight-only quantization: _fq_to_schedule remained completely empty, emitting UserWarning: step() called but no qat_schedule is configured on any module.
  • Joint weight & activation quantization: Activation fake-quantizers transitioned on schedule, while the weight fake-quantizer's observer never froze and fake-quantization never activated, silently corrupting training dynamics.

---

Automated Tests Added

Added to tests/quantization/test_qat_schedule.py:

• test_root_module_fake_quantize_mapping_and_schedule:
  • Parametrized across both ExecutionMode.EAGER and ExecutionMode.GRAPH.
  • Verifies _get_fake_quantize_modules() produces only the canonical root key "".
  • Verifies "parametrizations" is not present in fq_map.keys().
  • Verifies all fake-quantizers (weights + activations) are registered in _fq_to_schedule.
  • Steps through quantizer.training_mode() and verifies observer and fake-quantization transitions at each step (step 0: observer on, FQ off $\rightarrow$ step 2: FQ on).

---

Verification

• pytest tests/quantization/test_qat_schedule.py -k "test_root_module_fake_quantize_mapping_and_schedule": 2 passed in 4.72s.
• pytest tests/quantization/test_qat_schedule.py: 50 passed, 1 xfailed (known upstream xfail) in 19.74s.
• pytest tests/quantization/ -m "not slow" -n auto: 1,676 passed, 240 skipped, 37 xfailed in 149.63s.
• make check: All checks passed (26/26):
  • License header verification: Passed
  • Formatting & line-length (ruff, darker $\le$ 100): Passed
  • Linters (ruff-lint, checkmake, pymarkdown 0.9.40): Passed
  • Project rules (__all__, import aliasing): Passed
  • Towncrier changelog validation: Passed

---

Checklist

• Follows Conventional Commits formatting (fix(quantization): map root module weight FQs to root key for eager QAT scheduling (#100))
• Commit is cryptographically signed with SSH key (SHA256:0wr/nMrZodwxsOOHBGpcwRXPk8c3rePKishJLzmhj98)
• Includes Towncrier fragment (changelog.d/100.fixed)
• Passes make check without warnings or diff pollution
• Includes automated tests reproducing the bug and proving the fix

[u-simha]

couple small changes requested; thanks for opening up the PR.

[rohith500]

Hi @u-simha, thank you for the review!

I've addressed both feedback points in commit c057094:

1. Extracted Utility: Added extract_name_from_parameterization (along with extract_name_from_parametrization alias) in src/coreai_opt/_utils/torch_utils.py to robustly resolve module names from parametrization FQNs (supporting submodule, nested submodule, root module, and pass-through fallback), and added comprehensive unit tests in tests/test_utils/test_torch_utils.py.
2. Scoped Test: Simplified the test in tests/quantization/test_qat_schedule.py to test_root_module_fake_quantize_mapping, removing the schedule stepping loop to focus exclusively on verifying _get_fake_quantize_modules behavior.

All 26 make check gates and the full test suite pass cleanly.

[u-simha]

Thanks for the fix!

[u-simha]

I think we can just use the PyTorch standard spelling for the function extract_name_from_parametrization and skip the aliasing.

[rohith500]

Updated! Renamed the function directly to extract_name_from_parametrization across the codebase and removed the alias in commit 7477f88.