[fsdp] fix: make deferred gradient sync configurable

docs/perf/perf_tuning.rst

@@ -198,10 +198,9 @@ Reduce FSDP gradient synchronization during gradient accumulation
 
 When a PPO mini-batch is split into multiple micro-batches, the optimizer only
 steps after the final micro-batch, so gradients only need to be synchronized
-once per mini-batch. The FSDP engine automatically defers gradient
-synchronization on the non-final micro-batches and synchronizes only before the
-final backward. This applies to both the actor and the critic, and requires no
-configuration.
+once per mini-batch. By default, the FSDP engine defers gradient synchronization
+on the non-final micro-batches and synchronizes only before the final backward.
+This applies to both the actor and the critic.
 
 With :math:`M` micro-batches per mini-batch, this reduces gradient
 synchronization from :math:`M` rounds to one round. It does not remove parameter
@@ -211,9 +210,12 @@ numerically identical to synchronizing every micro-batch.
 
 .. note::
     Deferring synchronization retains unsharded gradients until the final
-    micro-batch, which slightly increases peak device memory during gradient
-    accumulation. Forward-only passes are unaffected and always keep the default
-    behavior.
+    micro-batch, which can substantially increase peak device memory during
+    gradient accumulation for large models or long packed sequences. Set
+    ``actor_rollout_ref.actor.fsdp_config.use_no_sync_for_gradient_accumulation=False``
+    (or the corresponding critic FSDP setting) to synchronize and reshard after
+    every micro-batch when memory headroom is limited. Forward-only passes are
+    unaffected.
 
 Migrating to FSDP2
 ----------------------

tests/workers/config/test_engine_config_on_cpu.py

@@ -51,6 +51,11 @@ def test_default_values(self):
         assert config.param_offload is False
         assert config.optimizer_offload is False
         assert config.fsdp_size == -1
