Repository: OpenNMT/OpenNMT-py

Issue:
 Error with lora_weights.py 
Got following error when I was trying to merge lora weights with last PR.

File “/home/m.barroso/OpenNMT_nllb/OpenNMT-py/tools/lora_weights.py”, line 45, in
base_checkpoint = load_checkpoint(opt.base_model, map_location=torch.device(“cpu”))
TypeError: load_checkpoint() got an unexpected keyword argument ‘map_location’

I had to comment map_location in the function call.

After that, I had to make changes to 3 lines of code that were trying to access to model generator in lora's checkpoint that is empty.

2 of them in model.saver.py

if  "0.weight" in checkpoint["generator"]:  ->  checkpoint["generator"] and "0.weight" in checkpoint["generator"]:
...
if "0.bias" in checkpoint["generator"]:  -> checkpoint["generator"] and "0.bias" in checkpoint["generator"]:

and 1 of them in model.py

name == "generator" and len(checkpoint["generator"].keys()) > 0  -> name == "generator" and checkpoint["generator"]



Discussion:

Resolving patch:
diff --git a/onmt/models/model.py b/onmt/models/model.py
--- a/onmt/models/model.py
+++ b/onmt/models/model.py
@@ -76,7 +76,9 @@ def load_state_dict(
                         param.data = checkpoint["model"][name + "." + param_name]
                         del checkpoint["model"][name + "." + param_name]
                     elif (
-                        name == "generator" and len(checkpoint["generator"].keys()) > 0
+                        "generator" in checkpoint.keys()
+                        and name == "generator"
+                        and len(checkpoint["generator"].keys()) > 0
                     ):
                         param.data = checkpoint["generator"][param_name]
                         del checkpoint["generator"][param_name]
diff --git a/onmt/models/model_saver.py b/onmt/models/model_saver.py
--- a/onmt/models/model_saver.py
+++ b/onmt/models/model_saver.py
@@ -13,7 +13,12 @@ def build_model_saver(model_opt, opt, model, vocabs, optim):
     os.makedirs(os.path.dirname(save_model_path), exist_ok=True)
 
     model_saver = ModelSaver(
-        opt.save_model, model, model_opt, vocabs, optim, opt.keep_checkpoint
+        opt.save_model,
+        model,
+        model_opt,
+        vocabs,
+        optim,
+        opt.keep_checkpoint,
     )
     return model_saver
 
@@ -23,25 +28,35 @@ def load_checkpoint(ckpt_path):
     checkpoint = None
     if ckpt_path:
         logger.info("Loading checkpoint from %s" % ckpt_path)
-        checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
-
-        # This preserves backward-compat for models using customed layernorm
-        def fix_key(s):
-            s = re.sub(r"(.*)\.layer_norm((_\d+)?)\.b_2", r"\1.layer_norm\2.bias", s)
-            s = re.sub(r"(.*)\.layer_norm((_\d+)?)\.a_2", r"\1.layer_norm\2.weight", s)
-            return s
+        checkpoint = torch.load(ckpt_path, map_location=torch.device("cpu"))
+
+        if "model" in checkpoint.keys():
+            # This preserves backward-compat for models using customed layernorm
+            def fix_key(s):
+                s = re.sub(
+                    r"(.*)\.layer_norm((_\d+)?)\.b_2", r"\1.layer_norm\2.bias", s
+                )
+                s = re.sub(
+                    r"(.*)\.layer_norm((_\d+)?)\.a_2", r"\1.layer_norm\2.weight", s
+                )
+                return s
+
+            checkpoint["model"] = {
+                fix_key(k): v for k, v in checkpoint["model"].items()
+            }
+            # Force add_ffnbias to True if bias found in model w_1 keys
+            for key in checkpoint["model"].keys():
+                if "w_1.bias" in key:
+                    checkpoint["opt"].add_ffnbias = True
 
-        checkpoint["model"] = {fix_key(k): v for k, v in checkpoint["model"].items()}
         # fix v2 compatibility
-        if "0.weight" in checkpoint["generator"]:
-            checkpoint["generator"]["weight"] = checkpoint["generator"].pop("0.weight")
-        if "0.bias" in checkpoint["generator"]:
-            checkpoint["generator"]["bias"] = checkpoint["generator"].pop("0.bias")
-
-        for key in checkpoint["model"].keys():
-            if "w_1.bias" in key:
-                checkpoint["opt"].add_ffnbias = True
-
+        if "generator" in checkpoint.keys():
+            if "0.weight" in checkpoint["generator"]:
+                checkpoint["generator"]["weight"] = checkpoint["generator"].pop(
+                    "0.weight"
+                )
+            if "0.bias" in checkpoint["generator"]:
+                checkpoint["generator"]["bias"] = checkpoint["generator"].pop("0.bias")
         # end of patch for backward compatibility
 
     return checkpoint
diff --git a/tools/lora_weights.py b/tools/lora_weights.py
--- a/tools/lora_weights.py
+++ b/tools/lora_weights.py
@@ -42,11 +42,9 @@
 
     init_logger()
 
-    base_checkpoint = load_checkpoint(opt.base_model, map_location=torch.device("cpu"))
+    base_checkpoint = load_checkpoint(opt.base_model)
 
-    lora_checkpoint = load_checkpoint(
-        opt.lora_weights, map_location=torch.device("cpu")
-    )
+    lora_checkpoint = load_checkpoint(opt.lora_weights)
 
     vocabs = dict_to_vocabs(lora_checkpoint["vocab"])
 


Test patch: