resume training does not work for multi-gpus training

[forever208]

I add --resume_checkpoint $path_to_checkpoint$ to continue the training, it works for a single GPU, but does not work for multi-gpus

the code gets stuck here:

Logging to /proj/ihorse_2021/users/x_manni/guided-diffusion/log9
creating model and diffusion...
creating data loader...
start training...
loading model from checkpoint: /proj/ihorse_2021/users/x_manni/guided-diffusion/log9/model200000.pt...

[VigneshSrinivasan10]

@forever208 I have the same problem and the code gets stuck forever.

On further investigation, I found that the test script image_sample.py reloads the model.
Here is the difference:
The test script reloads the model before placing the model on cuda.
However, the training script already has the model on cuda and this leads to the problem of getting stuck there.
Upon debugging, I found the code to get stuck on this line in dist_util.py:

guided-diffusion/guided_diffusion/dist_util.py

Line 67 in 27c20a8

MPI.COMM_WORLD.bcast(data[i : i + chunk_size])

It is not clear to me why this fails.
Any pointers in fixing this problem would be greatly appreciated.
Thanks in advance.

[bahjat-kawar]

@VigneshSrinivasan10 Not sure if it's the same issue, but I also found the code to stop at the same line when running classifier_train.py.
I fixed the issue by indenting load_state_dict out of the if dist.get_rank() == 0 block on line 51.
Hope this helps.

[VigneshSrinivasan10]

@bahjat-kawar Thanks for tip and sorry for the delay in my response.
Your suggestion fixed the problem.

[VigneshSrinivasan10]

@bahjat-kawar Although the model reloading was successful, I still face loss values going to NAN after retraining for a few iterations. All the three .pt files were reloaded, but this issue still persists. I assumed the opt.pt file should have some information of the optimizer parameters which should help continue the training.

Did you also face this issue?

[JiamingLiu-Jeremy]

@VigneshSrinivasan10 I met the similar problem. Any progress on fixing this issue ?

[forever208]

@VigneshSrinivasan10 I met the similar problem. Any progress on fixing this issue ?

solution: remove if dist.get_rank() == 0 in script train_util.py when loading checkpints, because each GPU need to load checkpoint

[ONobody]

@forever208 Hello, do you use opt, model or ema's. pt file when using resume_checkpoint?
Or put their. pt files in a folder as the path of resume_checkpoint.

[forever208]

@on I use model to do resume training (where both ema and opt will be loaded). use ema to do sampling

[ONobody]

@forever208 When I continue to train,
such as python image_train.py --resume_checkpoint path/modelXX.pt? think you very much

[forever208]

@ONobody exactly
If you have further trouble, take a look at this pull request Fix resumed model training for Multi-GPUs

[ONobody]

@forever208 Thank you very much.