Issue with saving accelerator state with FSDP

[amarazad]

System Info

python  3.9
accelerate 0.21.0
pytorch                   2.0.1           py3.9_cuda11.7_cudnn8.5.0_0

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

Hi @pacman100 , I am using FSDP with full sharding. I use the following to save sate so that I can resume with the last state:

accelerator.wait_for_everyone()
if accelerator.is_main_process:
    if config["SAVE_STATE"] :
            accelerator.save_state(save_state_dir)

And to resume:

model = accelerator.prepare(model)
optimizer =  accelerator.prepare(optimizer) 
...
if config["RESUME_STATE"]:
              accelerator.wait_for_everyone()
              accelerator.load_state(save_state_dir)

However, while saving, it hangs at accelerator.save_state(save_state_dir) and after a long time throws the following error:

INFO:accelerate.accelerator:Saving FSDP model
[E ProcessGroupNCCL.cpp:828] [Rank 2] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=18364, OpType=ALLGATHER, Timeout(ms)=4800000) ran for 4804648 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:828] [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=18364, OpType=ALLGATHER, Timeout(ms)=4800000) ran for 4804768 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:828] [Rank 3] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=18364, OpType=ALLGATHER, Timeout(ms)=4800000) ran for 4804777 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:828] [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=18338, OpType=_ALLGATHER_BASE, Timeout(ms)=4800000) ran for 4807653 milliseconds before timing out.
[E ProcessGroupNCCL.cpp:455] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[E ProcessGroupNCCL.cpp:460] To avoid data inconsistency, we are taking the entire process down.
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 2820777 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 2820778 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 2820780 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: -6) local_rank: 0 (pid: 2820776) of binary: /home/apa/anaconda3/envs/py39_llm/bin/python
Traceback (most recent call last): ..
....

Expected behavior

As per document, it should save state to be able to resume from there.

[pacman100]

Hello, save_state should be called on all processes as shown below

accelerator.wait_for_everyone()
- if accelerator.is_main_process:
    if config["SAVE_STATE"] :
            accelerator.save_state(save_state_dir)

[amarazad]

Thanks for the correction @pacman.
After calling save_state on all the processes, it moved ahead but still stuck with FSDP optimizer step.
See below:

INFO:__main__: overwriting accelerator state to : model/save_state
INFO:accelerate.accelerator:Saving current state to model/save_state
INFO:accelerate.accelerator:Saving FSDP model
INFO:accelerate.utils.fsdp_utils:Saving model to model/save_state/pytorch_model_0
2023-09-22 02:35:43,456 _dedup_tensors.py:44 INFO p:MainProcess t:MainThread: Duplicate keys to remove: {}
INFO:accelerate.utils.fsdp_utils:Model saved to model/save_state/pytorch_model_0
INFO:accelerate.accelerator:FSDP Model saved to output dir model/save_state
INFO:accelerate.accelerator:Saving FSDP Optimizer

[pacman100]

Hello, can you share a minimal reproducible example for us to look into this.

[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.

[JingyunLiang]

@amarazad

If you are using FULL_STATE_DICT, you should use if accelerator.is_main_process (run on rank 0);
If you are using SHARDED_STATE_DICT, you should remove if accelerator.is_main_process (run on all ranks).

If you meet timeout errors, there are two possible cases:
1, for FULL_STATE_DICT, it's because pytorch need to accumulate the model from other GPUs and then save it to disk on rank0. Other ranks may need to wait for a long time.
2, for SHARDED_STATE_DICT, it's possibly due to that you only run accelerator.save_state(save_state_dir) on rank0. Rank0 is acutally waiting for other ranks to join the saving process.

[Xingyuxingyu]

@amarazad

If you are using FULL_STATE_DICT, you should use if accelerator.is_main_process (run on rank 0); If you are using SHARDED_STATE_DICT, you should remove if accelerator.is_main_process (run on all ranks).

If you meet timeout errors, there are two possible cases: 1, for FULL_STATE_DICT, it's because pytorch need to accumulate the model from other GPUs and then save it to disk on rank0. Other ranks may need to wait for a long time. 2, for SHARDED_STATE_DICT, it's possibly due to that you only run accelerator.save_state(save_state_dir) on rank0. Rank0 is acutally waiting for other ranks to join the saving process.

I use FULL_STATE_DICT and if accelerator.is_main_process (run on rank 0). But it also gets the same issue.