Fix WOQ huggingface model loading (#1400)

Author: changwangss

intel_extension_for_transformers/transformers/modeling/modeling_auto.py

@@ -83,6 +83,7 @@ def recover_export_model(model, current_key_name=None):
  Return optimum format model.
  """
  from ..llm.quantization.nn.modules import QuantizedLinearQBits
+
  for name, module in model.named_children():
  if current_key_name is None:
  current_key_name = []
@@ -194,8 +195,13 @@ def save_low_bit(
  )
  return
 
- if self.quantization_config.weight_dtype not in \
- ["fp8_e5m2", "fp8_e4m3", "nf4", "fp4", "int4_fullrange"]:
+ if self.quantization_config.weight_dtype not in [
+ "fp8_e5m2",
+ "fp8_e4m3",
+ "nf4",
+ "fp4",
+ "int4_fullrange",
+ ]:
  convert_model_to_public(self)
  os.makedirs(save_directory, exist_ok=True)
  # use transformers original `save_pretrained` function
@@ -336,7 +342,27 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
  return_unused_kwargs=True,
  **kwargs,
  )
- if hasattr(config, "quantization_config"):
+
+ if kwargs.get("use_llm_runtime", None) is not None:
+ use_neural_speed = kwargs.pop("use_llm_runtime", True) and not use_xpu
+ logger.warning(
+ "use_llm_runtime is deprecated in version 1.3.2, please use_neural_speed instead."
+ )
+ elif kwargs.get("use_neural_speed", None) is not None:
+ use_neural_speed = kwargs.pop("use_neural_speed", True) and not use_xpu
+ else:
+ if hasattr(config, "model_type") == False:
+ logger.error(
+ "Can't get the model_type. Please check the correct model_type"
+ )
+ exit(0)
+
+ if config.model_type in cls.model_type_list and not use_xpu:
+ use_neural_speed = True
+ else:
+ use_neural_speed = False
+
+ if hasattr(config, "quantization_config") and not use_neural_speed:
  if config.quantization_config is None:
  logger.warning(
  "Quantization_config loading failed. If you want to load saved "
@@ -369,26 +395,6 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
  "Saved low bit model loading failed, please check your model."
  )
  exit(0)
- if kwargs.get("use_llm_runtime", None) is not None:
- use_neural_speed = kwargs.pop("use_llm_runtime", True) and not use_xpu
- logger.warning(
- "use_llm_runtime is deprecated in version 1.3.2, please use_neural_speed instead."
- )
- elif kwargs.get("use_neural_speed", None) is not None:
- use_neural_speed = kwargs.pop("use_neural_speed", True) and not use_xpu
- else:
- if hasattr(config, "model_type") == False:
- logger.error(
- "Can't get the model_type. Please check the correct model_type"
- )
- exit(0)
-
- if config.model_type in cls.model_type_list and not use_xpu:
- logger.info("Using Neural Speed...")
- use_neural_speed = True
- else:
- logger.info("Using Pytorch...")
- use_neural_speed = False
 
  import intel_extension_for_transformers.transformers.modeling.modeling_map
 
@@ -437,7 +443,9 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
  if quantization_config is None:
  if use_neural_speed:
  # use wnf4_sfp32_cfp32_g32_sym by default
- quantization_config = RtnConfig(compute_dtype="fp32", weight_dtype="nf4")
+ quantization_config = RtnConfig(
+ compute_dtype="fp32", weight_dtype="nf4"
+ )
  else:
  quantization_config = RtnConfig(
  bits=4,
@@ -502,7 +510,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
  ):
  logger.info("Applying Weight Only Quantization.")
  if use_neural_speed:
- logger.info("Using LLM runtime.")
+ logger.info("Using Neural Speed.")
  quantization_config.post_init_runtime()
  from neural_speed import Model
 
@@ -966,6 +974,7 @@ def load_low_bit(cls, pretrained_model_name_or_path, *model_args, **kwargs):
  kwargs["torch_dtype"] = "auto"
  config = kwargs.pop("config", None)
  quantization_config = config.quantization_config
+
  if quantization_config["quant_method"] == "rtn":
  quantization_config = RtnConfig.from_dict(quantization_config)
  elif quantization_config["quant_method"] == "awq":
@@ -976,7 +985,6 @@ def load_low_bit(cls, pretrained_model_name_or_path, *model_args, **kwargs):
  quantization_config = GPTQConfig.from_dict(quantization_config)
  elif quantization_config["quant_method"] == "autoround":
  quantization_config = AutoRoundConfig.from_dict(quantization_config)
-
  assert (
  quantization_config is not None
  ), "Detect this model is not a low-bit model."
@@ -1170,8 +1178,13 @@ def load_low_bit(cls, pretrained_model_name_or_path, *model_args, **kwargs):
  model = model_class(config, *model_args, **kwargs)
  else:
  model = model_class(config, *model_args, **kwargs)
- if config.quantization_config["weight_dtype"] not in \
- ["fp8_e5m2", "fp8_e4m3", "fp4", "nf4", "int4_fullrange"]:
+ if config.quantization_config["weight_dtype"] not in [
+ "fp8_e5m2",
+ "fp8_e4m3",
+ "fp4",
+ "nf4",
+ "int4_fullrange",
+ ]:
  model = build_woq_model(model, quantization_config)
  else:
  model = replace_linear(
@@ -1221,8 +1234,12 @@ def load_low_bit(cls, pretrained_model_name_or_path, *model_args, **kwargs):
 
  # Set model in evaluation mode to deactivate DropOut modules by default
  model.eval()
- if config.quantization_config["weight_dtype"] not in \
- ["fp8_e5m2", "fp8_e4m3", "nf4", "fp4" "int4_fullrange"]:
+ if config.quantization_config["weight_dtype"] not in [
+ "fp8_e5m2",
+ "fp8_e4m3",
+ "nf4",
+ "fp4" "int4_fullrange",
+ ]:
  model = replace_linear(
  model,
  quantization_config=quantization_config,

tests/CI/test_weight_only.py

@@ -188,7 +188,7 @@ def test_auto_model_saving_loading(self):
  module_list.append(name)
  self.assertTrue(len(module_list) > 0)
  model.save_pretrained(self.workspace, safe_serialization=False)
- loaded_model = AutoModelForCausalLM.from_pretrained(self.workspace)
+ loaded_model = AutoModelForCausalLM.from_pretrained(self.workspace, use_neural_speed=False)
  for name, module in loaded_model.named_modules():
  if isinstance(module, QuantizedLinearQBits):
  module_list.append(name)