@jteuwen I see, thanks for the added context!
It indeed seems like workflow may be OOM-prone given that a 34GB model, corresponding gradients and optimizer states, and an input sample must fit into 40GB. Is my understanding of these memory sizes correct? There has been some work on DDP to reduce the memory overhead, perhaps @mrshenli may be able to shed some more light on that.
As an aside, torchelastic is a great way of recovering from errors and restarting training from a checkpoint. I’m curious whether it will result in the GPU tensors being freed (which could replace the failure recovery script shared above) cc @Kiuk_Chung