Merge lora gpu support

examples/merge_lora.py

@@ -41,6 +41,12 @@ class MergeLoraArguments:
             "help": "output merged full model path"
         },
     )
+    local_rank: Optional[int] = field(
+        default=-1,
+        metadata={
+            "help": "local rank for deepspeed",
+        },
+    )
 
 
 def main():

scripts/run_merge_lora.sh

@@ -54,7 +54,7 @@ elif [ ${device} == "gpu" ]; then
         --lora_model_path ${lora_model_path} \
         --output_model_path ${output_model_path} \
         --device ${device} \
-        --ds_config configs/ds_config_eval.json
+        --ds_config configs/ds_config_zero3_for_eval.json
 else
     echo "error: unknown device \"${device}\"" 1>&2
     exit 1

src/lmflow/models/hf_decoder_model.py

@@ -202,6 +202,7 @@ def __init__(
             if model_args.torch_dtype in ["auto", None]
             else getattr(torch, model_args.torch_dtype)
         )
+        logger.debug(f"torch_dtype on init: {torch_dtype}")
 
         config_kwargs = {
             "cache_dir": model_args.cache_dir,
@@ -852,7 +853,13 @@ def save(self, dir, save_full_model=False, *args, **kwargs):
         """
         self.get_tokenizer().save_pretrained(dir)
         if save_full_model and self.model_args.use_lora:
-            self.backend_model_full.save_pretrained(dir)
+            save_dtype = (
+                torch.float16
+                if self.model_args.torch_dtype in ["auto", None]
+                else getattr(torch, self.model_args.torch_dtype)
+            )
+            self.backend_model_full.to(dtype=save_dtype).save_pretrained(dir)
+            logger.warning(f"Save full model with dtype: {save_dtype}")
         else:
             self.get_backend_model().save_pretrained(dir)