Skip to content

Commit d8a717c

Browse files
authored
Support bool arguments for actorder (#150)
1 parent 0f5823c commit d8a717c

File tree

2 files changed

+8
-4
lines changed

2 files changed

+8
-4
lines changed

src/compressed_tensors/quantization/quant_args.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
9393
strategy: Optional[QuantizationStrategy] = None
9494
block_structure: Optional[str] = None
9595
dynamic: bool = False
96-
actorder: Optional[ActivationOrdering] = None
96+
actorder: Union[ActivationOrdering, bool, None] = None
9797
observer: str = Field(
9898
default="minmax",
9999
description=(
@@ -151,6 +151,9 @@ def validate_strategy(cls, value) -> Union[QuantizationStrategy, None]:
151151

152152
@field_validator("actorder", mode="before")
153153
def validate_actorder(cls, value) -> Optional[ActivationOrdering]:
154+
if isinstance(value, bool):
155+
return ActivationOrdering.GROUP if value else None
156+
154157
if isinstance(value, str):
155158
return ActivationOrdering(value.lower())
156159

tests/test_quantization/test_quant_args.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -82,11 +82,12 @@ def test_actorder():
8282
with pytest.raises(ValueError):
8383
QuantizationArgs(strategy="tensor", actorder="weight")
8484

85-
# test boolean defaulting
85+
# test boolean and none defaulting
8686
assert (
87-
QuantizationArgs(group_size=1, actorder="weight").actorder
88-
== ActivationOrdering.WEIGHT
87+
QuantizationArgs(group_size=1, actorder=True).actorder
88+
== ActivationOrdering.GROUP
8989
)
90+
assert QuantizationArgs(group_size=1, actorder=False).actorder is None
9091
assert QuantizationArgs(group_size=1, actorder=None).actorder is None
9192

9293

0 commit comments

Comments
 (0)