[LLM] Support WOQ scheme asym (#1266)

Signed-off-by: Wang, Chang <chang1.wang@intel.com>

Author: changwangss

intel_extension_for_transformers/llm/quantization/autograd/functions.py

@@ -76,6 +76,7 @@ def forward(
  compute_dtype=None,
  weight_dtype=None,
  scale_dtype=None,
+ scheme=None,
  ):
  # # 1. Dequantize
  # B_dequant = torch.zeros(out.shape[-1], A.shape[-1], dtype=torch.float)
@@ -114,7 +115,7 @@ def forward(
  compute_dtype,
  weight_dtype,
  scale_dtype,
- False,
+ False if scheme == "sym" else True,
  )
  else:
  out = qbits_woq_linear_ref_impl(
@@ -154,7 +155,7 @@ def backward(ctx, grad_output):
  None,
  )
 
- req_gradA, _, _, req_gradBias, _, _, _ = ctx.needs_input_grad
+ req_gradA, _, _, req_gradBias, _, _, _, _ = ctx.needs_input_grad
  A, B = ctx.tensors
  grad_A, grad_B, grad_bias = None, None, None
 
@@ -176,7 +177,7 @@ def backward(ctx, grad_output):
  if req_gradA:
  grad_A = torch.matmul(grad_output, B.to(grad_output.dtype))
 
- return grad_A, grad_B, None, grad_bias, None, None, None
+ return grad_A, grad_B, None, grad_bias, None, None, None, None
 
 
 def matmul_kbit(
@@ -187,11 +188,13 @@ def matmul_kbit(
  compute_dtype,
  weight_dtype,
  scale_dtype,
+ scheme,
  do_dequant=False,
 ):
+
  if do_dequant:
  return MatMulKBit.apply(
- A, B, out, bias, compute_dtype, weight_dtype, scale_dtype
+ A, B, out, bias, compute_dtype, weight_dtype, scale_dtype, scheme
  )
  else:
  qbits_debug_flag = os.getenv('QBITS_DEBUG', 'NULL')
@@ -206,7 +209,7 @@ def matmul_kbit(
  compute_dtype,
  weight_dtype,
  scale_dtype,
- False,
+ False if scheme == "sym" else True,
  )
  else:
  out = qbits_woq_linear_ref_impl(

intel_extension_for_transformers/llm/quantization/nn/modules.py

@@ -129,6 +129,7 @@ def forward(self, x: torch.Tensor):
  self.compute_dtype if self.compute_dtype is not None else "fp32",
  self.weight_dtype,
  self.scale_dtype if self.scale_dtype is not None else "fp32",
+ self.scheme,
  do_dequant=self.training,
  )
  shape[-1] = self.out_features
@@ -146,7 +147,7 @@ def set_weights_bias(self, weight_data, bias=None):
  self.compute_dtype if self.compute_dtype is not None else "fp32",
  self.weight_dtype,
  self.scale_dtype if self.scale_dtype is not None else "fp32",
- False,
+ False if self.scheme == "sym" else True,
  )
  self.weight = ParamsQBits(
  data=weight,
@@ -315,7 +316,7 @@ def merge(self, safe_merge: bool = False) -> None:
  self.compute_dtype,
  self.weight_dtype,
  self.scale_dtype,
- False,
+ False if self.scheme == "sym" else True,
  )
 
  self.weight = ParamsQBits(
@@ -360,7 +361,7 @@ def unmerge(self) -> None:
  self.compute_dtype,
  self.weight_dtype,
  self.scale_dtype,
- False,
+ False if self.scheme == "sym" else True,
  )
 
  self.weight = ParamsQBits(

intel_extension_for_transformers/transformers/utils/config.py

@@ -142,6 +142,11 @@ def post_init(self):
  if not isinstance(self.scheme, str):
  raise ValueError("scheme must be a string")
 
+ if self.scheme == "asym" and (self.compute_dtype == "int8" or self.weight_dtype.startswith("fp") \
+ or self.weight_dtype.startswith("nf") or self.scale_dtype != "fp32"):
+ raise ValueError("WeightOnlyQuantization doesn't support asym with \
+ compute_dtype int8 or weight_dtype float or scale_dtype non-fp32 now, \
+ please use sym scheme")
  self.use_llm_runtime = False
 
  def post_init_xpu(self):

tests/CI/test_quantization.py

@@ -356,6 +356,7 @@ def test_quantization_for_llm(self):
  quantization_config=woq_config,
  use_llm_runtime=False
  )
+ woq_model.eval()
  output = woq_model(dummy_input)
  print("output:", float(output[0][0][0][0]))
  self.assertTrue(isclose(float(output[0][0][0][0]), 0.16387596726417542, rel_tol=1e-04))
@@ -368,6 +369,7 @@ def test_quantization_for_llm(self):
  quantization_config=woq_config,
  use_llm_runtime=False
  )
+ woq_model.eval()
  output = woq_model(dummy_input)
  print("output:", float(output[0][0][0][0]))
  self.assertTrue(isclose(float(output[0][0][0][0]), 0.17239853739738464, rel_tol=1e-04))
@@ -380,23 +382,25 @@ def test_quantization_for_llm(self):
  quantization_config=woq_config,
  use_llm_runtime=False
  )
+ woq_model.eval()
  output = woq_model(dummy_input)
- # fp8
+ #fp8
  woq_config = WeightOnlyQuantConfig(weight_dtype="fp8_e5m2", scale_dtype="fp8_e8m0")
  woq_model = AutoModelForCausalLM.from_pretrained(
  model_name_or_path, quantization_config=woq_config, use_llm_runtime=False
  )
+ woq_model.eval()
  output = woq_model(dummy_input)
  self.assertTrue(
  isclose(float(output[0][0][0][0]), 0.16162332892417908, rel_tol=1e-04)
  )
-
  # amp
  amp_config = MixedPrecisionConfig()
  amp_model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
  quantization_config=amp_config,
  use_llm_runtime=False
  )
+ amp_model.eval()
  output = amp_model(dummy_input)
  self.assertTrue(isclose(float(output[0][0][0][0]), 0.1689453125, rel_tol=1e-04))
  # bitsandbytes, for cpu is fp32 model
@@ -410,6 +414,7 @@ def test_quantization_for_llm(self):
  load_in_4bit=True,
  use_llm_runtime=False
  )
+ bit4_model.eval()
  output = bit4_model(dummy_input)
  print("output:", float(output[0][0][0][0]))
  self.assertTrue(isclose(float(output[0][0][0][0]), 0.18726778030395508, rel_tol=1e-04))
@@ -420,6 +425,7 @@ def test_quantization_for_llm(self):
  use_llm_runtime=False,
  device_map="cpu"
  )
+ bit8_model.eval()
  output = bit8_model(dummy_input)
  print("output:", float(output[0][0][0][0]))
  self.assertTrue(isclose(float(output[0][0][0][0]), 0.1675747185945511, rel_tol=1e-04))
@@ -441,6 +447,7 @@ def test_quantization_for_llm(self):
  quantization_config=woq_config,
  use_llm_runtime=False
  )
+ woq_model.eval()
  output = woq_model(dummy_input)
  print("output:", float(output[0][0][0][0]))
  self.assertTrue(isclose(float(output[0][0][0][0]), 0.17126554250717163, rel_tol=1e-04))