fix lora merge with TP and avoid OOM

onmt/models/model_saver.py

@@ -210,6 +210,8 @@ def _save(self, step, model):
             ws = 1
         if ws > 1:
             full_model = [None for _ in range(ws)]
+            for key, value in model_state_dict.items():
+                model_state_dict[key] = value.cpu()
             torch.distributed.all_gather_object(full_model, model_state_dict)
             fm_sd = {}
             for key in full_model[0].keys():
@@ -297,6 +299,8 @@ def _st_save(self, step, model):
             ws = 1
         if ws > 1:
             full_model = [None for _ in range(ws)]
+            for key, value in model_state_dict.items():
+                model_state_dict[key] = value.cpu()
             torch.distributed.all_gather_object(full_model, model_state_dict)
             fm_sd = {}
             for key in full_model[0].keys():

onmt/trainer.py

@@ -563,7 +563,7 @@ def _maybe_report_training(self, step, num_steps, learning_rate, report_stats):
                 if self.earlystopper is None
                 else self.earlystopper.current_tolerance,
                 report_stats,
-                multigpu=self.n_gpu > 1,
+                multigpu=self.n_gpu > 1 and self.parallel_mode == "data_parallel",
             )
 
     def _report_step(self, learning_rate, step, valid_stats=None, train_stats=None):

tools/lora_weights.py

@@ -62,6 +62,7 @@
     lora_opt = lora_checkpoint["opt"]
 
     lora_opt.quant_layers = []  # we need to remove any quantization to merge weights
+    lora_opt.parallel_mode= 'data_parallel'
 
     model = build_base_model(lora_opt, vocabs)