Fix chatglm multi-reference problem

intel-analytics · Aug 8, 2024 · 7b9ff3d · 7b9ff3d
1 parent 07210c6
commit 7b9ff3d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/python/llm/src/ipex_llm/transformers/convert.py b/python/llm/src/ipex_llm/transformers/convert.py
@@ -331,6 +331,11 @@ def _replace_with_low_bit_linear(model, qtype, modules_to_not_convert=None,
         if any(key in full_module_name for key in modules_to_not_convert):
             continue
 
+        if is_linear and getattr(model_config, "model_type", None) == "chatglm" and name == "lm_head":
+            # Now we re-reference it to output_layer
+            model._modules[name] = model._modules["transformer"]._modules["output_layer"]
+            continue
+
         if is_linear and not isinstance(module, LowBitLinear):
             in_features, out_features, mp_group = linear_args
             optimize_lm_head = False