modelscope / modelscope/ms-swift
FSDP2 doesn't save HF checkpoint and doesn't use parameters from fsdp_config
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 15.7k
- Forks
- 1.7k
- Avg merge
- 1d 16h
- Merged PRs (30d)
- 136
Description
Checklist / 检查清单
- I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues,确认这是一个新的 bug report。
Bug Description / Bug 描述
when flag --fsdp fsdp2 used, there is no HF checkpoint saved during saving, only pytorch_model_fsdp_0 and optimizer_0 along with rng_states and scheduler.pt
Also I've tried to pass my own json with fsdp_config to change state_dict_type:
"fsdp": "shard_grad_op auto_wrap",
"fsdp_config": {
"fsdp_version": 2,
"reshard_after_forward": false,
"auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
"cpu_ram_efficient_loading": true,
"state_dict_type": "FULL_STATE_DICT",
"activation_checkpointing": true
}
But nothing was changed in terms of speed, memory consumption and saving behavior, so I've checked source code and noticed that only parameter that used is fsdp_version
After that I've tried to change state_dict_type via env variable
os.environ["FSDP_STATE_DICT_TYPE"] = "FULL_STATE_DICT"
And it worked for common models like qwen, but when I used my custom model that I've registered via ms-swift, my model_modeling.py file was not saved. Such behavior only occurs when I use fsdp2, and persistent in both versions of ms-swift 4.0.0 and 3.12.6, when I use deepspeed everything works fine. Also in version 3.12.6 I could pass config for fsdp1 which looked like that
fsdp: "shard_grad_op auto_wrap"
fsdp_config:
fsdp_version: 1
backward_prefetch: "backward_pre"
forward_prefetch: false
limit_all_gathers: true
use_orig_params: true
sync_module_states: true
cpu_ram_efficient_loading: true
activation_checkpointing: true
min_num_params: 0
And it also worked as intended and saved hf checkpoint along everything else
How to Reproduce / 如何复现
just launch example of train with fsdp2 and wait for saving
Additional Information / 补充信息
This bug persistent in newest 4.0.0 and previous 3.12.6 versions of ms-swift
In both cases I've used torch==2.9.1 and cuda==12.8
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the training example using --fsdp fsdp2, then trace how fsdp_config and the FSDP_STATE_DICT_TYPE environment variable are handled. Compare behavior with the provided FSDP1 configuration and a registered custom model's model_modeling.py; done means the configuration affects saving and an HF checkpoint is produced alongside the existing shard files.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100