[fsdp] fix: preserve Hugging Face fp32-keep modules in FSDP2

[chengcuiping]

What does this PR do?

Fixes #7092.

verl's FSDP2 path did not honor Hugging Face's _keep_in_fp32_modules and _keep_in_fp32_modules_strict declarations. As a result, parameters and floating-point buffers that Hugging Face intentionally keeps in FP32 could be silently reduced during model construction, FSDP2 parameter materialization/all-gather, or full-state-dict loading.

This PR preserves the Hugging Face FP32-keep parameter contract throughout the FSDP2 lifecycle:

• Replaces the blanket FSDP2 build-time cast with a keep-aware cast, preserving matched parameters and floating-point buffers in FP32.
• Wraps eligible FP32-keep modules as child FSDP2 parameter units with param_dtype=None, so their parameters materialize/all-gather in their original FP32 dtype.
• Does not force floating activations to FP32 at the child FSDP2 boundary.
• Preserves all other fields from the caller's mixed-precision policy.
• Aligns full-state-dict tensors with destination parameter and persistent-buffer dtypes before broadcast and assignment.
• Fails early for layouts that FSDP2 cannot represent safely instead of silently reducing precision.

FSDP1 behavior is unchanged. This PR introduces no user-facing API, CLI, or configuration changes.

AI assistance was used in preparing this change; I reviewed every changed line and independently executed all validation reported below.

Checklist Before Starting

• Searched for similar open PRs using the issue number and keep_in_fp32_modules keywords; no existing PR implements this contract.

• Reviewed related PR #4275, which adds dtype-aware FSDP2 wrapping for LoRA modules but does not implement Hugging Face _keep_in_fp32_modules / _keep_in_fp32_modules_strict semantics.

• Formatted the title as [{modules}] {type}: {description}:

  • [fsdp] fix: preserve Hugging Face fp32-keep modules in FSDP2

Test

Command / coverage	Result
pytest tests/utils/test_keep_in_fp32_utils_on_cpu.py -q	67 passed
pytest tests/utils/test_fsdp_wrap_policy_on_cpu.py tests/workers/test_fsdp_gradient_accumulation_sync_on_cpu.py tests/utils/test_fsdp2_peft_wrapping.py -q	14 passed; 1 existing NPU router-replay warning
CUDA_VISIBLE_DEVICES=0 torchrun --standalone --nproc-per-node=1 tests/special_distributed/test_fsdp2_keep_in_fp32.py	PASS on 1 × A100-SXM4-80GB
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc-per-node=2 tests/special_distributed/test_fsdp2_keep_in_fp32.py	PASS on 2 × A100-SXM4-80GB
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc-per-node=2 tests/special_distributed/test_fsdp2_cpu_offload_state_dict.py	PASS
BF16/FP16 × engine-autocast/direct reference transparency	Activation input/output dtypes match the unsharded reference; logits, loss, and every gradient are bit-exact; FP32-keep weights materialize/all-gather in FP32
CPU helper suite with Transformers 4.56.1, 5.3.0, 5.5.3, and 5.10.0	67 passed on every installed version
Negative control: temporarily restore child param_dtype=torch.float32 in an isolated worktree	Expected non-zero exit; caught the BF16 activation boundary changing from the reference's BF16 to FP32
pre-commit run --all-files --show-diff-on-failure --color=always	Final run: 13/13 hooks passed; no files modified
Ruff, Ruff format, and mypy	PASS via the repository pre-commit configuration
git diff --check for the worktree and complete PR diff	PASS

Test environment:

• Python 3.12.3
• PyTorch 2.8.0+cu128
• CUDA 12.8
• NCCL 2.27.3
• NVIDIA A100-SXM4-80GB
• Transformers 4.56.1, 5.3.0, 5.5.3, and 5.10.0

Transformers 5.10.0 was reported as yanked by the installer, but the requested wheel was installed successfully and its CPU helper suite completed with 67 passed tests.

Two distinct regressions are covered:

