Rm

src/transformers/trainer.py

@@ -69,7 +69,10 @@ from .models.auto.modeling_auto import (
     MODEL_MAPPING_NAMES,
 )
 from .optimization import Adafactor, get_scheduler
-from .pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
+from .pytorch_utils import (
+    ALL_LAYERNORM_LAYERS,
+    is_torch_greater_or_equal_than_1_13,
+)
 from .tokenization_utils_base import PreTrainedTokenizerBase
 from .trainer_callback import (
     CallbackHandler,
@@ -205,6 +208,7 @@ if is_accelerate_available():
     from accelerate import Accelerator, skip_first_batches
     from accelerate import __version__ as accelerate_version
     from accelerate.utils import (
+        AutocastKwargs,
         DistributedDataParallelKwargs,
         DistributedType,
         GradientAccumulationPlugin,
@@ -3094,6 +3098,8 @@ class Trainer:
         """
         if self.use_cpu_amp:
             ctx_manager = torch.cpu.amp.autocast(cache_enabled=cache_enabled, dtype=self.amp_dtype)
+        elif self.accelerator is not None:
+            ctx_manager = self.accelerator.autocast(autocast_handler=AutocastKwargs(cache_enabled=cache_enabled))
         else:
             ctx_manager = contextlib.nullcontext()