Tighten ModeloptBaseConfig mapping semantics

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Author: shengliangxu

modelopt/torch/opt/config.py

@@ -60,8 +60,15 @@ def ModeloptField(default: Any = PydanticUndefined, **kwargs):  # noqa: N802
 class ModeloptBaseConfig(BaseModel, MutableMapping[str, Any]):
     """Our config base class for mode configuration.
 
-    The base class extends the capabilities of pydantic's BaseModel to provide additional methods
-    and properties for easier access and manipulation of the configuration.
+    The base class extends pydantic's BaseModel with a mapping interface so schema-backed
+    config objects can keep the dict-style access patterns used by older ModelOpt code.
+
+    This is intentionally a fixed-key mutable mapping instead of a general dict. The mapping
+    keys are the model fields exposed through their aliases when present, and lookups accept
+    either a field name or its alias. Values are read and written through the pydantic model, so
+    assignment validation still applies. New keys cannot be inserted, and existing keys cannot
+    be deleted because the schema defines the complete key set; callers that need omission
+    semantics should use model_dump(exclude_unset=True) or the explicit_* helpers.
     """
 
     model_config = PyDanticConfigDict(extra="forbid", validate_assignment=True)
@@ -111,44 +118,42 @@ def __contains__(self, key: str) -> bool:
 
     def __getitem__(self, key: str) -> Any:
         """Get the value for the given key (can be name or alias of field)."""
-        return getattr(self, self.get_field_name_from_key(key))
+        try:
+            return getattr(self, self.get_field_name_from_key(key))
+        except AttributeError as e:
+            raise KeyError(key) from e
 
     def __setitem__(self, key: str, value: Any) -> None:
-        """Set the value for the given key (can be name or alias of field)."""
-        setattr(self, self.get_field_name_from_key(key), value)
+        """Set an existing field by name or alias, preserving pydantic assignment validation."""
+        try:
+            field_name = self.get_field_name_from_key(key)
+        except AttributeError as e:
+            raise KeyError(key) from e
+        if field_name not in type(self).model_fields:
+            raise KeyError(key)
+        setattr(self, field_name, value)
 
     def __delitem__(self, key: str) -> None:
-        """Unset the given key so exclude_unset dumps omit it."""
+        """Reject deletion because ModeloptBaseConfig exposes a fixed schema key set."""
         try:
-            field_name = self.get_field_name_from_key(key)
+            self.get_field_name_from_key(key)
         except AttributeError as e:
             raise KeyError(key) from e
-        if field_name in self._iterable_model_extra:
-            assert self.model_extra is not None
-            del self.model_extra[field_name]
-            self.model_fields_set.discard(field_name)
-            return
-
-        field_info = type(self).model_fields[field_name]
-        default = field_info.get_default(call_default_factory=True)
-        if default is PydanticUndefined:
-            raise KeyError(f"Key {key} cannot be unset because it has no default.")
-        self.__dict__[field_name] = default
-        self.model_fields_set.discard(field_name)
+        raise TypeError("Config mapping keys are fixed and cannot be deleted.")
 
     def get(self, key: str, default: Any = None) -> Any:
         """Get the value for the given key (can be name or alias) or default if not found."""
         try:
             return self[key]
-        except AttributeError:
+        except KeyError:
             return default
 
     def __len__(self) -> int:
-        """Return the length of the config."""
-        return len(self.model_fields) + len(self._iterable_model_extra)
+        """Return the number of schema and extra keys exposed by the mapping."""
+        return len(type(self).model_fields) + len(self._iterable_model_extra)
 
     def __iter__(self) -> Iterator[str]:
-        """Iterate over aliases (or name if alias is not defined) of fields."""
+        """Iterate over schema keys, preferring aliases over field names."""
         for field_name, field_info in type(self).model_fields.items():
             yield field_info.alias or field_name
         yield from self._iterable_model_extra
@@ -157,6 +162,29 @@ def _get_kv_dict(self) -> dict[str, Any]:
         """Return a dictionary with keys as aliases if possible."""
         return {k: self[k] for k in self}
 