1. On the unmodified base, stored parameter values can remain correct while a matched weight materializes/all-gathers in the low-precision policy dtype instead of FP32. The weight.dtype forward hook observes only the materialized/all-gathered parameter dtype; it is not evidence of an operator's execution dtype.
2. The earlier child param_dtype=torch.float32 policy forces a different activation boundary. The isolated negative control changes the matched module's BF16 activation input/output to FP32 and fails reference-transparency assertions as expected.

The distributed regression coverage includes:

• Regular and strict FP32-keep declarations.
• FP16 and BF16 dtype-gating behavior.
• Nested matching, parent/child overlap, and duplicate declarations.
• Tied parameter aliases and FP32 buffers.
• Forward, backward, and gradient correctness.
• Full-state-dict round trips.
• Non-callable containers and unsupported fail-closed layouts.
• Models without an FP32-keep declaration.
• Separate observations for activation input/output dtype, internal operator/result tensor dtypes, and materialized/all-gathered parameter dtype.
• Bit-exact logits, loss, and gradient parity against an unsharded reference in identical BF16/FP16 autocast and direct-forward contexts.

API and Usage Example

No user action is required. Existing Hugging Face declarations are honored automatically when verl uses FSDP2:

class MyModel(PreTrainedModel):
    _keep_in_fp32_modules_strict = ["sensitive"]

With strategy=fsdp2 and a low-precision mixed_precision.param_dtype, eligible matched parameters remain stored in FP32 and materialize/all-gather in FP32. The FSDP2 boundary does not redefine module activation input/output or execution dtypes; those remain determined by the caller's autocast context and the module's own casting logic.

Design & Code Changes

Honor Hugging Face declaration semantics

The resolver applies the same dtype gates as Hugging Face Transformers:

• _keep_in_fp32_modules applies when the model dtype is FP16.
• _keep_in_fp32_modules_strict applies when the model dtype is FP16 or BF16.

Matching follows the installed Transformers capability: whole-segment matching for supported 4.x versions and unanchored glob matching for supported 5.x versions. The helper suite was executed with Transformers 4.56.1, 5.3.0, 5.5.3, and 5.10.0.

Preserve FP32 before sharding

FSDPEngine._build_module uses a keep-aware cast for FSDP2. Matched floating-point parameters and buffers remain FP32, including all aliases of tied tensors. Other floating-point states are cast to the configured model dtype.

The FSDP1 construction path is unchanged.

Isolate FP32 parameter units without changing activation boundaries

Each topmost eligible matched module is wrapped before its parent using:

dataclasses.replace(mp_policy, param_dtype=None)

The keep-aware build cast has already preserved the matched parameters in FP32. Leaving the child policy's param_dtype unset makes FSDP2 materialize/all-gather those parameters in their original FP32 dtype without inserting an FP32 activation conversion at the child boundary. Activation input/output and internal execution dtypes remain controlled by the caller's autocast context and the module's own casting logic. All other policy fields are retained.

Reject unsafe or ambiguous layouts

A selected FP32 child absorbs its subtree from the standard wrap-target list to avoid duplicate wrapping. Containers without their own callable forward resolve to eligible callable descendants.

The implementation raises a descriptive error before wrapping for layouts that cannot safely preserve the requested precision, including:

• Root-owned FP32-keep parameters.
• Tied parameters spanning incompatible FSDP units.
• FP32-keep parameters already degraded before wrapping.
• Units that mix kept FP32 parameters with lower-precision parameters.

This fail-closed behavior prevents silent precision loss.

Make destination dtype authoritative during state-dict loading

fsdp2_load_full_state_dict aligns incoming parameters and persistent buffers with their destination dtypes before broadcast and distribution.

The loader:

• Processes every alias with remove_duplicate=False.
• Preserves state-dict metadata.
• Operates on a shallow copy instead of mutating the caller's mapping.

Validation Scope and Limitations

