modelscope / modelscope/ms-swift

FSDP2 doesn't save HF checkpoint and doesn't use parameters from fsdp_config

Open
#8,198 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug stale
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.