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..f93bb4c0513 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: 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 diff --git a/verl/trainer/config/engine/fsdp.yaml b/verl/trainer/config/engine/fsdp.yaml index 6ed61604c72..a14ec2af573 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: 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 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