model loading

mttl/models/utils.py

@@ -106,7 +106,8 @@ def model_loader_helper(
         raise ValueError("Specify either 'load_in_4bit' or 'load_in_8bit' or neither.")
 
     from transformers import (
-        AutoModel,
+        AutoModelForCausalLM,
+        AutoModelForSeq2SeqLM,
         BitsAndBytesConfig,
         LlamaForCausalLM,
         PreTrainedModel,
@@ -144,17 +145,23 @@ def model_loader_helper(
 
         model_name = os.environ["PHI_PATH"]
         logger.info(f"Loading phi-2 model from {os.environ['PHI_PATH']}")
-    try:
-        model_object = AutoModel.from_pretrained(
-            model_name,
-            device_map=device_map,
-            trust_remote_code=True,
-            attn_implementation=attn_implementation,
-            quantization_config=bnb_config,
-            torch_dtype=torch_dtype,
-        )
-    except Exception as e:
-        logger.error(f"loading model: {e}")
+    else:
+        model_object = None
+        for klass in [AutoModelForCausalLM, AutoModelForSeq2SeqLM]:
+            try:
+                model_object = klass.from_pretrained(
+                    model_name,
+                    device_map=device_map,
+                    trust_remote_code=True,
+                    attn_implementation=attn_implementation,
+                    quantization_config=bnb_config,
+                    torch_dtype=torch_dtype,
+                )
+                break
+            except:
+                continue
+        if model_object is None:
+            raise ValueError(f"Couldn't load {model_name}!")
 
     if bnb_config is not None:
         model_object = prepare_model_for_kbit_training(model_object)