ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

up (#3391)

Browse source
This commit is contained in:
Scott Roy 2025-10-01 19:14:32 -07:00 • committed by GitHub
commit 291987113b
1 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
unsloth/models/_utils.py
View file

@ -1660,11 +1660,12 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.
from torchao.quantization import (
Float8DynamicActivationInt4WeightConfig,
Float8DynamicActivationFloat8WeightConfig,
Int8DynamicActivationInt4WeightConfig,
Int8DynamicActivationIntxWeightConfig,
Int4WeightOnlyConfig,
PerRow,
quantize_,
)
from torchao.quantization.granularity import PerGroup
from torchao.quantization.qat import QATConfig
filter_fn = None
if qat_scheme == "fp8-int4":
@ -1675,7 +1676,7 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.
base_config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
elif qat_scheme == "int8-int4":
group_size = 32
base_config = Int8DynamicActivationInt4WeightConfig(group_size=group_size)
base_config = Int8DynamicActivationIntxWeightConfig(weight_dtype=torch.int4, weight_granularity=PerGroup(group_size))
filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
elif qat_scheme == "int4":
group_size = 128