+    def iter_explicit_keys(self) -> Iterator[str]:
+        """Iterate over explicitly set schema keys, preferring aliases over field names."""
+        for field_name, field_info in type(self).model_fields.items():
+            if field_name in self.model_fields_set:
+                yield field_info.alias or field_name
+        yield from self._iterable_model_extra
+
+    def _get_explicit_kv_dict(self) -> dict[str, Any]:
+        """Return explicitly set key-value pairs with keys as aliases if possible."""
+        return {k: self[k] for k in self.iter_explicit_keys()}
+
+    def explicit_keys(self) -> KeysView[str]:
+        """Return the explicitly set keys of the config."""
+        return self._get_explicit_kv_dict().keys()
+
+    def explicit_values(self) -> ValuesView[Any]:
+        """Return the explicitly set values of the config."""
+        return self._get_explicit_kv_dict().values()
+
+    def explicit_items(self) -> ItemsView[str, Any]:
+        """Return the explicitly set items of the config with keys as aliases if possible."""
+        return self._get_explicit_kv_dict().items()
+
     def keys(self) -> KeysView[str]:
         """Return the keys (aliases prioritized over names) of the config."""
         return self._get_kv_dict().keys()

modelopt/torch/quantization/config.py

@@ -543,8 +543,8 @@ class QuantizerCfgEntry(ModeloptBaseConfig):
             "Attributes to apply to matched quantizers. A list configures a sequential quantizer."
         ),
     )
