How to continue training with DDP when loss is nan causing by amp

[ptrblck]

@xjm

In a multi-GPU DDP environment, if the loss on one rank is NaN while the others are normal, could this cause the all-reduce to hang?

No, the NaN loss will cause all gradient to be NaNs on this rank as well. The gradient synchronization will then create NaN gradients on all ranks and the scaler.step(optimizer) call will skip this update. The following scaler.update() step will decrease the scaling factor.