415 lines
17 KiB
Python
415 lines
17 KiB
Python
#
|
|
# Copyright 2016 The BigDL Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import os
|
|
import copy
|
|
import types
|
|
import warnings
|
|
import torch
|
|
import transformers
|
|
from typing import List
|
|
from unittest.mock import patch
|
|
from transformers.dynamic_module_utils import get_imports
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
|
|
from ipex_llm.utils.common.log4Error import invalidInputError
|
|
from ipex_llm.transformers.utils import logger
|
|
from ipex_llm.transformers.npu_models.convert import optimize_llm, optimize_llm_post
|
|
|
|
|
|
def patch_flash_attn_import(filename: str) -> List[str]:
|
|
"""Work around for https://huggingface.co/microsoft/phi-1_5/discussions/72."""
|
|
imports = get_imports(filename)
|
|
if "flash_attn" in imports:
|
|
imports.remove("flash_attn")
|
|
return imports
|
|
|
|
|
|
def ignore_argument(kwargs: dict, key: 'str'):
|
|
arg = kwargs.pop(key, None)
|
|
if arg is not None:
|
|
warnings.warn(f"argument `{key}={arg}` will be ignored")
|
|
|
|
|
|
def save_low_bit(self, model_dir: str, *args, **kwargs):
|
|
origin_device = self.device
|
|
kwargs['safe_serialization'] = False
|
|
self.save_pretrained(model_dir, *args, **kwargs)
|
|
import json
|
|
import os
|
|
# We conveniently save all the keys of the model to have them on hand,
|
|
# so that when using 'low_cpumem load',
|
|
# it's not necessary to load the entire model to extract its keys
|
|
# and we can avoid gc not triggered potentially.
|
|
load_keys = {"all_checkpoint_keys": list(self.state_dict().keys())}
|
|
with open(os.path.join(model_dir, "load_keys.json"), "w") as json_file:
|
|
json.dump(load_keys, json_file)
|
|
if origin_device != 'cpu':
|
|
self.to(origin_device)
|
|
|
|
|
|
class _BaseAutoModelClass:
|
|
HF_MODEL = None
|
|
|
|
@classmethod
|
|
@patch("transformers.dynamic_module_utils.get_imports", patch_flash_attn_import)
|
|
def from_pretrained(cls,
|
|
*args,
|
|
**kwargs):
|
|
"""
|
|
Load a model from a directory or the HF Hub. Use load_in_low_bit parameter to convert
|
|
model to low-bit format, like int4 and int8.
|
|
The loaded model will run supported OPs on NPU, then run other OPs on CPU.
|
|
|
|
Three new arguments are added to extend Hugging Face's from_pretrained method as follows:
|
|
:param load_in_low_bit: str value, options are ``'sym_int4'``, ``'sym_int8'``,
|
|
``'fp16'``, ``'fp32'``.
|
|
Relevant low bit optimizations will be applied to the model.
|
|
:return: a model instance
|
|
"""
|
|
if kwargs.get('device_map', None) not in [None, 'cpu', 'auto']:
|
|
warnings.warn("`device_map` will be ignored")
|
|
kwargs['device_map'] = 'cpu'
|
|
|
|
if kwargs.get('torch_dtype', None) not in [None, 'auto', torch.float, torch.float16]:
|
|
warnings.warn("`torch_dtype` will be ignored, `torch.float` will be used")
|
|
kwargs['torch_dtype'] = torch.float
|
|
|
|
low_bit = kwargs.pop('load_in_low_bit', 'sym_int4')
|
|
qtype_map = {
|
|
'sym_int4': "sym_int4_rtn",
|
|
'sym_int8': "sym_int8_rtn",
|
|
}
|
|
|
|
invalidInputError(low_bit in qtype_map.keys(),
|
|
f"unsupported low_bit: {low_bit}, "
|
|
f"only {list(qtype_map.keys())} are supported")
|
|
qtype = qtype_map[low_bit]
|
|
|
|
kwargs["low_cpu_mem_usage"] = True
|
|
|
|
# ignore following arguments
|
|
ignore_argument(kwargs, "model_hub")
|
|
ignore_argument(kwargs, "lightweight_bmm")
|
|
ignore_argument(kwargs, "load_in_4bit")
|
|
ignore_argument(kwargs, "load_in_8bit")
|
|
ignore_argument(kwargs, "imatrix")
|
|
ignore_argument(kwargs, "mixed_precision")
|
|
ignore_argument(kwargs, "cpu_embedding")
|
|
ignore_argument(kwargs, "embedding_qtype")
|
|
ignore_argument(kwargs, "optimize_model")
|
|
ignore_argument(kwargs, "modules_to_not_convert")
|
|
ignore_argument(kwargs, "quantization_config")
|
|
ignore_argument(kwargs, "speculative")
|
|
pipeline_parallel_stages = kwargs.pop("pipeline_parallel_stages", 1)
|
|
|
|
_args = copy.deepcopy(args)
|
|
_kwargs = copy.deepcopy(kwargs)
|
|
try:
|
|
# To handle the input CUDA setting (such as 'device_map={"":0}'), ignore it
|
|
kwargs.pop('device_map', None)
|
|
model = cls.HF_Model.from_pretrained(*args, **kwargs)
|
|
except NotImplementedError:
|
|
logger.info("Failed to load models with `low_cpu_mem_usage` specified, "
|
|
"will fall to traditional load method with higher memory consumption.")
|
|
_kwargs["low_cpu_mem_usage"] = False
|
|
model = cls.HF_Model.from_pretrained(*_args, **_kwargs)
|
|
model.config.update({"bigdl_lcmu_enabled": False})
|
|
|
|
logger.info(f"Converting model, it may takes up to several minutes ...")
|
|
|
|
if pipeline_parallel_stages > 1:
|
|
invalidInputError(torch.distributed.get_world_size() == pipeline_parallel_stages,
|
|
"Please make sure world size is same as `pipeline_parallel_stages`")
|
|
kwargs['torch_dtype'] = torch.float16
|
|
from .npu_models.pipeline_parallel import pipeline_parallel, pipeline_parallel_generate
|
|
model = pipeline_parallel(model, pipeline_parallel_stages,
|
|
kwargs["torch_dtype"], device="cpu")
|
|
|
|
# add pipeline_parallel_generate to pretrained model dynamically
|
|
model.pipeline_parallel_generate = types.MethodType(pipeline_parallel_generate,
|
|
model)
|
|
|
|
from intel_npu_acceleration_library.compiler import create_npu_kernels
|
|
with torch.no_grad():
|
|
optimize_llm(model)
|
|
if pipeline_parallel_stages == 1:
|
|
cls.load_convert(qtype, model, 'cpu', *args, **kwargs)
|
|
create_npu_kernels(model)
|
|
else:
|
|
cls.load_convert(qtype, model.model, 'cpu', *args, **kwargs)
|
|
create_npu_kernels(model.model)
|
|
optimize_llm_post(model)
|
|
model = model.eval()
|
|
|
|
logger.info(f"Finish to convert model")
|
|
|
|
model.config.update({"bigdl_transformers_low_bit": qtype})
|
|
|
|
# add save_low_bit to pretrained model dynamically
|
|
model.save_low_bit = types.MethodType(save_low_bit, model)
|
|
|
|
return model
|
|
|
|
@classmethod
|
|
def load_convert(cls, q_k, optimize_model, device, *arg, **kwarg):
|
|
from ipex_llm.transformers.npu_models.convert import replace_with_QuantizedLinear
|
|
replace_with_QuantizedLinear(optimize_model, q_k, device=device)
|
|
|
|
@classmethod
|
|
@patch("transformers.dynamic_module_utils.get_imports", patch_flash_attn_import)
|
|
def load_low_bit(cls, pretrained_model_name_or_path: str, *model_args, **kwargs):
|
|
if kwargs.pop('torch_dtype', None) not in [None, 'auto', torch.float]:
|
|
warnings.warn("`torch_dtype` will be ignored, `torch.float` will be used")
|
|
|
|
# ignore following arguments
|
|
ignore_argument(kwargs, "model_hub")
|
|
ignore_argument(kwargs, "lightweight_bmm")
|
|
ignore_argument(kwargs, "cpu_embedding")
|
|
ignore_argument(kwargs, "embedding_qtype")
|
|
ignore_argument(kwargs, "optimize_model")
|
|
ignore_argument(kwargs, "modules_to_not_convert")
|
|
ignore_argument(kwargs, "speculative")
|
|
ignore_argument(kwargs, "pipeline_parallel_stages")
|
|
|
|
from transformers.models.auto.configuration_auto import AutoConfig
|
|
from transformers.modeling_utils import no_init_weights, get_state_dict_dtype
|
|
from transformers.dynamic_module_utils import resolve_trust_remote_code, \
|
|
get_class_from_dynamic_module
|
|
from transformers.models.auto.auto_factory import _get_model_class
|
|
from transformers.utils.generic import ContextManagers
|
|
from transformers.generation.configuration_utils import GenerationConfig
|
|
from ipex_llm.transformers.utils import extract_local_archive_file, get_local_shard_files, \
|
|
load_state_dict
|
|
from accelerate.big_modeling import init_empty_weights
|
|
|
|
trust_remote_code = kwargs.pop("trust_remote_code", None)
|
|
kwargs_orig = copy.deepcopy(kwargs)
|
|
|
|
config, kwargs = AutoConfig.from_pretrained(
|
|
pretrained_model_name_or_path,
|
|
return_unused_kwargs=True,
|
|
trust_remote_code=trust_remote_code,
|
|
**kwargs,
|
|
)
|
|
|
|
# if torch_dtype=auto was passed here, ensure to pass it on
|
|
if kwargs_orig.get("torch_dtype", None) == "auto":
|
|
kwargs["torch_dtype"] = "auto"
|
|
|
|
# Maybe needed when extract_local_archive_file
|
|
subfolder = kwargs.get("subfolder", "")
|
|
variant = kwargs.get("variant", None)
|
|
offload_folder = kwargs.pop("offload_folder", None)
|
|
offload_state_dict = kwargs.pop("offload_state_dict", False)
|
|
torch_dtype = kwargs.pop("torch_dtype", "auto")
|
|
sharded_metadata = None
|
|
|
|
config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path)
|
|
qtype = config_dict.pop("bigdl_transformers_low_bit", False)
|
|
bigdl_lcmu_enabled = config_dict.pop("bigdl_lcmu_enabled", True)
|
|
|
|
invalidInputError(qtype,
|
|
"Detect this model is not a low-bit model, Please use from_pretrained"
|
|
" with load_in_4bit or load_in_low_bit to get a low-bit model , and "
|
|
" serialize the model using save_low_bit first.")
|
|
|
|
invalidInputError(qtype in ["sym_int8_rtn", "sym_int4_rtn"],
|
|
f"Unknown bigdl_transformers_low_bit value: {qtype},"
|
|
f" expected: sym_int4, asym_int4, sym_int5, asym_int5 or sym_int8.")
|
|
|
|
has_remote_code = hasattr(config, "auto_map") and cls.HF_Model.__name__ in config.auto_map
|
|
has_local_code = type(config) in cls.HF_Model._model_mapping.keys()
|
|
trust_remote_code = resolve_trust_remote_code(
|
|
trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code
|
|
)
|
|
if has_remote_code and trust_remote_code:
|
|
class_ref = config.auto_map[cls.HF_Model.__name__]
|
|
model_class = get_class_from_dynamic_module(
|
|
class_ref, pretrained_model_name_or_path, **kwargs
|
|
)
|
|
if os.path.isdir(pretrained_model_name_or_path):
|
|
model_class.register_for_auto_class(cls.HF_Model.__name__)
|
|
else:
|
|
cls.HF_Model.register(config.__class__, model_class, exist_ok=True)
|
|
elif type(config) in cls.HF_Model._model_mapping.keys():
|
|
model_class = _get_model_class(config, cls.HF_Model._model_mapping)
|
|
|
|
resolved_archive_file, is_sharded = extract_local_archive_file(
|
|
pretrained_model_name_or_path,
|
|
subfolder,
|
|
variant)
|
|
|
|
if is_sharded:
|
|
resolved_archive_file, sharded_metadata = \
|
|
get_local_shard_files(pretrained_model_name_or_path,
|
|
resolved_archive_file,
|
|
subfolder=subfolder)
|
|
|
|
# set dtype to instantiate the model under:
|
|
# 1. If torch_dtype is not None, we use that dtype
|
|
# 2. If torch_dtype is "auto", we auto-detect dtype from the loaded state_dict,
|
|
# by checking its first weights entry that is of a floating type
|
|
# - we assume all floating dtype weights are of the same dtype
|
|
# we also may have config.torch_dtype available, but we won't rely on it till v5
|
|
dtype_orig = None
|
|
|
|
if torch_dtype is not None:
|
|
if isinstance(torch_dtype, str):
|
|
if torch_dtype == "auto":
|
|
if hasattr(config, "torch_dtype") and config.torch_dtype is not None:
|
|
torch_dtype = config.torch_dtype
|
|
|
|
else:
|
|
if is_sharded and "dtype" in sharded_metadata:
|
|
torch_dtype = sharded_metadata["dtype"]
|
|
else:
|
|
one_state_dict = load_state_dict(resolved_archive_file[0])
|
|
torch_dtype = get_state_dict_dtype(one_state_dict)
|
|
del one_state_dict # free CPU memory
|
|
else:
|
|
invalidInputError(False,
|
|
f'`torch_dtype` can be either `torch.dtype` or `"auto"`,'
|
|
'but received {torch_dtype}')
|
|
dtype_orig = model_class._set_default_torch_dtype(torch_dtype)
|
|
|
|
# Pretrained Model
|
|
_fast_init = kwargs.pop("_fast_init", True)
|
|
init_contexts = [no_init_weights(_enable=_fast_init)]
|
|
init_contexts.append(init_empty_weights())
|
|
|
|
if bigdl_lcmu_enabled:
|
|
with ContextManagers(init_contexts):
|
|
if config.architectures is not None and config.architectures[0] in \
|
|
["ChatGLMModel", "ChatGLMForConditionalGeneration"]:
|
|
|
|
"""
|
|
ChatGLMModel uses skip_init by default, which will force modules placed on cpu
|
|
if the device is not specified. This will further cause replaced linear
|
|
allocating memory on cpu.
|
|
"""
|
|
kwargs["device"] = "meta"
|
|
model = model_class(config, *model_args, **kwargs)
|
|
else:
|
|
model = model_class(config, *model_args, **kwargs)
|
|
|
|
# Loading args may differ based on their usage
|
|
quant_device = "meta" if bigdl_lcmu_enabled else "cpu"
|
|
logger.info(f"Converting model, it may takes up to several minutes ...")
|
|
from intel_npu_acceleration_library.compiler import create_npu_kernels
|
|
with torch.no_grad():
|
|
optimize_llm(model)
|
|
cls.load_convert(qtype, model, quant_device, *model_args, **kwargs)
|
|
create_npu_kernels(model)
|
|
|
|
model = model.eval()
|
|
|
|
if is_sharded:
|
|
loaded_state_dict_keys = sharded_metadata["all_checkpoint_keys"]
|
|
else:
|
|
import os
|
|
import json
|
|
with open(os.path.join(pretrained_model_name_or_path,
|
|
"load_keys.json"), "r") as json_file:
|
|
loaded_data = json.load(json_file)
|
|
loaded_state_dict_keys = loaded_data["all_checkpoint_keys"]
|
|
|
|
# restore default dtype
|
|
if dtype_orig is not None:
|
|
torch.set_default_dtype(dtype_orig)
|
|
|
|
(
|
|
model,
|
|
missing_keys,
|
|
unexpected_keys,
|
|
mismatched_keys,
|
|
offload_index,
|
|
error_msgs,
|
|
) = model_class._load_pretrained_model(
|
|
model,
|
|
None,
|
|
loaded_state_dict_keys, # XXX: rename?
|
|
resolved_archive_file,
|
|
pretrained_model_name_or_path,
|
|
sharded_metadata=sharded_metadata,
|
|
_fast_init=False, # always false to avoid pre-init behaviors
|
|
low_cpu_mem_usage=bigdl_lcmu_enabled,
|
|
offload_folder=offload_folder,
|
|
offload_state_dict=offload_state_dict,
|
|
dtype=torch_dtype,
|
|
keep_in_fp32_modules=[],
|
|
)
|
|
|
|
# make sure token embedding weights are still tied if needed
|
|
model.tie_weights()
|
|
|
|
# Set model in evaluation mode to deactivate DropOut modules by default
|
|
model.eval()
|
|
|
|
# If it is a model with generation capabilities, attempt to load the generation config
|
|
if model.can_generate():
|
|
try:
|
|
model.generation_config = GenerationConfig.from_pretrained(
|
|
pretrained_model_name_or_path,
|
|
subfolder=subfolder,
|
|
**kwargs,
|
|
)
|
|
except (OSError, TypeError):
|
|
pass
|
|
for param in model.parameters():
|
|
param.requires_grad_(False)
|
|
|
|
return model
|
|
|
|
|
|
class AutoModelForCausalLM(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForCausalLM
|
|
|
|
|
|
class AutoModel(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModel
|
|
|
|
|
|
class AutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForSpeechSeq2Seq
|
|
|
|
|
|
class AutoModelForSeq2SeqLM(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForSeq2SeqLM
|
|
|
|
|
|
class AutoModelForSequenceClassification(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForSequenceClassification
|
|
|
|
|
|
class AutoModelForMaskedLM(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForMaskedLM
|
|
|
|
|
|
class AutoModelForQuestionAnswering(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForQuestionAnswering
|
|
|
|
|
|
class AutoModelForNextSentencePrediction(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForNextSentencePrediction
|
|
|
|
|
|
class AutoModelForMultipleChoice(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForMultipleChoice
|
|
|
|
|
|
class AutoModelForTokenClassification(_BaseAutoModelClass):
|
|
HF_Model = transformers.AutoModelForTokenClassification
|