Repository: huggingface/transformers
Pull request: Fix sync gradient

# What does this PR do?

This PR enables to use `sync_each_batch` argument when passing `gradient_accumulation_kwargs` in `AcceleratorConfig`. I'm also removing `adjust_scheduler` docstring as it is not used/enabled for now as we don't prepare the scheduler with accelerate in Trainer. 

Fixes https://github.com/huggingface/transformers/issues/43899

Changed files:
filename: src/transformers/trainer.py
patch: @@ -218,6 +218,7 @@
         DataLoaderConfiguration,
         DistributedDataParallelKwargs,
         DistributedType,
+        GradientAccumulationPlugin,
         load_fsdp_model,
         load_fsdp_optimizer,
         release_memory,
@@ -1764,14 +1765,16 @@ def _inner_training_loop(
                     if step % args.gradient_accumulation_steps == 0:
                         self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
 
-                    # We explicitly want to avoid relying on `accelerator.accumulate` for generation training
-                    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():
+                    # We sync the gradients in the following cases: 1. sync_each_batch set to True 2. Using deepspeed 3. when we are at the last batch sample
+                    if (
+                        self.accelerator.gradient_state.plugin_kwargs.get("sync_each_batch", False)
+                        or self.accelerator.distributed_type == DistributedType.DEEPSPEED
+                        or i == len(batch_samples) - 1
+                    ):
+                        sync_context = contextlib.nullcontext
+                    else:
+                        sync_context = functools.partial(self.accelerator.no_sync, model=model)
+                    with sync_context():
                         tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
 
                     if (
@@ -4009,6 +4012,13 @@ def create_accelerator_and_postprocess(self):
                 )
             else:
                 self.args.gradient_accumulation_steps = grad_acc_kwargs["num_steps"]
+        else:
+            grad_acc_kwargs["num_steps"] = self.args.gradient_accumulation_steps
+
+        # Just making sure that gradient_state have the correct values passed.
+        # We don't rely on `accumulate` from accelerate to set sync_gradients in gradient_state.
+        # Rather, we do it ourselves by setting self.accelerator.gradient_state._set_sync_gradients.
+        gradient_accumulation_plugin = GradientAccumulationPlugin(**grad_acc_kwargs)
 
         accelerator_config = self.args.accelerator_config.to_dict()
 
@@ -4040,6 +4050,7 @@ def create_accelerator_and_postprocess(self):
             "dataloader_config": dataloader_config,
             "fsdp_plugin": fsdp_plugin,
             "deepspeed_plugin": self.args.deepspeed_plugin,
+            "gradient_accumulation_plugin": gradient_accumulation_plugin,
         }
 
         # We defer compatibility checks to accelerator

filename: src/transformers/trainer_pt_utils.py
patch: @@ -1144,8 +1144,6 @@ class AcceleratorConfig:
             Any of the following (optional) keys are acceptable:
               num_steps (`int`): Will take precedence over [`~.TrainingArguments.gradient_accumulation_steps`] if
                 the latter is set to 1, otherwise an exception will be raised.
-              adjust_scheduler (`bool`): Whether to adjust the scheduler steps to account for [`~.TrainingArguments.gradient_accumulation_steps`].
-                The [`accelerate.utils.GradientAccumulationPlugin`] default is `True`.
               sync_each_batch (`bool`): Whether to synchronize the gradients at each data batch.
                 The [`accelerate.utils.GradientAccumulationPlugin`] default is `False`.
         non_blocking (`bool`, *optional*, defaults to `False`):
@@ -1212,8 +1210,6 @@ class AcceleratorConfig:
             "Any of the following (optional) keys are acceptable: "
             "  num_steps (`int`): Will take precedence over [`~.TrainingArguments.gradient_accumulation_steps`] if "
             "    the latter is set to 1, otherwise an exception will be raised. "
-            "  adjust_scheduler (`bool`): Whether to adjust the scheduler steps to account for [`~.TrainingArguments.gradient_accumulation_steps`]. "
-            "    The [`accelerate.utils.GradientAccumulationPlugin`] default is `True`. "
             "  sync_each_batch (`bool`): Whether to synchronize the gradients at each data batch. "
             "    The [`accelerate.utils.GradientAccumulationPlugin`] default is `False`."
         },