loading weight is so slow for fsdp2

[jiaqiw09]

System Info

main branch

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

I am trying to train a model using LLaMAFactory or MsSwift. When using FSDP2 with Accelerate, the process gets stuck during fsdp2_load_full_state_dict, specifically at this line:
https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py#L516

full_param = full_param.detach().to(device_mesh.device_type)

At this point, the available memory keeps decreasing. The process remains stuck here for nearly half an hour without completing.

and the fsdp yaml looks like

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_reshard_after_forward: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_version: 2
machine_rank: 0
main_training_function: main
mixed_precision: bf16  # or fp16
num_machines: 1  # the number of nodes
num_processes: 8  # the number of GPUs in all nodes
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Expected behavior

None

[iavinas]

FULL_STATE_DICT is an FSDP1 thing — FSDP2 doesn't support it. PyTorch's own docs say so explicitly. Accelerate should really warn you about this but currently doesn't, which is a bug.
The fix is just changing one line:

fsdp_state_dict_type: SHARDED_STATE_DICT

Pytorch docs - https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html

Compared to PyTorch FSDP1 (FullyShardedDataParallel):

FSDP2 simplifies some of the API surface: e.g. FSDP2 does not directly support full state dicts. Instead, users can reshard the sharded state dicts containing DTensor s to full state dicts themselves using DTensor APIs like DTensor.full_tensor() or by using higher-level APIs like PyTorch Distributed Checkpoint ‘s distributed state dict APIs. Also, some other args have been removed; see here for details.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.