• The regression suite uses a locally constructed toy PreTrainedModel; no large FP32-keep checkpoint was downloaded or executed.
• Validation covers one A100 and DP=2 across two A100s.
• No end-to-end PPO or SFT run is claimed.
• No Inkling 975B, Qwen3-Next, or B300 validation is claimed.
• FSDP1, Megatron, VeOmni, AutoModel, NPU, HSDP/2-D meshes, and other untested backends receive no new support claim from this validation.
• The general FSDP2 CPU-offload state-dict regression passed, but CPU offload combined specifically with FP32-keep child units was not exercised.
• Unsupported FSDP2 layouts now fail explicitly instead of silently degrading precision. This is an intentional behavior change for those layouts.

Checklist Before Submitting

• Read the Contribute Guide.
• Applied the required pre-commit checks: pre-commit run --all-files --show-diff-on-failure --color=always.
• Confirmed that a documentation update is not required because this PR adds no user-facing API, CLI flag, or configuration option.
• Added CPU and distributed regression coverage.
• Registered the distributed regression in tests/special_distributed/run_all.sh.
• Confirmed that the CPU regression is collected by the CPU test workflow.
• Request CI in the ci-request Slack channel when the PR is ready for CI, or use the Feishu group if Slack is unavailable.
• Recipe submodule update is not applicable.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[MagellaX]

Hey, so I just saw this substantial work here on this specific issue... I independently reproduced the original precision-loss bug and confirmed that the overall isolation/loading approach works on both one and two GPUs.

I found one merge-blocking semantic issue: the FP32 child policy currently uses param_dtype=torch.float32.

The matched parameters are already stored in FP32, so param_dtype=None still all-gathers them using their original FP32 dtype. Setting it to torch.float32 additionally casts floating inputs to FP32 at the child boundary, which is not part of Hugging Face’s FP32-keep contract.

This changes the motivating Inkling behavior directly. Inkling’s short-convolution module records its incoming activation dtype, computes internally in FP32, and converts its output back to the original dtype. In an A100 reproduction with PyTorch 2.8:

• param_dtype=None: BF16 input and output; bit-identical to the unsharded reference.
• param_dtype=torch.float32: FP32 input and output; does not match the reference.

The current distributed test misses this because it runs outside FSDPEngine’s autocast context and labels weight.dtype as the operation’s “compute dtype.” The weight dtype proves the all-gathered parameter dtype, but not the actual autocast execution dtype.

I prepared a focused correction that:

• Uses param_dtype=None for FP32-keep child units.
• Runs the regression under the engine-style autocast context.
• Verifies Inkling-style activation-boundary parity.
• Describes the observed weight dtype accurately as the all-gather dtype.
• Fixes the hardcoded Transformers 4.x assertion, which would fail under the current 5.x CI environment.

Corrective commit:
MagellaX@dbefb12

Validation:

• 67 CPU tests passed on Transformers 4.56.1.
• 67 CPU tests passed on Transformers 5.3.0.
• 67 CPU tests passed on Transformers 5.5.3.
• 67 CPU tests passed on Transformers 5.10.0.
• Distributed regression passed on one and two A100-SXM4-40GB GPUs.
• Existing two-GPU FSDP2 CPU-offload regression passed.
• Restoring param_dtype=torch.float32 caused the new activation-boundary regression to fail as expected.
• Ruff, Ruff formatting, and mypy passed.

The correction is based on current main plus this PR. Only commit dbefb124 needs to be cherry-picked.

[chengcuiping]

Thanks @MagellaX — I reviewed the PyTorch 2.8 mixed-precision semantics and your corrective diff, and I agree this is a merge-blocking issue.

My original reasoning conflated the all-gathered parameter dtype with the activation and execution dtype at the FSDP2 unit boundary. The keep-aware build cast already preserves the matched parameters in FP32, so param_dtype=None all-gathers them in their original FP32 dtype. Setting param_dtype=torch.float32 additionally enables FSDP2's forward-input cast and changes the module's activation boundary, which is outside the Hugging Face keep-in-fp32 contract and is incorrect for Inkling-style modules that manage their own internal casting.

The existing distributed test did not catch this because its forward hook observed weight.dtype, which proves the materialized parameter dtype but not the operator execution dtype, and the test ran outside verl's engine-style autocast context. The no-self-cast A/B result therefore was not a valid reason to prefer torch.float32.

