fix running without ema

[feiloo]

afaict the model has to be unwrapped from accelerates distributed wrapping to not wait for the other gpus forever.
this happens when not using ema.

also since validation only runs on the main process and gpu, the other gpus timeout (nccl) waiting for it.
synchronizing within the validation function is a okay-ish liveliness-check and prevents timeouts or having to increase them to long durations.

i have done some rough and manual testing with updated dependencies (newer pytorch, accelerate, ...) on a single node with multiple gpus.

this should work with the original conda deps too, afaict

[flukeskywalker]

Thanks @feiloo!

I was aware of the nccl sync issue after porting the code to this simplified repo, but I decided that increasing the timeout or reducing val iters to not spend more too long on validation is a simpler solution for the purposes of this repo. IIRC the default timeout is 30 mins so if the val takes longer than that nccl will throw an error.

As you demonstrate, taking care of this complicates what should be a rather simple validation function. So I prefer not to do that. But perhaps it would be good to discuss this in the README to avoid surprises for users.

Thanks for catching the bug when EMA is off. I think this script wasn't fully tested in that setting since all our experiments used EMA.

[feiloo]

feel free to rewrite and adapt however you seem fit and close the pr.

the nccl default timeout is 10 minutes, other backends are 30.
it was/is a bit more costly to have a script fail late.
i think my fix is the more correct one, but i'll leave the choice up to you, ofc.

in a few days, i will adapt the pr to the em fix only for you, if you haven't by then.
Anyways, we appreciate your code.