🅰️ Remove apex (#4139)

Author: qgallouedec

trl/trainer/nash_md_trainer.py

@@ -30,7 +30,7 @@
 )
 from transformers.trainer_utils import EvalPrediction
 from transformers.training_args import OptimizerNames
-from transformers.utils import is_apex_available, is_peft_available
+from transformers.utils import is_peft_available
 
 from ..data_utils import is_conversational, maybe_apply_chat_template
 from ..models.modeling_base import GeometricMixtureWrapper
@@ -47,10 +47,6 @@
 )
 
 
-if is_apex_available():
-    from apex import amp
-
-
 if is_peft_available():
     from peft import PeftModel
 
@@ -501,10 +497,6 @@ def training_step(
         if self.args.n_gpu > 1:
             loss = loss.mean()  # mean() to average on multi-gpu parallel training
 
-        if self.use_apex:
-            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
-                scaled_loss.backward()
-        else:
-            self.accelerator.backward(loss, **kwargs)
+        self.accelerator.backward(loss, **kwargs)
 
         return loss.detach() / self.args.gradient_accumulation_steps

trl/trainer/online_dpo_trainer.py

@@ -43,7 +43,6 @@
     ProcessorMixin,
     Trainer,
     TrainerCallback,
-    is_apex_available,
 )
 from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
 from transformers.trainer_utils import EvalPrediction, seed_worker
@@ -78,9 +77,6 @@
 if is_peft_available():
     from peft import PeftConfig, PeftModel
 
-if is_apex_available():
-    from apex import amp
-
 
 if is_sagemaker_mp_enabled():
     from smdistributed.modelparallel import __version__ as SMP_VERSION
@@ -1457,11 +1453,7 @@ def training_step(
         if self.args.n_gpu > 1:
             loss = loss.mean()  # mean() to average on multi-gpu parallel training
 
-        if self.use_apex:
-            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
-                scaled_loss.backward()
-        else:
-            self.accelerator.backward(loss, **kwargs)
+        self.accelerator.backward(loss, **kwargs)
 
         return loss.detach() / self.args.gradient_accumulation_steps
 

trl/trainer/xpo_trainer.py

@@ -27,7 +27,6 @@
     PreTrainedTokenizerBase,
     ProcessorMixin,
     TrainerCallback,
-    is_apex_available,
 )
 from transformers.trainer_utils import EvalPrediction
 from transformers.training_args import OptimizerNames
@@ -47,10 +46,6 @@
 from .xpo_config import XPOConfig
 
 
-if is_apex_available():
-    from apex import amp
-
-
 if is_peft_available():
     from peft import PeftModel
 
@@ -545,10 +540,6 @@ def training_step(
         if self.args.n_gpu > 1:
             loss = loss.mean()  # mean() to average on multi-gpu parallel training
 
-        if self.use_apex:
-            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
-                scaled_loss.backward()
-        else:
-            self.accelerator.backward(loss, **kwargs)
+        self.accelerator.backward(loss, **kwargs)
 
         return loss.detach() / self.args.gradient_accumulation_steps