Merge remote-tracking branch 'upstream/main'

Author: dushyantbehl

scripts/run_inference.py

@@ -16,7 +16,7 @@
 # Third Party
 from peft import AutoPeftModelForCausalLM
 from tqdm import tqdm
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 
@@ -156,16 +156,22 @@ def load(
         tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
         # Apply the configs to the adapter config of this model; if no overrides
         # are provided, then the context manager doesn't have any effect.
-        with AdapterConfigPatcher(checkpoint_path, overrides):
-            try:
-                peft_model = AutoPeftModelForCausalLM.from_pretrained(checkpoint_path)
-            except OSError as e:
-                print("Failed to initialize checkpoint model!")
-                raise e
+        try:
+            with AdapterConfigPatcher(checkpoint_path, overrides):
+                try:
+                    model = AutoPeftModelForCausalLM.from_pretrained(checkpoint_path)
+                except OSError as e:
+                    print("Failed to initialize checkpoint model!")
+                    raise e
+        except FileNotFoundError:
+            print("No adapter config found! Loading as a merged model...")
+            # Unable to find the adapter config; fall back to loading as a merged model
+            model = AutoModelForCausalLM.from_pretrained(checkpoint_path)
+
         device = "cuda" if torch.cuda.is_available() else None
         print(f"Inferred device: {device}")
-        peft_model.to(device)
-        return cls(peft_model, tokenizer, device)
+        model.to(device)
+        return cls(model, tokenizer, device)
 
     def run(self, text: str, *, max_new_tokens: int) -> str:
         """Runs inference on an instance of this model.
@@ -198,7 +204,7 @@ def main():
         description="Loads a tuned model and runs an inference call(s) through it"
     )
     parser.add_argument(
-        "--model", help="Path to tuned model to be loaded", required=True
+        "--model", help="Path to tuned model / merged model to be loaded", required=True
     )
     parser.add_argument(
         "--out_file",
@@ -207,7 +213,7 @@ def main():
     )
     parser.add_argument(
         "--base_model_name_or_path",
-        help="Override for base model to be used [default: value in model adapter_config.json]",
+        help="Override for base model to be used for non-merged models [default: value in model adapter_config.json]",
         default=None,
     )
     parser.add_argument(

tuning/sft_trainer.py

@@ -47,34 +47,41 @@ def __init__(self, logger):
         self.logger = logger
 
     def on_log(self, args, state, control, logs=None, **kwargs):
-        """Checks if this log contains keys of interest, e.g., los, and if so, creates
+        """Checks if this log contains keys of interest, e.g., loss, and if so, creates
         train_loss.jsonl in the model output dir (if it doesn't already exist),
         appends the subdict of the log & dumps the file.
         """
         # All processes get the logs from this node; only update from process 0.
         if not state.is_world_process_zero:
             return
 
+        # separate evaluation loss with train loss
         log_file_path = os.path.join(args.output_dir, "train_loss.jsonl")
+        eval_log_file_path = os.path.join(args.output_dir, "eval_loss.jsonl")
         if logs is not None and "loss" in logs and "epoch" in logs:
-            try:
-                # Take the subdict of the last log line; if any log_keys aren't part of this log
-                # object, asssume this line is something else, e.g., train completion, and skip.
-                log_obj = {
-                    "name": "loss",
-                    "data": {
-                        "epoch": round(logs["epoch"], 2),
-                        "step": state.global_step,
-                        "value": logs["loss"],
-                        "timestamp": datetime.isoformat(datetime.now()),
-                    },
-                }
-            except KeyError:
-                return
-
-            # append the current log to the jsonl file
-            with open(log_file_path, "a") as log_file:
-                log_file.write(f"{json.dumps(log_obj, sort_keys=True)}\n")
+            self._track_loss("loss", log_file_path, logs, state)
+        elif logs is not None and "eval_loss" in logs and "epoch" in logs:
+            self._track_loss("eval_loss", eval_log_file_path, logs, state)
+
+    def _track_loss(self, loss_key, log_file, logs, state):
+        try:
+            # Take the subdict of the last log line; if any log_keys aren't part of this log
+            # object, assume this line is something else, e.g., train completion, and skip.
+            log_obj = {
+                "name": loss_key,
+                "data": {
+                    "epoch": round(logs["epoch"], 2),
+                    "step": state.global_step,
+                    "value": logs[loss_key],
+                    "timestamp": datetime.isoformat(datetime.now()),
+                },
+            }
+        except KeyError:
+            return
+
+        # append the current log to the jsonl file
+        with open(log_file, "a") as f:
+            f.write(f"{json.dumps(log_obj, sort_keys=True)}\n")
 
 
 def train(