Skip to content

Commit

Permalink
feat: add check for quantized model (#913)
Browse files Browse the repository at this point in the history
* feat: add check for quantized model

* chore: refactor and add another check

* Update src/axolotl/utils/models.py

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>
  • Loading branch information
NanoCode012 and winglian committed Dec 4, 2023
1 parent 992e742 commit a581e9f
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions src/axolotl/utils/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,27 @@
LOG = logging.getLogger("axolotl")


def check_model_config(cfg: DictDefault, model_config: AutoConfig):
quant_config_exists = hasattr(model_config, "quantization_config")
quant_config_method_is_gptq = (
quant_config_exists
and "quant_method" in model_config.quantization_config
and model_config.quantization_config["quant_method"] == "gptq"
)

if cfg.gptq and not quant_config_method_is_gptq:
raise ValueError(
"model_config.quantization_config is not set or quant_method is not set to gptq. "
"Please make sure to point to a GPTQ model."
)

if not cfg.gptq and quant_config_exists:
raise ValueError(
"model_config.quantization_config is set but `gptq` flag is not. "
"Please use the `gptq` flag to train quantized model or point to a non-quantized model."
)


def load_model_config(cfg):
model_config_name = cfg.base_model_config or cfg.base_model
trust_remote_code = cfg.trust_remote_code is True
Expand All @@ -38,6 +59,8 @@ def load_model_config(cfg):
for key, val in cfg.model_config.items():
setattr(model_config, key, val)

check_model_config(cfg, model_config)

return model_config


Expand Down

0 comments on commit a581e9f

Please sign in to comment.