[fsdp] fix: make deferred gradient sync configurable

[Mengyuyang]

What does this PR do?

Complete the configuration contract for deferred FSDP gradient synchronization introduced by #7095.

The merged implementation currently defers synchronization on every non-final micro-batch, but it does not expose the documented use_no_sync_for_gradient_accumulation setting. This PR adds the missing FSDP engine option and keeps its default true to preserve the current runtime behavior. Memory-constrained jobs can explicitly set it to false to synchronize and reshard gradients after every micro-batch.

The change supports both FSDP1 and FSDP2, preserves forward-only behavior, and retains the current deferred-sync behavior for legacy engine config objects that do not define the new field.

Related: #7095, #6010.

Checklist Before Starting

• Searched for similar open PRs: use_no_sync_for_gradient_accumulation, FSDP no_sync, and gradient sync accumulation. No open duplicate was found at the time of submission preparation.
• Formatted the PR title as [fsdp] fix: make deferred gradient sync configurable.
• Verified that the latest main still has the unconditional deferred-sync path and does not define use_no_sync_for_gradient_accumulation.

Test

Focused CPU tests that do not depend on the local two-rank Gloo backend:

python -m pytest -q \
  tests/workers/test_fsdp_gradient_accumulation_sync_on_cpu.py \
  tests/workers/config/test_engine_config_on_cpu.py \
  -k "not distributed_accumulation_matches_per_micro_batch_sync"

Result:

21 passed, 1 deselected

Additional validation completed:

• Ruff lint passed for all changed Python files.
• Ruff format check passed for all changed Python files.
• Mypy passed for the changed-file pre-commit run.
• Hydra composition verified true for the actor, reference model, and critic defaults.
• A Hydra actor override was verified to resolve to false.
• Generated trainer configuration verification passed.
• Documentation, license, device API, DataProto, naming, and compileall hooks passed.
• git diff --check passed.

The real two-rank Gloo equivalence test was attempted on Windows, but the local PyTorch backend failed during init_process_group() with makeDeviceForHostname(): unsupported gloo device, before entering the FSDP code under test. It remains enabled for Linux CI.

The complete pre-commit suite is not marked as passed locally because the current upstream tree contains root-level tests/test_* files rejected by the repository's validate-structure hook. This PR does not add or modify those files.

API and Usage Example

No existing CLI or Python API is removed. The new FSDP engine option defaults to true, matching the current merged runtime behavior.

# Default: defer gradient synchronization until the final micro-batch.
actor_rollout_ref.actor.fsdp_config.use_no_sync_for_gradient_accumulation=True

# Memory-constrained actor training: synchronize and reshard after every micro-batch.
actor_rollout_ref.actor.fsdp_config.use_no_sync_for_gradient_accumulation=False

# Apply the same memory-oriented behavior to critic training when needed.
critic.fsdp.use_no_sync_for_gradient_accumulation=False

With M micro-batches, true reduces gradient synchronization from M rounds to one, but temporarily retains unsharded gradients and can increase peak device memory. Setting the option to false trades additional gradient communication for lower peak memory. When there is only one micro-batch, both settings follow the normal synchronized backward path.

Design & Code Changes

• Add use_no_sync_for_gradient_accumulation: bool = True to FSDPEngineConfig.
• Add the option to the FSDP Hydra YAML and regenerate the flattened PPO reference configuration.
• Gate the existing FSDP1 no_sync() and FSDP2 set_requires_gradient_sync(False) paths on the new option.
• Use a getattr(..., True) compatibility fallback so legacy or subclass-specific engine config objects preserve the current deferred-sync behavior.
• Keep the final micro-batch and all forward-only execution synchronized exactly as before.
• Add CPU coverage for the default, explicit opt-out, FSDP1, FSDP2, final-micro-batch, exception restoration, and legacy-config behavior.
• Document the communication-versus-memory trade-off and the actor/critic override paths.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply the complete pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always. Relevant changed-file hooks passed locally; see the test limitation above.
• Updated the FSDP performance-tuning documentation.
• Added CPU unit tests covering the configuration and both FSDP synchronization paths. The existing Linux two-rank test remains enabled for CI.
• Once the PR is ready for CI, request CI in the ci-request channel in the verl Slack workspace or the linked Feishu group.
• The recipe submodule is not affected.
• AI assistance was used for repository inspection, implementation review, test execution, and drafting this description. The human submitter must review every changed line and be able to explain the synchronization and memory trade-off before requesting review.

[lihy11]

failed my training because of the same issue。。。

[alanhuangyoo]

Measured what the deferral actually costs, in case it helps the docs paragraph land more concretely.

The retained buffer is fp32, not the bf16 gradient. MixedPrecisionPolicy defaults to reduce_dtype=fp32 against bf16 parameters, so the early return in FSDP2's to_accumulated_grad_if_needed never fires:

if (self.reduce_dtype is None or unsharded_param is None
        or unsharded_param.grad is None
        or unsharded_param.grad.dtype == self.reduce_dtype):
    return
self.unsharded_accumulated_grad = unsharded_grad.to(self.reduce_dtype)

bf16 never equals fp32, so every parameter gets an upcast copy.

The part that surprised me is that it does not scale out. Two gloo ranks up to eight, 262,144-parameter model, bf16 params / fp32 reduce, measured after the first non-final micro-batch:

        parameter shard    retained accumulated grad
ws=1        1024.0 KiB              1024.0 KiB
ws=2         512.0 KiB              1024.0 KiB
ws=4         256.0 KiB              1024.0 KiB
ws=8         128.0 KiB              1024.0 KiB

The shard halves per rank as FSDP promises; the retained gradient stays at the whole model in fp32, so at ws=8 it is 8x the parameter shard on the same card. That works out to 4 bytes per parameter regardless of topology — 108 GiB a rank at 27B, which is more than the "slightly increases peak device memory" the current docs suggest and is why this flag matters.

Reproducer if useful:

model.set_requires_gradient_sync(False)
model(x).sum().backward()
held = sum(
    p.unsharded_accumulated_grad.numel() * p.unsharded_accumulated_grad.element_size()
    for m in model.modules()
    if (s := getattr(m, "_get_fsdp_state", None)) and s()._fsdp_param_group
    for p in s()._fsdp_param_group.fsdp_params
    if p.unsharded_accumulated_grad is not None
)

Worth asserting the 4-bytes-per-parameter figure in the test file rather than only the context-manager behaviour? It would pin the reason the flag exists, and it runs on CPU.