-    enable: bool | None = ModeloptField(
-        default=None,
+    enable: bool = ModeloptField(
+        default=True,
         title="Quantizer enable flag.",
         description="Optional on/off toggle for matched quantizers, independent of cfg.",
     )
@@ -556,11 +556,6 @@ def validate_quantizer_cfg_entry(cls, values):
         if not isinstance(values, Mapping):
             return values
 
-        if "cfg" not in values and "enable" not in values:
-            raise ValueError(
-                "Each quant_cfg entry must specify 'cfg', 'enable', or both. "
-                "An entry with only 'quantizer_name' has no effect."
-            )
         if "cfg" in values and values["cfg"] is None:
             raise ValueError("cfg must be omitted or a valid mapping/list, not null.")
         if "enable" in values and values["enable"] is None:
@@ -1038,25 +1033,21 @@ def normalize_quant_cfg_list(
     - Legacy ``nn.*``-scoped format: ``{"nn.<Class>": {"<quantizer_name>": <cfg>}}`` - converted
       to a new-format entry with ``parent_class`` set.
 
-    **Validation** - an entry is rejected if it carries no instruction, i.e. it specifies neither
-    ``cfg`` nor ``enable``.  Concretely, the following are invalid:
+    **Validation** - an entry is rejected if its shape is invalid.  Concretely, the following
+    are invalid:
 
     - An empty entry ``{}``.
-    - An entry with only ``quantizer_name`` and no other keys - the only effect would be an
-      implicit ``enable=True``, which must be stated explicitly.
     - An entry with ``enable=True`` (explicit or implicit) whose ``cfg`` is not a non-empty
       ``dict`` or ``list`` - e.g. ``{"quantizer_name": "*", "cfg": {}}`` or
       ``{"quantizer_name": "*", "cfg": 42}``.  An enabled quantizer must have a valid
       configuration.
 
-    **Normalization** - after conversion and validation every entry is put into canonical form:
-
-    - ``enable`` is set to ``True`` if not explicitly specified.
-    - ``cfg`` is set to ``None`` if not present in the entry.
-
-    For dict and legacy inputs, every returned entry is guaranteed to have
-    ``quantizer_name``, ``enable``, and ``cfg`` set (plus optionally ``parent_class``). Typed
-    :class:`QuantizerCfgEntry` inputs are assumed to be already parsed and are preserved.
+    **Normalization** - after conversion and validation every entry is parsed as a
+    :class:`QuantizerCfgEntry`. Schema defaults are available through mapping access, so ``enable``
+    defaults to ``True`` and ``cfg`` defaults to ``None`` when omitted. Omitted defaults are not
+    marked as explicitly set, so ``model_dump(exclude_unset=True)`` preserves the user's sparse
+    input shape. Typed :class:`QuantizerCfgEntry` inputs are assumed to be already parsed and are
+    preserved.
 
     Args:
         v: A list of raw quant_cfg entries in any supported format, or a legacy flat dict.
@@ -1066,9 +1057,8 @@ def normalize_quant_cfg_list(
         typed entries are preserved.
 
     Raises:
-        ValueError: If any entry has only ``quantizer_name`` with neither ``cfg`` nor ``enable``,
-            if ``enable=True`` with an empty or non-dict/list ``cfg``, or if the entry format
-            is not recognized.
+        ValueError: If ``enable=True`` with an empty or non-dict/list ``cfg``, or if the entry
+            format is not recognized.
     """
     if isinstance(v, list) and all(isinstance(raw, QuantizerCfgEntry) for raw in v):
         return cast("list[QuantizerCfgEntry]", v)
@@ -1164,14 +1154,6 @@ def _dict_to_entry(key: str, value: Any) -> list[dict[str, Any]]:
             raise ValueError(f"Invalid quant_cfg entry: {raw!r}.")
 
         for entry in entries:
-            # Validate: must carry at least one instruction beyond the path selector.
-            if "cfg" not in entry and "enable" not in entry:
-                raise ValueError(
-                    f"Invalid quant_cfg entry: {raw!r} - each entry must specify 'cfg', 'enable', "
-                    "or both. An entry with only 'quantizer_name' has no effect (implicit "
-                    "enable=True is not allowed; set it explicitly)."
-                )
-
             # Validate: when cfg is present and enable=True, cfg must be a non-empty
             # dict or list.  An empty cfg would attempt to create a
             # QuantizerAttributeConfig with no actual configuration.
@@ -1211,9 +1193,6 @@ def _dict_to_entry(key: str, value: Any) -> list[dict[str, Any]]:
                         "explicitly."
                     )
 
-            # Normalize: make enable explicit. cfg remains omitted when it is intentionally unset.
-            entry.setdefault("enable", True)
-
             result.append(QuantizerCfgEntry.model_validate(entry))
     return result
 
@@ -1981,22 +1960,20 @@ def _cfg_to_dict(cfg):
     for entry in quant_cfg:
         name = entry["quantizer_name"]
         raw_cfg = entry.get("cfg")
-        enable = entry.get("enable")
+        enable = entry["enable"]
         if "weight_quantizer" in name:
             # We don't calibrate weight quantizer
             continue
         # Sequential quantizers (e.g. W4A8) have a list of cfg dicts
         if isinstance(raw_cfg, list):
             for _config in raw_cfg:
                 cfg = _cfg_to_dict(_config)
-                if enable is not None:
-                    cfg["enable"] = enable
+                cfg["enable"] = enable
                 if _not_dynamic(cfg):
                     return True
             continue
         cfg = _cfg_to_dict(raw_cfg)
-        if enable is not None:
-            cfg["enable"] = enable
+        cfg["enable"] = enable
         if _not_dynamic(cfg):
             return True
 

modelopt/torch/quantization/conversion.py

@@ -254,7 +254,7 @@ def set_quantizer_by_cfg(quant_model: nn.Module, quant_cfg: QuantizeQuantCfgInpu
     for entry in quant_cfg:
         quantizer_name: str = entry["quantizer_name"]
         cfg = entry["cfg"]  # None, QuantizerAttributeConfig, or list after normalization
-        enable = entry["enable"] if entry["enable"] is not None else True
+        enable = entry["enable"]
         parent_class_name = entry.get("parent_class")
         if parent_class_name:
             try:

tests/unit/torch/opt/test_config.py

@@ -72,7 +72,7 @@ def _run_test(is_new_registered):
         assert config[lin_name] == lin_expected_value
         assert config[lin_alias] == lin_expected_value
         assert getattr(config, lin_name) == lin_expected_value
-        with nullcontext() if is_new_registered else pytest.raises(AttributeError):
+        with nullcontext() if is_new_registered else pytest.raises(KeyError):
             config[new_name]
 
         # get

tests/unit/torch/quantization/test_config_validation.py

@@ -59,6 +59,7 @@ def test_need_calibration():
 
 def test_need_calibration_with_quantize_config_type():
     """need_calibration accepts schema-backed QuantizeConfig objects."""
+    assert need_calibration(QuantizeConfig())
     assert need_calibration(QuantizeConfig.model_validate(FP8_DEFAULT_CFG))
     assert not need_calibration(QuantizeConfig.model_validate(FP8_PER_CHANNEL_PER_TOKEN_CFG))
 
@@ -100,16 +101,27 @@ def test_quantizer_cfg_entry_is_pydantic_and_dict_like():
     entry = QuantizerCfgEntry(quantizer_name="*", enable=False)
     assert isinstance(entry, ModeloptBaseConfig)
     assert entry["quantizer_name"] == "*"
-    assert entry.get("cfg") is None
+    assert entry["cfg"] is None
+    assert "cfg" in entry
+    assert list(entry) == ["quantizer_name", "parent_class", "cfg", "enable"]
+    assert dict(entry.items()) == {
+        "quantizer_name": "*",
+        "parent_class": None,
+        "cfg": None,
+        "enable": False,
+    }
+    assert dict(entry.explicit_items()) == {"quantizer_name": "*", "enable": False}
+    with pytest.raises(KeyError):
+        entry["unknown"] = 1
     assert entry.model_dump(exclude_unset=True) == {"quantizer_name": "*", "enable": False}
 
     cfg_entry = QuantizerCfgEntry(quantizer_name="*weight_quantizer", cfg={"num_bits": 8})
     assert isinstance(cfg_entry["cfg"], QuantizerAttributeConfig)
     assert _cfg_to_dict(cfg_entry["cfg"]) == {"num_bits": 8}
 
 
-def test_quantizer_cfg_entry_mutable_mapping_delitem_unsets_field():
-    """Deleting a config key resets it to unset for exclude_unset dumps."""
+def test_quantizer_cfg_entry_mutable_mapping_rejects_key_deletion():
+    """ModeloptBaseConfig mappings have a fixed key set and reject deletion."""
     entry = QuantizerCfgEntry(quantizer_name="*weight_quantizer", cfg={"num_bits": 8}, enable=True)
     assert isinstance(entry, MutableMapping)
     assert entry.model_dump(exclude_unset=True) == {
@@ -118,11 +130,14 @@ def test_quantizer_cfg_entry_mutable_mapping_delitem_unsets_field():
         "enable": True,
     }
 
-    del entry["cfg"]
+    with pytest.raises(TypeError):
+        del entry["cfg"]
 
-    assert entry["cfg"] is None
+    assert "cfg" in entry
+    assert entry["cfg"] is not None
     assert entry.model_dump(exclude_unset=True) == {
         "quantizer_name": "*weight_quantizer",
+        "cfg": {"num_bits": 8},
         "enable": True,
     }
 
@@ -196,10 +211,13 @@ def test_quantizer_cfg_entry_rejects_explicit_null_values(raw, match):
         normalize_quant_cfg_list([raw])
 
 
-def test_quantizer_cfg_entry_rejects_no_effect_entry():
-    """Direct QuantizerCfgEntry construction rejects entries with no cfg or enable."""
-    with pytest.raises(ValidationError, match="must specify 'cfg', 'enable'"):
-        QuantizerCfgEntry(quantizer_name="*")
+def test_quantizer_cfg_entry_defaults_enable_true():
+    """Direct QuantizerCfgEntry construction uses enable=True when omitted."""
+    entry = QuantizerCfgEntry(quantizer_name="*")
+    assert entry["enable"] is True
+    assert entry["cfg"] is None
+    assert dict(entry.explicit_items()) == {"quantizer_name": "*"}
+    assert entry.model_dump(exclude_unset=True) == {"quantizer_name": "*"}
 
 
 def test_quantizer_cfg_entry_rejects_empty_name():
@@ -231,7 +249,8 @@ def test_new_format_passthrough(self):
         assert result[0]["quantizer_name"] == "*weight_quantizer"
         assert isinstance(result[0]["cfg"], QuantizerAttributeConfig)
         assert _cfg_to_dict(result[0]["cfg"]) == {"num_bits": 8, "axis": 0}
-        assert result[0]["enable"] is True  # defaulted
+        assert result[0]["enable"] is True  # schema default
+        assert "enable" not in dict(result[0].explicit_items())
 
     def test_typed_entry_list_passthrough(self):
         """Already-parsed QuantizerCfgEntry lists are returned unchanged."""
@@ -256,6 +275,7 @@ def test_mixed_typed_and_dict_entries_normalize_to_typed_entries(self):
         assert isinstance(result[1], QuantizerCfgEntry)
         assert _cfg_to_dict(result[1]["cfg"]) == {"num_bits": 8}
         assert result[1]["enable"] is True
+        assert "enable" not in dict(result[1].explicit_items())
 
     def test_new_format_enable_false(self):
         """Explicit enable=False is preserved."""
@@ -277,7 +297,8 @@ def test_legacy_single_key_dict(self):
         result = normalize_quant_cfg_list(raw)
         assert result[0]["quantizer_name"] == "*weight_quantizer"
         assert _cfg_to_dict(result[0]["cfg"]) == {"num_bits": 8, "axis": 0}
-        assert result[0]["enable"] is True  # defaulted
+        assert result[0]["enable"] is True  # schema default
+        assert "enable" not in dict(result[0].explicit_items())
 
     def test_legacy_single_key_dict_with_enable(self):
         """Legacy {'*path': {'enable': False}} splits enable out from cfg."""
@@ -296,17 +317,19 @@ def test_legacy_nn_class_scoped(self):
         assert result[0]["enable"] is False
 
     def test_normalization_cfg_defaults_to_none(self):
-        """Entries without cfg get cfg=None after normalization."""
+        """Entries without cfg expose the default mapping key but keep it unset."""
         raw = [{"quantizer_name": "*lm_head*", "enable": False}]
         result = normalize_quant_cfg_list(raw)
         assert "cfg" in result[0]
         assert result[0]["cfg"] is None
+        assert "cfg" not in dict(result[0].explicit_items())
 
     def test_normalization_enable_defaults_to_true(self):
-        """Entries with cfg but no enable get enable=True after normalization."""
+        """Entries with cfg but no enable read as enable=True without marking it explicit."""
         raw = [{"quantizer_name": "*", "cfg": {"num_bits": 4}}]
         result = normalize_quant_cfg_list(raw)
         assert result[0]["enable"] is True
+        assert "enable" not in dict(result[0].explicit_items())
 
     def test_empty_list(self):
         """Empty list is returned unchanged."""
@@ -322,10 +345,13 @@ def test_multiple_entries_order_preserved(self):
         assert result[0]["quantizer_name"] == "*"
         assert result[1]["quantizer_name"] == "*weight_quantizer"
 
-    def test_error_on_quantizer_name_only(self):
-        """Entry with only quantizer_name and no cfg or enable is rejected."""
-        with pytest.raises(ValueError, match="must specify 'cfg', 'enable'"):
-            normalize_quant_cfg_list([{"quantizer_name": "*"}])
+    def test_quantizer_name_only_defaults_enable_true(self):
+        """Entry with only quantizer_name uses enable=True from the schema default."""
+        result = normalize_quant_cfg_list([{"quantizer_name": "*"}])
+        assert result[0]["enable"] is True
+        assert result[0]["cfg"] is None
+        assert dict(result[0].explicit_items()) == {"quantizer_name": "*"}
+        assert result[0].model_dump(exclude_unset=True) == {"quantizer_name": "*"}
 
     def test_error_on_empty_dict(self):
         """An empty dict entry is rejected."""
@@ -421,6 +447,7 @@ def test_legacy_flat_dict_conversion(self):
         assert result[1]["quantizer_name"] == "*weight_quantizer"
         assert _cfg_to_dict(result[1]["cfg"]) == {"num_bits": 8, "axis": 0}
         assert result[1]["enable"] is True
+        assert "enable" not in dict(result[1].explicit_items())
 
     def test_legacy_enable_only_produces_cfg_none(self):
         """Legacy {'*': {'enable': False}} should produce cfg=None, not cfg={}."""