Skip to content

[fsdp] feat: support fsdp2 training and inference in fsdp_workers #1026

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PeterSH6 merged 11 commits into from
May 2, 2025
Merged

[fsdp] feat: support fsdp2 training and inference in fsdp_workers #1026

Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
d6297a6
support fsdp2
lxg2015 Apr 10, 2025
2ce75dd
fsdp2 load full state dict only on rank 0
lxg2015 Apr 11, 2025
3240532
fix indentation
lixiaoguang12 Apr 11, 2025
b58a93d
inspect fsdp wrap layer with _no_split_module
lxg2015 Apr 15, 2025
82aaae4
fsdp2 support cpu_offload_policy mode
lxg2015 Apr 23, 2025
95ce32e
remove DTensor check for fsdp2
lxg2015 Apr 24, 2025
87550dc
expose reshard_after_forward for fsdp2
lxg2015 Apr 28, 2025
d361184
fix pre_commit
lxg2015 May 1, 2025
08a60c3
add exp record at ppo.rst
lxg2015 May 1, 2025
1efd9c2
fix wrap_policy usage in apply_fsdp2 func config
PeterSH6 May 2, 2025
97a0f23
fix wrap_policy usage in test
PeterSH6 May 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
fix wrap_policy usage in test
  • Loading branch information
@PeterSH6
PeterSH6 committed May 2, 2025
commit 97a0f239f9fb0af5bd640d47035da6bba2a68221
4 changes: 2 additions & 2 deletions verl/utils/fsdp_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,8 +418,8 @@ def apply_fsdp2(model, fsdp_kwargs, config):
assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)"

default_transformer_cls_names_to_wrap = getattr(model, "_no_split_modules", None)
fsdp_transformer_layer_cls_to_wrap = config.wrap_policy.get("transformer_layer_cls_to_wrap", default_transformer_cls_names_to_wrap)

fsdp_transformer_layer_cls_to_wrap = config.get("wrap_policy", {}).get("transformer_layer_cls_to_wrap", default_transformer_cls_names_to_wrap)
if isinstance(fsdp_transformer_layer_cls_to_wrap, str):
fsdp_transformer_layer_cls_to_wrap = [fsdp_transformer_layer_cls_to_wrap]

Expand Down