forbid extra on models, add tests

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -393,7 +393,7 @@ def compress_model(self, model: Module):
 
             if prefix in module_to_scheme or prefix in sparse_compression_targets:
                 module_device = get_execution_device(module).type
-                is_meta = (module_device == "meta")
+                is_meta = module_device == "meta"
 
                 exec_device = "meta" if is_meta else "cpu"
                 onloading_device = "meta" if is_meta else module_device

src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py

@@ -178,9 +178,13 @@ def sparse24_bitmask_compress(
 
     if tensor.is_meta:
         num_rows, num_cols = tensor.shape
-        compressed_values = torch.empty((num_rows, num_cols // 2), dtype=tensor.dtype, device="meta")
+        compressed_values = torch.empty(
+            (num_rows, num_cols // 2), dtype=tensor.dtype, device="meta"
+        )
         packed_cols = (num_cols + 7) // 8
-        bitmasks_packed = torch.empty((num_rows, packed_cols), dtype=torch.uint8, device="meta")
+        bitmasks_packed = torch.empty(
+            (num_rows, packed_cols), dtype=torch.uint8, device="meta"
+        )
         return compressed_values, bitmasks_packed
 
     bytemasks = get_24_bytemasks(tensor=tensor)

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -189,7 +189,12 @@ def _initialize_scale_zero_point(
     else:
         # TODO: consider erroring out in the future as if the dtype if not one of these,
         # there is likely bug
-        if scale_dtype not in [torch.float16, torch.bfloat16, torch.float32, torch.float64]:
+        if scale_dtype not in [
+            torch.float16,
+            torch.bfloat16,
+            torch.float32,
+            torch.float64,
+        ]:
             scale_dtype = torch.float16
         zp_dtype = quantization_args.pytorch_dtype()
 

src/compressed_tensors/quantization/quant_args.py

@@ -19,7 +19,7 @@
 import torch
 from compressed_tensors.utils import Aliasable
 from compressed_tensors.utils.helpers import deprecated
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 
 __all__ = [
@@ -186,6 +186,7 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
             "Observers constructor excluding quantization range or symmetry"
         ),
     )
+    model_config = ConfigDict(extra="forbid")
 
     @field_validator("type", mode="before")
     def validate_type(cls, value) -> QuantizationType:

src/compressed_tensors/quantization/quant_config.py

@@ -26,7 +26,7 @@
     module_type,
     parse_out_kv_cache_args,
 )
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 from torch.nn import Module
 
 
@@ -142,6 +142,7 @@ class QuantizationConfig(BaseModel):
     quantization_status: QuantizationStatus = QuantizationStatus.INITIALIZED
     global_compression_ratio: Optional[float] = None
     ignore: Optional[List[str]] = Field(default_factory=list)
+    model_config = ConfigDict(extra="forbid")
 
     def model_post_init(self, __context):
         """

src/compressed_tensors/quantization/quant_scheme.py

@@ -13,15 +13,15 @@
 # limitations under the License.
 
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import List, Optional
 
 from compressed_tensors.quantization.quant_args import (
     DynamicType,
     QuantizationArgs,
     QuantizationStrategy,
     QuantizationType,
 )
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, ConfigDict, model_validator
 
 
 __all__ = [
@@ -47,6 +47,7 @@ class QuantizationScheme(BaseModel):
     weights: Optional[QuantizationArgs] = None
     input_activations: Optional[QuantizationArgs] = None
     output_activations: Optional[QuantizationArgs] = None
+    model_config = ConfigDict(extra="forbid")
 
     @model_validator(mode="after")
     def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":

src/compressed_tensors/transform/transform_args.py

@@ -15,7 +15,7 @@
 from enum import Enum
 from typing import List
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 
 __all__ = ["TransformArgs", "TransformLocation"]
@@ -61,6 +61,7 @@ class TransformArgs(BaseModel):
     location: TransformLocation
     inverse: bool = Field(default=False)
     ignore: List[str] = Field(default_factory=list)
+    model_config = ConfigDict(extra="forbid")
 
     @field_validator("targets", "ignore", mode="before")
     @classmethod

src/compressed_tensors/transform/transform_config.py

@@ -15,7 +15,7 @@
 from typing import Dict
 
 from compressed_tensors.transform import TransformArgs, TransformScheme
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 __all__ = ["TransformConfig"]
@@ -31,6 +31,7 @@ class TransformConfig(BaseModel):
     """
 
     config_groups: Dict[str, TransformScheme]
+    model_config = ConfigDict(extra="forbid")
 
 
 # quip / quip sharp

src/compressed_tensors/transform/transform_scheme.py

@@ -15,7 +15,7 @@
 from typing import List
 
 from compressed_tensors.transform import TransformArgs
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 
 __all__ = ["TransformScheme"]
@@ -40,3 +40,4 @@ class TransformScheme(BaseModel):
     apply: List[TransformArgs] = Field(default_factory=list)
     randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)
+    model_config = ConfigDict(extra="forbid")

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -446,10 +446,7 @@ def test_compress_model_meta(model_stub, q_format, s_config):
         cpu_model, s_config, q_format
     )
     # Only stores dtype because meta model does not store values
-    expected = {
-        k: v.dtype
-        for k, v in reference_compressor.compress(cpu_model).items()
-    }
+    expected = {k: v.dtype for k, v in reference_compressor.compress(cpu_model).items()}
 
     # Load model on meta device
     meta_model = AutoModelForCausalLM.from_pretrained(