don't use no_sync when deepspeed doesn't support it for certain zero …

…stages (#35157) * don't use no_sync when deepspeed doesn't support it for certain zero stages * chore: lint * fix no_sync context for deepspeed across all zero types * chore: lint
huggingface · Dec 13, 2024 · 1b6cb1e · 1b6cb1e
1 parent b1d5d6d
commit 1b6cb1e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -2517,6 +2517,7 @@ def _inner_training_loop(
                     context = (
                         functools.partial(self.accelerator.no_sync, model=model)
                         if i != len(batch_samples) - 1
+                        and self.accelerator.distributed_type != DistributedType.DEEPSPEED
                         else contextlib.nullcontext
                     )
                     with context():