should DeepSpeedEngine.save_checkpoint be only under main_process

[better629]

Describe the bug
If DeepSpeedEngine.save_checkpoint should be run only under main_process(rank=0) ？

I have seen the save the model only for the main process when using the distributed training mode in pytorch inside https://pytorch.org/tutorials/intermediate/ddp_tutorial.html, but it seems that there is no need to be under main_process when using save_checkpoint. Does any know why ?

If the save_checkpoint is under main_process in huggingface diffusers, the save stage will hands on.

System info (please complete the following information):
Ubuntu 20.04
Nvidia GTX 3090
CUDA Version: 11.7
Torch: 1.13.1
Diffusers: 0.15.0.dev0
deepspeed: 0.8.1
xformers: 0.0.17.dev466
accelerate: 0.16.0

[tjruwase]

@better629, for DP training the ranks have exactly the same model state so it would be wasteful for all the ranks to write out the same checkpoint state. Alos, using rank=0 is just convention.

[better629]

@tjruwase Yes, I do think so.
But as the issue mentioned, if i judge rank==0 before save checkpoint, the process will hang on under DeepSpeed. It's a little confusing.

application related issues are huggingface/accelerate#1171 and huggingface/diffusers#2606

[tjruwase]

@better629, sorry I think I misunderstood. The documented expectation that all ranks should call DeepSpeedEngine.save_checkpoint(), even if in reality only DP rank 0 will actually create a checkpoint file. Otherwise, you will observe the hang as there is a barrier in that call. Can you please try that to see if it fixes the issue?

[better629]

@tjruwase I was confused with why all ranks should call save_checkpoint which is different with pytorch ddp. Since there is a api description, I will try to find it out from code.

[tjruwase]

@better629, optimizations such as the following are possible because all data parallel ranks call save_checkpoint: #1419.

[better629]

@tjruwase
yes, all ranks call save_checkpoint method, but seems only rank=0 save model params. which mentioned in

# deepspeed/runtime/engine.py
        # All procs then call module_state_dict(), but only procs of data
        # parallel rank 0 save the general model params.

and also the inner code has a rank=0 judgement.

[tjruwase]

@tjruwase
yes, all ranks call save_checkpoint method, but seems only rank=0 save model params. which mentioned in

This is true for the DeepSpeedEngine, PipelineEngine uses all ranks to save model params, as in the PR. Does that help clarify things?

[better629]

@tjruwase yes. The code is the truth.