Gradient synchronization in data-parallel trainers

[cgarciae]

Hey, great job with nanodl!

I was just looking through the code and noticed that when in Lambda's Trainer the gradients are not being averaged across devices here:

nanodl/nanodl/__src/models/lamda.py

Lines 564 to 565 in 18c7f8e

	loss, grads = jax.value_and_grad(loss_fn)(state.params)
	state = state.apply_gradients(grads=grads)

Not sure if this is happening elsewhere but usually to keep the weights in sync you apply a jax.lax.pmean over the gradients before passing them to apply_gradients, e.g.

grads = jax.lax.pmean(grads, axis_name='devices')

[HMUNACHI]

Thanks for noticing this! It's often challenging to test these portions due to the unavailability of a personal multi-GPU setup for development. However, I will be accessing 2 GPUs around 10th March. Will immediately examine this but you are more than welcome to make corrections from your end if convenient, I would in fact very much appreciate that.