Skip to content

Commit

Permalink
[Fix] Allow to create SparseAutoModelForCausalLM with `trust_remote…
Browse files Browse the repository at this point in the history
…_code=True` (#2349)

* initial commit

* better comments

---------

Co-authored-by: bogunowicz@arrival.com <bogunowicz@arrival.com>
  • Loading branch information
dbogunowicz and bogunowicz@arrival.com authored Jul 3, 2024
1 parent f0a3692 commit 179fd90
Showing 1 changed file with 14 additions and 1 deletion.
15 changes: 14 additions & 1 deletion src/sparseml/transformers/sparsification/sparse_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,16 +101,29 @@ def skip(*args, **kwargs):
)

# instantiate compressor from model config
compressor = ModelCompressor.from_pretrained(pretrained_model_name_or_path)
compressor = ModelCompressor.from_pretrained(
pretrained_model_name_or_path, **kwargs
)

# temporarily set the log level to error, to ignore printing out long missing
# and unexpected key error messages (these are EXPECTED for quantized models)
logger = logging.getLogger("transformers.modeling_utils")
restore_log_level = logger.getEffectiveLevel()
logger.setLevel(level=logging.ERROR)

if kwargs.get("trust_remote_code"):
# By artifically aliasing
# class name SparseAutoModelForCausallLM to
# AutoModelForCausalLM we can "trick" the
# `from_pretrained` method into properly
# resolving the logic when
# (has_remote_code and trust_remote_code) == True
cls.__name__ = AutoModelForCausalLM.__name__

model = super(AutoModelForCausalLM, cls).from_pretrained(
pretrained_model_name_or_path, *model_args, **kwargs
)

if model.dtype != model.config.torch_dtype:
_LOGGER.warning(
f"The dtype of the loaded model: {model.dtype} is different "
Expand Down

0 comments on commit 179fd90

Please sign in to comment.