huggingface / accelerate

🚀 A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support
https://huggingface.co/docs/accelerate
Apache License 2.0
7.34k stars 875 forks source link

Auto create dir when merging FSDP weights #2854

Closed helloworld1 closed 3 weeks ago

helloworld1 commented 3 weeks ago

What does this PR do?

If the merge save_dir does not exist, the accelerate merge-weight will crash after a long time (Weights already merged in memory)

$ accelerate merge-weights checkpoint-2000/pytorch_model_fsdp_0/ checkpoint-2000-merged-weight
Traceback (most recent call last):
  File "/home/jobuser/.local/bin/accelerate", line 8, in <module>
    sys.exit(main())
  File "/home/jobuser/.local/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py", line 48, in main
    args.func(args)
  File "/home/jobuser/.local/lib/python3.10/site-packages/accelerate/commands/merge.py", line 27, in merge_command
    merge_fsdp_weights(
  File "/home/jobuser/.local/lib/python3.10/site-packages/accelerate/utils/fsdp_utils.py", line 270, in merge_fsdp_weights
    save_path = _distributed_checkpoint_to_merged_weights(checkpoint_dir, output_path, safe_serialization)
  File "/home/jobuser/.local/lib/python3.10/site-packages/accelerate/utils/fsdp_utils.py", line 237, in _distributed_checkpoint_to_merged_weights
    save(state_dict, save_path, safe_serialization=safe_serialization)
  File "/home/jobuser/.local/lib/python3.10/site-packages/accelerate/utils/other.py", line 205, in save
    save_func(obj, f)
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/serialization.py", line 627, in save
    with _open_zipfile_writer(f) as opened_zipfile:
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/serialization.py", line 501, in _open_zipfile_writer
    return container(name_or_buffer)
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/serialization.py", line 472, in __init__
    super().__init__(torch._C.PyTorchFileWriter(self.name))
RuntimeError: Parent directory checkpoint-2000-merged-weight does not exist.

This PR automatically create the dir if it does not exist

Fixes # (issue)

Before submitting

Who can review?

@muellerzr

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

HuggingFaceDocBuilderDev commented 3 weeks ago

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.