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

Update _utils.py

Browse source
This commit is contained in:
Daniel Han 2025-10-16 04:48:55 -07:00
commit f1e8a0468a
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/_utils.py
View file

@ -1664,7 +1664,7 @@ pass
class TorchAOConfig:
qat_scheme : str = "int4"
base_config : AOBaseConfig = Int4WeightOnlyConfig
filter_fn : Callable = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
filter_fn : Callable = None
group_size : int = 128
pass
@ -1694,11 +1694,11 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: Union[str, TorchA
filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
elif qat_scheme == "fp8-fp8":
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
base_config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
base_config = Float8DynamicActivationFloat8WeightConfig(granularity = PerRow())
elif qat_scheme == "int8-int4":
from torchao.quantization import Int8DynamicActivationIntxWeightConfig
group_size = 32
base_config = Int8DynamicActivationIntxWeightConfig(weight_dtype=torch.int4, weight_granularity=PerGroup(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":
from torchao.quantization import Int4WeightOnlyConfig