fix gradient checkpointing warnings

turbo-llm · Oct 27, 2024 · 392185d · 392185d
1 parent 21c3693
commit 392185d
Show file tree

Hide file tree

Showing 7 changed files with 33 additions and 12 deletions.
diff --git a/turbo_alignment/cherry_picks/chat.py b/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)

diff --git a/turbo_alignment/pipelines/train/classification.py b/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}')

diff --git a/turbo_alignment/pipelines/train/dpo.py b/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:

diff --git a/turbo_alignment/pipelines/train/kto.py b/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:

diff --git a/turbo_alignment/pipelines/train/rm.py b/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,

diff --git a/turbo_alignment/pipelines/train/sft.py b/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,

diff --git a/turbo_alignment/trainers/dpo.py b/turbo_alignment/trainers/dpo.py
@@ -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