DeepSpeed Training Deadlock

[eaplatanios]

Bug description

I'm using PyTorch Lightning combined with the DeepSpeed strategy (stage 2) to train on 8 V100 GPUs on a single node and I'm running into the following deadlock issue. If I get an OOM on one of the GPUs that does not correspond to the main process that started the job, then that "secondary process" will terminate but all other processes will end up waiting on it (i.e., 100% GPU utilization and nothing happening). This seems like a deadlock issue that lightning should be managing but I'm not sure where that is supposed to be handled and how. Aside from a potential resolution, pointers to how lightning is meant to handle this kind of deadlock and where that happens in the codebase would be super helpful. Thank you!

Environment

Current environment

#- Lightning Component: Trainer
#- PyTorch Lightning Version: 1.9.0
#- PyTorch Version: 1.13.0
#- Python version: 3.7
#- OS: Ubuntu 20
#- CUDA version: 11.6
#- GPU models and configuration: 8 V100s with 32GB each

cc @Borda @awaelchli @justusschock

[awaelchli]

@eaplatanios have you tried running with PL_RECONCILE_PROCESS=1? This isn't documented and not well known. I'm not sure if we have tested this with the deepspeed strategy.
Also relevant: #16410

[eaplatanios]

Thanks @awaelchli ! I did but I realized this isn't enabled in the DeepSpeed strategy because it overrides setup which prepares it. For now I have subclassed the DeepSpeed strategy in our codebase and patched it so that we do something similar which worked. It is not ideal because it does not currently allow the master process to wrap things up if it wants to (it just kills it if a secondary process initiates the reconciliation step) and so I didn't put up a PR in this repo, but I might do that if I get some time to think this through more. Making this a plugin like in your PR seems like the right way to go about this. :)

[awaelchli]

this isn't enabled in the DeepSpeed strategy because it overrides setup which prepares it.

You are right, I didn't see this but it's fully the reason as you said.

Making this a plugin like in your PR seems like the right way to go about this. :)

Yep. One reason why we removed the code in master is because the logic was a bit too much hardcoded to the ddp strategy. The issue you are facing with deepspeed is exactly a weak point of the original code, as it becomes very fragile through the inheritance. So this motivates the plugin through composition even more. cc @carmocca