I will integrate [dbefb124](https://github.com/MagellaX/verl/commit/dbefb124d94b85322cdef8d9fe66951ff2ca68cd) into this PR branch while preserving your authorship, rather than asking maintainers to cherry-pick across forks. Before pushing the update, I will independently rerun:

• the CPU helper and related existing FSDP tests;
• the corrected regression on one GPU and DP=2;
• the existing DP=2 FSDP2 CPU-offload regression;
• the activation-boundary comparison under the same autocast context as FSDPEngine; and
• pre-commit and diff checks.

I will also audit direct/non-autocast forward callers and extend the activation-boundary parity coverage to FP16. The acceptance criterion will be transparency against the unsharded reference under the same execution context; I will not use an FSDP boundary cast to make an otherwise unsupported path behave differently from that reference.

The Transformers-version assertion fix is independently valid and will be included. I will also update the implementation comments, test terminology, and PR description to distinguish the all-gathered parameter dtype from the activation/compute dtype and remove the claim that the matched module is necessarily computed in FP32.

Thank you for the independent reproduction and for preparing a focused, well-validated correction. I will post the exact revalidation results here after updating the branch.

[chengcuiping]

@MagellaX, thank you for identifying and correcting the activation-boundary semantics.

Your original fix dbefb124 has been integrated as ef686cffa6b1e352cb7bf44e3ac5d929f93c87ba. The original author, yash solanki, and the original Signed-off-by / Co-authored-by trailers are preserved. The follow-up f28d3debffd50663145b00fb1802a234549136c5 adds FP16 coverage and direct/non-autocast reference parity.

Caller audit: FSDPEngineWithLMHead.forward_step() is the production engine entry that calls self.module(...); low-precision execution is wrapped by the engine autocast context. train_batch, infer_batch, and Tinker split-training all flow through forward_backward_batch → forward_step. I found no production path that bypasses that context to call FSDPEngine.module directly.

Independent validation:

• CPU helper suite: 67 passed.
• Related CPU FSDP suites: 14 passed, with one existing NPU router-replay warning.
• 1-GPU FSDP2 regression: PASS on A100-SXM4-80GB.
• DP=2 FSDP2 regression: PASS on two A100-SXM4-80GB GPUs.
• DP=2 FSDP2 CPU-offload state-dict regression: PASS.
• BF16/FP16 × engine-autocast/direct reference transparency: activation input/output dtypes matched; logits, loss, and every gradient were bit-exact.
• Transformers 4.56.1, 5.3.0, 5.5.3, and 5.10.0: 67 passed for each version. The installer reported 5.10.0 as yanked, but the specified wheel installed and completed the suite successfully.
• Isolated param_dtype=torch.float32 negative control: expected non-zero exit after detecting the BF16 activation boundary changed to FP32.
• Final pre-commit: 13/13 hooks passed, including Ruff, Ruff format, and mypy.
• Worktree and complete-PR diff checks: PASS.

The final contract is deliberately narrow: FP32-keep parameters remain stored/materialized/all-gathered in FP32, while the FSDP2 child boundary does not change module activation input/output dtypes. Final PR head: f28d3debffd50663145b00fb1802a234549136c5.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[chengcuiping]

Hi @eric-haibin-lin @vermouth1992 — PR #7165 has been rebased and pushed at f621a305. GitHub now reports it as mergeable.

The rebase preserves the param_dtype=None activation-boundary fix and upstream’s test entries. I also added BF16/FP16 reference-parity coverage across autocast/direct execution with forward prefetch enabled and disabled.

Local validation passed: 85 CPU tests, the FP32-keep regression on 1 and 2 GPUs, the 2-GPU CPU-offload and pinned-transfer regressions, and full pre-commit checks. These ran on A100 with PyTorch 2.8.0+cu128 and Transformers 4.56.1; the current CI environment remains unvalidated.

All 36 workflow runs for this head are currently action_required with no jobs. Could a maintainer please approve the workflows so CI can run? Thank you!