Thanks, @iffiX
That definitely helps leverage the workers gradients - the model is able to find solutions A LOT faster now the model is actually able to explore different possibilities.
I tried it without the backward step at the PS and it causes predictions to become nan. However, if I keep the backward step in and actually (now) accumulate the gradients from the worker and the single batch in the PS, then the model is able to find solutions.
Once setting the gradient directly into the global model, technically I should be able to perform an optimizer step and update the global weights.
Any idea why the predictions result in nan when there isn’t a backward step in the PS?