update reranker trainer

Author: 545999961

FlagEmbedding/finetune/reranker/decoder_only/base/trainer.py

@@ -2,7 +2,7 @@
 import torch
 import logging
 from typing import Optional
-from transformers.deepspeed import is_deepspeed_zero3_enabled
+# from transformers.deepspeed import is_deepspeed_zero3_enabled
 
 from FlagEmbedding.abc.finetune.reranker import AbsRerankerTrainer
 from peft import get_peft_model_state_dict
@@ -29,13 +29,13 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
 
         torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
 
-        if is_deepspeed_zero3_enabled():
-            if state_dict is None:
-                state_dict = self.model.state_dict()
-            prefix = 'model.'
-            assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys())
-            state_dict = {k[len(prefix):]: v for k, v in state_dict.items()}
-            lora_state_dict = get_peft_model_state_dict(self.model.model, state_dict)
-            if self.args.process_index <= 0:
-                torch.save(lora_state_dict, os.path.join(output_dir, "adapter_model.bin"))
-                print(f"Save adapter model at {output_dir}")
+        # if is_deepspeed_zero3_enabled():
+        #     if state_dict is None:
+        #         state_dict = self.model.state_dict()
+        #     prefix = 'model.'
+        #     assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys())
+        #     state_dict = {k[len(prefix):]: v for k, v in state_dict.items()}
+        #     lora_state_dict = get_peft_model_state_dict(self.model.model, state_dict)
+        #     if self.args.process_index <= 0:
+        #         torch.save(lora_state_dict, os.path.join(output_dir, "adapter_model.bin"))
+        #         print(f"Save adapter model at {output_dir}")

FlagEmbedding/finetune/reranker/decoder_only/layerwise/trainer.py

@@ -2,7 +2,7 @@
 import torch
 import logging
 from typing import Optional
-from transformers.deepspeed import is_deepspeed_zero3_enabled
+# from transformers.deepspeed import is_deepspeed_zero3_enabled
 from peft import get_peft_model_state_dict
 
 from FlagEmbedding.abc.finetune.reranker import AbsRerankerTrainer
@@ -29,13 +29,13 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
 
         torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
 
-        if is_deepspeed_zero3_enabled():
-            if state_dict is None:
-                state_dict = self.model.state_dict()
-            prefix = 'model.'
-            assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys())
-            state_dict = {k[len(prefix):]: v for k, v in state_dict.items()}
-            lora_state_dict = get_peft_model_state_dict(self.model.model, state_dict)
-            if self.args.process_index <= 0:
-                torch.save(lora_state_dict, os.path.join(output_dir, "adapter_model.bin"))
-                print(f"Save adapter model at {output_dir}")
+        # if is_deepspeed_zero3_enabled():
+        #     if state_dict is None:
+        #         state_dict = self.model.state_dict()
+        #     prefix = 'model.'
+        #     assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys())
+        #     state_dict = {k[len(prefix):]: v for k, v in state_dict.items()}
+        #     lora_state_dict = get_peft_model_state_dict(self.model.model, state_dict)
+        #     if self.args.process_index <= 0:
+        #         torch.save(lora_state_dict, os.path.join(output_dir, "adapter_model.bin"))
+        #         print(f"Save adapter model at {output_dir}")