Update run_mlm_flax.py

run_mlm_flax.py

@@ -750,41 +750,39 @@ def main():
     # Setup train state
     state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer)
 
-    # Define gradient update step fn
     def train_step(state, batch, dropout_rng):
-        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
-
-        def loss_fn(params):
-            labels = batch.pop("labels")
-
-            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
-
-            # compute loss, ignore padded input tokens
-            label_mask = jnp.where(labels!=-100, 1.0, 0.0)
-            loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
-
-            # take average
-            loss = loss.sum()
-            num_labels = label_mask.sum()
-
-            return loss, num_labels
-
-        grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
-        (loss, num_labels), grad = grad_fn(state.params)
-        num_labels = jax.lax.psum(num_labels, "batch")
-
-        # true loss = total loss / total samples
-        loss = jax.lax.psum(loss, "batch")
-        loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)
-
-        # true grad = total grad / total samples
-        grad = jax.lax.psum(grad, "batch")
-        grad = jax.tree_util.tree_map(lambda x: x / num_labels, grad)
-        new_state = state.apply_gradients(grads=grad)
-
-        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
-
-        return new_state, metrics, new_dropout_rng
+            dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+    
+            def loss_fn(params):
+                labels = batch.pop("labels")
+                logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+    
+                # compute loss, ignore padded input tokens
+                label_mask = jnp.where(labels!=-100, 1.0, 0.0)
+                loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
+                
+                 # take average on the per device loss
+                loss = loss.sum()
+                num_labels = label_mask.sum()
+               
+                return loss, num_labels
+    
+            grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+            (loss, num_labels), grad = grad_fn(state.params)
+            # Sum number of labels
+            num_labels = jax.lax.psum(num_labels, "batch")
+            
+            # Sum loss over devices, but only AFTER dividing by the number of labels
+            loss = jax.lax.psum(loss, "batch") / num_labels
+    
+            # true grad = total grad / total samples
+            grad = jax.lax.psum(grad, "batch")
+            grad = jax.tree_util.tree_map(lambda x: x / num_labels, grad)
+            new_state = state.apply_gradients(grads=grad)
+    
+            metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+    
+            return new_state, metrics, new_dropout_rng
 
     # Create parallel version of the train step
     p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))