Skip to content

Commit

Permalink
add model validator to quantization scheme
Browse files Browse the repository at this point in the history
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
  • Loading branch information
kylesayrs committed Nov 19, 2024
1 parent ff121cc commit dfc531e
Showing 1 changed file with 22 additions and 3 deletions.
25 changes: 22 additions & 3 deletions src/compressed_tensors/quantization/quant_scheme.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@
# limitations under the License.

from copy import deepcopy
from typing import List, Optional
from typing import Any, Dict, List, Optional

from compressed_tensors.quantization.quant_args import (
QuantizationArgs,
QuantizationStrategy,
QuantizationType,
)
from pydantic import BaseModel
from pydantic import BaseModel, model_validator


__all__ = [
Expand All @@ -47,12 +47,31 @@ class QuantizationScheme(BaseModel):
input_activations: Optional[QuantizationArgs] = None
output_activations: Optional[QuantizationArgs] = None

@model_validator(mode="after")
def validate_model_after(model: "QuantizationArgs") -> Dict[str, Any]:
inputs = model.input_activations
weights = model.weights
outputs = model.output_activations

if inputs is not None:
if inputs.actorder is not None:
raise ValueError("Cannot apply actorder to input activations")

if weights is not None:
if weights.dynamic:
raise ValueError("Cannot apply dynamic quantization to weights")

if outputs is not None:
if outputs.actorder is not None:
raise ValueError("Cannot apply actorder to output activations")

return model

@classmethod
def default_scheme(
cls,
targets: Optional[List[str]] = None,
):

if targets is None:
# default to quantizing all Linear layers
targets = ["Linear"]
Expand Down

0 comments on commit dfc531e

Please sign in to comment.