☘️ Fix RM trainer with PEFT + DeepSpeed Stage3

turbo_alignment/common/tf/loaders/model/model.py

@@ -1,13 +1,15 @@
 import torch
 from peft import PeftModel, get_peft_model, prepare_model_for_int8_training
 from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
 
 from turbo_alignment.common.tf.loaders.model.registry import (
     PeftConfigRegistry,
     TransformersAutoModelRegistry,
 )
 from turbo_alignment.settings.model import (
     ModelForPeftSettings,
+    ModelType,
     PreTrainedAdaptersModelSettings,
     PreTrainedModelSettings,
 )
@@ -77,4 +79,13 @@ def load_model(
         # creating learnable adapters and freezing non-training parameters
         model = _prepare_model_for_peft(model, model_settings.peft_settings)
 
+        # deepspeed stage3 is currently doens't work with seq_cls head and peft
+        if model_settings.model_type == ModelType.SEQ_CLS and is_deepspeed_zero3_enabled():
+            model.base_model.model.score = torch.nn.Linear(
+                in_features=model.base_model.model.score.original_module.in_features,
+                out_features=model.base_model.model.score.original_module.out_features,
+                bias=model.base_model.model.score.original_module.bias,
+            )
+            model.base_model.model.score.weight.requires_grad = True
+
     return model

turbo_alignment/trainers/rm.py

@@ -1,9 +1,14 @@
+import os
+from pathlib import Path
 from typing import Any
 
 import torch
+from peft import PeftModel
 from torch import nn
 from transformers import PreTrainedModel
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
 from transformers.trainer_pt_utils import nested_detach
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
 from transformers.utils import logging
 
 from turbo_alignment.trainers.multigpu import MultiGPUCherryPicksTrainer
@@ -64,3 +69,16 @@ def prediction_step(
         labels = labels.long()
 
         return loss, logits, labels
+
+    def _save_checkpoint(self, model, trial, metrics=None):
+        if isinstance(model, PeftModel) and is_deepspeed_zero3_enabled():
+            logger.info('Running custom _save_checkpoint')
+            checkpoint_folder = f'{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}'
+            run_dir = self._get_output_dir(trial=trial)
+            output_dir = Path(os.path.join(run_dir, checkpoint_folder))
+
+            (output_dir / 'cls_head').mkdir(parents=True, exist_ok=True)
+
+            torch.save(model.base_model.model.score.state_dict(), output_dir / 'cls_head' / 'cls_head.pt')
+
+        return super()._save_checkpoint(model=model, trial=trial, metrics=metrics)  # pylint: disable=no-member