🎨 Fix gradient checkpointing warnings

turbo_alignment/cherry_picks/chat.py

@@ -36,6 +36,9 @@ def _get_dataset_metrics(
         ref_model: dict = kwargs.get('ref_model', None)
         sft_model: dict = kwargs.get('sft_model', None)
 
+        if model.is_gradient_checkpointing:
+            model.config.use_cache = False
+
         generator = ChatGenerator(
             model=model,
             tokenizer=tokenizer,
@@ -45,6 +48,8 @@ def _get_dataset_metrics(
             return_logits=True,
         )
 
+        model.config.use_cache = True
+
         batch_size = self._generator_transformers_settings.num_return_sequences
 
         generations = generator.generate_from_dataset(dataset)

turbo_alignment/common/tf/callbacks/logging.py

@@ -70,10 +70,7 @@ def _log(self, logs: dict[str, Any], state: TrainerState) -> None:
         self._wandb_run.log({**rewritten_logs, 'train/global_step': state.global_step}, step=state.global_step)
 
     def _fix_table_type(self, logs: dict[str, Any]) -> dict[str, Any]:
-        return {
-            k: wandb.Table(dataframe=v) if isinstance(v, pd.DataFrame) else v
-            for k, v in logs.items()
-        }
+        return {k: wandb.Table(dataframe=v) if isinstance(v, pd.DataFrame) else v for k, v in logs.items()}
 
 
 class ClearMLLoggingCallback(LoggingCallback):

turbo_alignment/metrics/distinctness.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from turbo_alignment.metrics.metric import Metric

turbo_alignment/metrics/registry.py

@@ -1,4 +1,5 @@
 from enum import Enum
+
 from pydantic import field_validator
 
 from turbo_alignment.common.registry import Registrable

turbo_alignment/pipelines/train/classification.py

@@ -68,7 +68,9 @@ def _get_training_args(experiment_settings: ClassificationTrainExperimentSetting
             output_dir=str(experiment_settings.log_path / TRAINER_LOGS_FOLDER),
             label_names=['labels'],
             remove_unused_columns=False,
-            **experiment_settings.trainer_settings.dict(exclude={'loss_settings'}),
+            **experiment_settings.trainer_settings.dict(
+                exclude={'loss_settings', 'gradient_checkpointing_kwargs', 'gradient_checkpointing'},
+            ),
         )
 
     @staticmethod
@@ -81,6 +83,9 @@ def _get_trainer(
         val_dataset: Dataset,
         data_collator: DataCollatorMixin,
     ):
+        if experiment_settings.trainer_settings.gradient_checkpointing:
+            model.gradient_checkpointing_enable({'use_reentrant': True})
+
         if experiment_settings.trainer_settings.loss_settings.alpha == 'auto':
             experiment_settings.trainer_settings.loss_settings.alpha = auto_class_weights(train_dataset)
             logger.info(f'Auto computed class weights: {experiment_settings.trainer_settings.loss_settings.alpha}')

turbo_alignment/pipelines/train/dpo.py

@@ -59,7 +59,9 @@ def _get_training_args(experiment_settings: DPOTrainExperimentSettings) -> DPOTr
             output_dir=str(experiment_settings.log_path / TRAINER_LOGS_FOLDER),
             label_names=[],
             remove_unused_columns=False,
-            **experiment_settings.trainer_settings.dict(),
+            **experiment_settings.trainer_settings.dict(
+                exclude={'gradient_checkpointing_kwargs', 'gradient_checkpointing'},
+            ),
         )
 
     @staticmethod
@@ -72,7 +74,8 @@ def _get_trainer(
         val_dataset: Dataset,
         data_collator: Callable,
     ):
-        model.config.use_cache = not training_args.gradient_checkpointing
+        if experiment_settings.trainer_settings.gradient_checkpointing:
+            model.gradient_checkpointing_enable({'use_reentrant': True})
 
         extra_args = {}
         if experiment_settings.trainer_settings.use_ref_model:

turbo_alignment/pipelines/train/kto.py

@@ -59,7 +59,9 @@ def _get_training_args(experiment_settings: KTOTrainExperimentSettings) -> KTOTr
             output_dir=str(experiment_settings.log_path / TRAINER_LOGS_FOLDER),
             label_names=[],
             remove_unused_columns=False,
-            **experiment_settings.trainer_settings.dict(),
+            **experiment_settings.trainer_settings.dict(
+                exclude={'gradient_checkpointing_kwargs', 'gradient_checkpointing'},
+            ),
         )
 
     @staticmethod
