From f765e4f439c92c553334154c2dcc4f17e27a8e39 Mon Sep 17 00:00:00 2001 From: Mengyuyang <781650427@qq.com> Date: Sat, 15 Aug 2026 15:19:35 +0800 Subject: [PATCH 1/3] [fsdp] fix: make deferred gradient sync configurable Preserve the current communication-optimized default while allowing memory-constrained jobs to synchronize and reshard gradients after every micro-batch. Keep legacy engine configs on the current deferred-sync behavior. Signed-off-by: Mengyuyang <781650427@qq.com> --- docs/perf/perf_tuning.rst | 16 +++++---- .../config/test_engine_config_on_cpu.py | 5 +++ ..._fsdp_gradient_accumulation_sync_on_cpu.py | 36 ++++++++++++++++++- .../config/_generated_ppo_trainer.yaml | 3 ++ verl/trainer/config/engine/fsdp.yaml | 4 +++ verl/workers/config/engine.py | 4 +++ verl/workers/engine/fsdp/transformer_impl.py | 7 +++- 7 files changed, 66 insertions(+), 9 deletions(-) diff --git a/docs/perf/perf_tuning.rst b/docs/perf/perf_tuning.rst index ed6c407c0fa..da7ae265e9c 100644 --- a/docs/perf/perf_tuning.rst +++ b/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 ---------------------- diff --git a/tests/workers/config/test_engine_config_on_cpu.py b/tests/workers/config/test_engine_config_on_cpu.py index 1253f5c9ab9..2c6522e815f 100644 --- a/tests/workers/config/test_engine_config_on_cpu.py +++ b/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", diff --git a/tests/workers/test_fsdp_gradient_accumulation_sync_on_cpu.py b/tests/workers/test_fsdp_gradient_accumulation_sync_on_cpu.py index 7e080469591..bb5c1896c5d 100644 --- a/tests/workers/test_fsdp_gradient_accumulation_sync_on_cpu.py +++ b/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 diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index d6b3fb72dde..72bd3459d04 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/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: true 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: true 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: true pad_to_length: false pad_to_length_bucket: 1024 forward_only: false diff --git a/verl/trainer/config/engine/fsdp.yaml b/verl/trainer/config/engine/fsdp.yaml index 6ed61604c72..009e6a138c9 100644 --- a/verl/trainer/config/engine/fsdp.yaml +++ b/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: true + # 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 diff --git a/verl/workers/config/engine.py b/verl/workers/config/engine.py index 9ba309029d3..ce3a21158c6 100644 --- a/verl/workers/config/engine.py +++ b/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 diff --git a/verl/workers/engine/fsdp/transformer_impl.py b/verl/workers/engine/fsdp/transformer_impl.py index dd5cf3311da..cb32478bcf6 100644 --- a/verl/workers/engine/fsdp/transformer_impl.py +++ b/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 From 488c23de48d354bb6ae1019d1c9e7638fca754c4 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Fri, 28 Aug 2026 10:28:23 +0800 Subject: [PATCH 2/3] set use_no_sync_for_gradient_accumulation=false --- verl/trainer/config/_generated_ppo_trainer.yaml | 6 +++--- verl/trainer/config/engine/fsdp.yaml | 5 +---- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index 72bd3459d04..f93bb4c0513 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -43,7 +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: true + 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 @@ -251,7 +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: true + 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 @@ -565,7 +565,7 @@ critic: entropy_from_logits_chunk_size: 2048 use_torch_compile: true entropy_checkpointing: false - use_no_sync_for_gradient_accumulation: true + use_no_sync_for_gradient_accumulation: false pad_to_length: false pad_to_length_bucket: 1024 forward_only: false diff --git a/verl/trainer/config/engine/fsdp.yaml b/verl/trainer/config/engine/fsdp.yaml index 009e6a138c9..3178040b566 100644 --- a/verl/trainer/config/engine/fsdp.yaml +++ b/verl/trainer/config/engine/fsdp.yaml @@ -3,7 +3,6 @@ _target_: verl.workers.config.FSDPEngineConfig # policy for wrapping the model wrap_policy: - # Minimum number of parameters to trigger wrapping a layer with FSDP min_num_params: 0 @@ -58,7 +57,7 @@ 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: true +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 @@ -83,7 +82,6 @@ dtype: bfloat16 # ["bfloat16", "float16"] # QAT (Quantization-Aware Training) configuration qat: - # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs _target_: verl.workers.config.QATEngineConfig @@ -98,7 +96,6 @@ qat: # Patterns to ignore (e.g., lm_head, embed_tokens) ignore_patterns: - - "lm_head" - "embed_tokens" - "re:.*mlp.gate$" From 8f88ea9d913da9fe457e962c0989667ce589fae3 Mon Sep 17 00:00:00 2001 From: wuxibin Date: Fri, 28 Aug 2026 10:30:53 +0800 Subject: [PATCH 3/3] format --- verl/trainer/config/engine/fsdp.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/verl/trainer/config/engine/fsdp.yaml b/verl/trainer/config/engine/fsdp.yaml index 3178040b566..a14ec2af573 100644 --- a/verl/trainer/config/engine/fsdp.yaml +++ b/verl/trainer/config/engine/fsdp.yaml @@ -3,6 +3,7 @@ _target_: verl.workers.config.FSDPEngineConfig # policy for wrapping the model wrap_policy: + # Minimum number of parameters to trigger wrapping a layer with FSDP min_num_params: 0 @@ -82,6 +83,7 @@ dtype: bfloat16 # ["bfloat16", "float16"] # QAT (Quantization-Aware Training) configuration qat: + # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs _target_: verl.workers.config.QATEngineConfig @@ -96,6 +98,7 @@ qat: # Patterns to ignore (e.g., lm_head, embed_tokens) ignore_patterns: + - "lm_head" - "embed_tokens" - "re:.*mlp.gate$"