Add QA and MaskedLM task encoder architectures in run_fp8

Signed-off-by: Andrea Fasoli <[email protected]>

Author: andrea-fasoli

fms_mo/run_quant.py

@@ -36,7 +36,11 @@
 from datasets import load_from_disk
 from huggingface_hub.errors import HFValidationError
 from torch.cuda import OutOfMemoryError
-from transformers import AutoTokenizer
+from transformers import (
+    AutoModelForMaskedLM,
+    AutoModelForQuestionAnswering,
+    AutoTokenizer,
+)
 import torch
 import transformers
 
@@ -204,9 +208,23 @@ def run_fp8(model_args, data_args, opt_args, fp8_args):
 
     logger = set_log_level(opt_args.log_level, "fms_mo.run_fp8")
 
-    model = SparseAutoModelForCausalLM.from_pretrained(
-        model_args.model_name_or_path, torch_dtype=model_args.torch_dtype
-    )
+    if model_args.task_type == "lm":
+        model = SparseAutoModelForCausalLM.from_pretrained(
+            model_args.model_name_or_path,
+            torch_dtype=model_args.torch_dtype,
+        )
+    elif model_args.task_type == "qa":
+        model = AutoModelForQuestionAnswering.from_pretrained(
+            model_args.model_name_or_path,
+            torch_dtype=model_args.torch_dtype,
+        )
+    elif model_args.task_type == "mlm":
+        model = AutoModelForMaskedLM.from_pretrained(
+            model_args.model_name_or_path,
+            torch_dtype=model_args.torch_dtype,
+        )
+    else:
+        raise ValueError(f"Unsupported task: {model_args.task_type}")
     tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
 
     recipe = QuantizationModifier(

fms_mo/training_args.py

@@ -55,6 +55,17 @@ class ModelArguments(TypeChecker):
     """Dataclass for model related arguments."""
 
     model_name_or_path: str = field(default="facebook/opt-125m")
+    task_type: str = field(
+        default="lm",
+        metadata={
+            "choices": ["lm", "qa", "mlm"],
+            "help": (
+                "Instantiate model for selected task: 'lm' (language modeling), 'qa' "
+                "(question answering, for encoders), 'mlm' (masked language modeling, "
+                "for encoders)."
+            ),
+        },
+    )
     torch_dtype: str = field(default="bfloat16")
     device_map: Optional[str] = field(
         default=None,