@@ -72,7 +74,8 @@ def _get_trainer(
         val_dataset: Dataset,
         data_collator: Callable,
     ):
-        model.config.use_cache = not experiment_settings.trainer_settings.gradient_checkpointing
+        if experiment_settings.trainer_settings.gradient_checkpointing:
+            model.gradient_checkpointing_enable({'use_reentrant': True})
 
         extra_args = {}
         if experiment_settings.trainer_settings.use_ref_model:

turbo_alignment/pipelines/train/rm.py

@@ -61,7 +61,9 @@ def _get_training_args(experiment_settings: RMTrainExperimentSettings) -> Traini
             output_dir=str(experiment_settings.log_path / TRAINER_LOGS_FOLDER),
             label_names=[],
             remove_unused_columns=False,
-            **experiment_settings.trainer_settings.dict(),
+            **experiment_settings.trainer_settings.dict(
+                exclude={'gradient_checkpointing_kwargs', 'gradient_checkpointing'},
+            ),
         )
 
     @staticmethod
@@ -74,7 +76,10 @@ def _get_trainer(
         val_dataset: Dataset,
         data_collator: DataCollatorMixin,
         **_kwargs,
-    ):
+    ) -> RMTrainer:
+        if experiment_settings.trainer_settings.gradient_checkpointing:
+            model.gradient_checkpointing_enable({'use_reentrant': True})
+
         return RMTrainer(
             model=model,
             tokenizer=tokenizer,

turbo_alignment/pipelines/train/sft.py

@@ -58,7 +58,9 @@ def _get_cherry_pick_callback(
     def _get_training_args(experiment_settings: SftTrainExperimentSettings) -> TrainingArguments:
         return TrainingArguments(
             output_dir=str(experiment_settings.log_path / TRAINER_LOGS_FOLDER),
-            **experiment_settings.trainer_settings.dict(),
+            **experiment_settings.trainer_settings.dict(
+                exclude={'gradient_checkpointing_kwargs', 'gradient_checkpointing'},
+            ),
         )
 
     @staticmethod
@@ -72,7 +74,8 @@ def _get_trainer(
         data_collator: DataCollatorMixin,
         **_kwargs,
     ) -> MultiGPUCherryPicksTrainer:
-        model.config.use_cache = not experiment_settings.trainer_settings.gradient_checkpointing
+        if experiment_settings.trainer_settings.gradient_checkpointing:
+            model.gradient_checkpointing_enable({'use_reentrant': True})
 
         return MultiGPUCherryPicksTrainer(
             model=model,

turbo_alignment/trainers/dpo.py

@@ -26,9 +26,9 @@
 from turbo_alignment.settings.pipelines.train.dpo import (
     APODownLossSettings,
     APOZeroLossSettings,
+    ASFTLossSettings,
     CPOLossSettings,
     DPOLossesType,
-    ASFTLossSettings,
     HingeLossSettings,
     IPOLossSettings,
     KTOLossSettings,
@@ -749,9 +749,6 @@ def _compute_metrics(
         metrics[f'{prefix_name}grad_term'] = (
             (self.dpo_loss_registry.beta * F.sigmoid(rejected_rewards - chosen_rewards)).detach().cpu().mean().item()
         )
-        metrics[f'{prefix_name}grad_term_std'] = (
-            (self.dpo_loss_registry.beta * F.sigmoid(rejected_rewards - chosen_rewards)).detach().cpu().std().item()
-        )
 
         return metrics