[distributed] Trainer: reduce FSDP/HSDP gradients on the last micro-batch only (#50286)

* Reduce the gradients of FSDP/HSDP-sharded models on the last micro-batch only

The Trainer runs the non-last micro-batches of a gradient accumulation window
under `accelerator.no_sync`, which does nothing for a model sharded at load time
by `DistributedConfig`: Accelerate did not wrap it, so it reduced its gradients on
every micro-batch.

Defer the FSDP2 reduction to the last micro-batch instead. With HSDP, only the
all-reduce across replicas waits (`set_requires_all_reduce`), so the gradients stay
sharded. With FSDP, the reduce-scatter waits too (`set_requires_gradient_sync`),
like Accelerate's FSDP2 `no_sync`.

* Update tests/trainer/distributed/scripts/grad_accum_distributed_config.py

Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

* Use Trainer.end() and drop the redundant device binding in the grad accumulation test

---------

Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

Author: 3outeille

src/transformers/trainer.py

@@ -56,7 +56,7 @@
 from .configuration_utils import PreTrainedConfig
 from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator
 from .debug_utils import DebugOption, DebugUnderflowOverflow
-from .distributed.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft
+from .distributed.fsdp import get_fsdp_ckpt_kwargs, is_fsdp_managed_module, update_fsdp_plugin_peft
 from .feature_extraction_sequence_utils import SequenceFeatureExtractor
 from .feature_extraction_utils import FeatureExtractionMixin
 from .hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend
@@ -1900,6 +1900,9 @@ def _run_epoch(
                     or i == len(batch_samples) - 1
                 ):
                     sync_context = contextlib.nullcontext
+                elif self.is_distributed_loading_by_transformers and is_fsdp_managed_module(self.model):
+                    # Accelerate did not wrap a model sharded at load time, so its `no_sync` does not apply
+                    sync_context = self._defer_fsdp_gradient_reduction
                 else:
                     sync_context = functools.partial(self.accelerator.no_sync, model=model)
                 with sync_context():
@@ -2603,6 +2606,21 @@ def get_tp_size(self) -> int:
         # 4. Default fallback
         return 1
 
+    @contextlib.contextmanager
+    def _defer_fsdp_gradient_reduction(self):
+        """
+        `no_sync` for a model sharded at load time with FSDP2. With HSDP, only the all-reduce across replicas waits for
+        the last micro-batch: the reduce-scatter inside each replica keeps the gradients sharded.
+        """
+        if self.model.config.distributed_config.dp_replicate_size > 1:
+            self.model.set_requires_all_reduce(False)
+        else:
+            self.model.set_requires_gradient_sync(False)
+        try:
+            yield
+        finally:
+            self.model.set_requires_gradient_sync(True)
+
     def _sync_replicated_trainable_parameters(self, model: nn.Module) -> None:
         """
         Keep the trainable parameters that FSDP2 does not manage identical across ranks.

tests/trainer/distributed/scripts/grad_accum_distributed_config.py

@@ -0,0 +1,79 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Worker script: with gradient accumulation, the gradients of a model sharded at load time by `DistributedConfig` are
+only reduced on the last micro-batch. FSDP2 only sets the sharded `.grad` of a parameter once it reduced it.
+
+Launched via ``torchrun`` from ``test_trainer_distributed_transformers.py``.
+"""
+
+import argparse
+from datetime import timedelta
+
+import torch.distributed as dist
+from datasets import Dataset
+
+from transformers import AutoModelForCausalLM, Trainer, TrainerCallback, TrainingArguments
+from transformers.distributed import DistributedConfig
+from transformers.distributed.utils import get_distributed_backend
+from transformers.testing_utils import torch_device
+
+
+class GradientReductionCallback(TrainerCallback):
+    def __init__(self):
+        self.reduced = []
+
+    def on_substep_end(self, args, state, control, model=None, **kwargs):
+        self.reduced.append(any(p.grad is not None for p in model.parameters()))
+
+    def on_pre_optimizer_step(self, args, state, control, model=None, **kwargs):
+        self.reduced.append(all(p.grad is not None for p in model.parameters()))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_dir", required=True)
+    parser.add_argument("--dp_replicate_size", type=int, default=1)
+    script_args = parser.parse_args()
+
+    dist.init_process_group(backend=get_distributed_backend(torch_device), timeout=timedelta(seconds=60))
+    distributed_config = DistributedConfig(
+        dp_replicate_size=script_args.dp_replicate_size,
+        fsdp_size=dist.get_world_size() // script_args.dp_replicate_size,
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        "trl-internal-testing/tiny-Qwen3ForCausalLM", distributed_config=distributed_config
+    )
+    input_ids = [[i + 10] * 8 for i in range(32)]
+    args = TrainingArguments(
+        output_dir=script_args.output_dir,
+        max_steps=2,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=4,
+        save_strategy="no",
+    )
+    callback = GradientReductionCallback()
+    dataset = Dataset.from_dict({"input_ids": input_ids, "labels": input_ids})
+    trainer = Trainer(model=model, args=args, train_dataset=dataset, callbacks=[callback])
+    trainer.train()
+
+    # 3 micro-batches without reduction, then the reduced gradients of the optimizer step, twice
+    expected = [False, False, False, True] * 2
+    assert callback.reduced == expected, f"rank {dist.get_rank()}: {callback.reduced} != {expected}"
+    trainer.end()
+
+
+if __name__ == "__main__":
+    main()

tests/trainer/distributed/test_trainer_distributed_transformers.py

@@ -18,6 +18,8 @@
 
 import os
 
+from parameterized import parameterized
+
 from transformers.testing_utils import (
     TestCasePlus,
     execute_subprocess_async,
@@ -31,6 +33,7 @@
 
 
 PEFT_SAVE_SCRIPT = os.path.join(SCRIPTS_DIR, "peft_save_distributed_config.py")
+GRAD_ACCUM_SCRIPT = os.path.join(SCRIPTS_DIR, "grad_accum_distributed_config.py")
 
 
 @require_torch_greater_or_equal("2.7")
@@ -73,3 +76,9 @@ def test_peft_save_trained_bias(self):
             ],
         )
         execute_subprocess_async(cmd, env=self.get_env())
+
+    @parameterized.expand([("fsdp", 1), ("hsdp", 2)])
+    def test_grad_accumulation_reduces_on_last_micro_batch(self, _, dp_replicate_size):
+        """Gradients of a model sharded at load time are only reduced on the last micro-batch."""
+        script_args = ["--output_dir", self.get_auto_remove_tmp_dir(), "--dp_replicate_size", str(dp_replicate_size)]
+        execute_subprocess_async(self.get_torchrun_cmd(GRAD_ACCUM_SCRIPT, script_args), env=self.get_env())