+        assert config.use_no_sync_for_gradient_accumulation is True
+
+    def test_gradient_accumulation_sync_can_be_restored_per_micro_batch(self):
+        config = FSDPEngineConfig(use_no_sync_for_gradient_accumulation=False)
+        assert config.use_no_sync_for_gradient_accumulation is False
 
     @pytest.mark.parametrize(
         "offload_params",

tests/workers/test_fsdp_gradient_accumulation_sync_on_cpu.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from contextlib import contextmanager, nullcontext
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -48,9 +49,12 @@ def set_requires_gradient_sync(self, enabled):
         self.events.append(enabled)
 
 
-def _make_engine(module):
+def _make_engine(module, *, use_no_sync_for_gradient_accumulation=True):
     engine = object.__new__(FSDPEngine)
     engine.module = module
+    engine.engine_config = SimpleNamespace(
+        use_no_sync_for_gradient_accumulation=use_no_sync_for_gradient_accumulation,
+    )
     return engine
 
 
@@ -90,6 +94,36 @@ def test_gradient_sync_context_keeps_sync_for_final_micro_batch(monkeypatch):
     assert module.events == ["backward"]
 
 
+@pytest.mark.parametrize("version,module_cls", [(1, _FSDP1Module), (2, _FSDP2Module)])
+def test_gradient_sync_context_keeps_sync_when_disabled(monkeypatch, version, module_cls):
+    module = module_cls()
+    engine = _make_engine(module, use_no_sync_for_gradient_accumulation=False)
+    monkeypatch.setattr(transformer_impl, "fsdp_version", lambda _: version)
+
+    with engine._gradient_sync_context(is_last_micro_batch=False):
+        module.events.append("backward")
+
+    assert module.events == ["backward"]
+
+
+@pytest.mark.parametrize("version,module_cls", [(1, _FSDP1Module), (2, _FSDP2Module)])
+def test_gradient_sync_context_defaults_to_deferred_sync_for_legacy_config(
+    monkeypatch,
+    version,
+    module_cls,
+):
+    module = module_cls()
+    engine = _make_engine(module)
+    engine.engine_config = SimpleNamespace()
+    monkeypatch.setattr(transformer_impl, "fsdp_version", lambda _: version)
+
+    with engine._gradient_sync_context(is_last_micro_batch=False):
+        module.events.append("backward")
+
+    expected = ["enter", "backward", "exit"] if version == 1 else [False, "backward", True]
+    assert module.events == expected
+
+
 def test_forward_backward_batch_syncs_only_final_micro_batch(monkeypatch):
     engine = _make_engine(_FSDP2Module())
     engine.ulysses_sequence_parallel_size = 1

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -43,6 +43,7 @@ actor_rollout_ref:
       entropy_from_logits_chunk_size: ${oc.select:actor_rollout_ref.actor.entropy_from_logits_chunk_size,2048}
       use_torch_compile: true
       entropy_checkpointing: ${oc.select:actor_rollout_ref.actor.entropy_checkpointing,False}
+      use_no_sync_for_gradient_accumulation: false
       pad_to_length: ${oc.select:actor_rollout_ref.actor.pad_to_length,False}
       pad_to_length_bucket: 1024
       forward_only: false
@@ -250,6 +251,7 @@ actor_rollout_ref:
       entropy_from_logits_chunk_size: ${oc.select:actor_rollout_ref.ref.entropy_from_logits_chunk_size,2048}
       use_torch_compile: true
       entropy_checkpointing: ${oc.select:actor_rollout_ref.ref.entropy_checkpointing,False}
+      use_no_sync_for_gradient_accumulation: false
       pad_to_length: ${oc.select:actor_rollout_ref.ref.pad_to_length,False}
       pad_to_length_bucket: ${oc.select:actor_rollout_ref.actor.fsdp_config.pad_to_length_bucket,1024}
       forward_only: true
@@ -563,6 +565,7 @@ critic:
     entropy_from_logits_chunk_size: 2048
     use_torch_compile: true
     entropy_checkpointing: false
+    use_no_sync_for_gradient_accumulation: false
     pad_to_length: false
     pad_to_length_bucket: 1024
     forward_only: false

verl/trainer/config/engine/fsdp.yaml

@@ -56,6 +56,10 @@ use_torch_compile: true
 # Whether to use entropy checkpointing in fsdp.
 entropy_checkpointing: false
 
+# Whether to defer gradient synchronization until the final micro-batch. Disable this when the extra peak memory
+# from retaining unsharded gradients outweighs the communication savings.
+use_no_sync_for_gradient_accumulation: false
+
 # Round every packed micro-batch up to a multiple of `pad_to_length_bucket` tokens, so the packed
 # shape only takes a few distinct values instead of a new one per micro-batch (avoids kernel
 # recompilation / re-autotuning). Requires use_remove_padding=True. Pad tokens are stripped before

verl/workers/config/engine.py

@@ -260,6 +260,9 @@ class FSDPEngineConfig(EngineConfig):
             debugging.
         mixed_precision (Optional[dict[str, Any]]): Mixed precision configuration for FSDP, default None
         dtype (str): Mixed precision training param dtype, default "bfloat16"
+        use_no_sync_for_gradient_accumulation (bool): Whether to defer FSDP gradient synchronization until the
+            final micro-batch. Disabling this reduces peak memory by synchronizing and resharding gradients after
+            every micro-batch. default True
         pad_to_length (bool): Round every packed micro-batch up to a multiple of
             ``pad_to_length_bucket`` tokens, so the packed shape only takes a handful of distinct
             values instead of a new one per micro-batch, which avoids repeated kernel
@@ -293,6 +296,7 @@ class FSDPEngineConfig(EngineConfig):
     entropy_from_logits_chunk_size: int = 2048
     use_torch_compile: bool = True
     entropy_checkpointing: bool = False
+    use_no_sync_for_gradient_accumulation: bool = True
     strategy: str = "fsdp"
     pad_to_length: bool = False
     pad_to_length_bucket: int = 1024

verl/workers/engine/fsdp/transformer_impl.py

@@ -680,7 +680,12 @@ def _gradient_sync_context(self, *, is_last_micro_batch: bool):
         micro-batch to a single round, at the cost of temporarily retaining
         unsharded gradients until the final backward.
         """
-        if is_last_micro_batch:
+        defer_sync = getattr(
+            self.engine_config,
+            "use_no_sync_for_gradient_accumulation",
+            True,
+        )
+        if is_last_micro_batch or not defer_sync:
             yield
             return