Reimplement HF workaround only for llama

henk717 · Jul 22, 2023 · 7a5d813 · 7a5d813
1 parent 8dd7b93
commit 7a5d813
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/modeling/inference_models/generic_hf_torch/class.py b/modeling/inference_models/generic_hf_torch/class.py
@@ -110,6 +110,11 @@ def _load(self, save_model: bool, initial_load: bool) -> None:
             # Also, lazy loader doesn't support GPT-2 models
             self.lazy_load = False
 
+        if self.model_type == "llama":
+            tf_kwargs.update({
+                "pretraining_tp": 1 # Workaround recommended by HF to fix their mistake on the config.json tuners adopted
+            })
+
         logger.debug(
             "lazy_load: {} hascuda: {} breakmodel: {} nobreakmode: {}".format(
                 self.lazy_load,