Update run_mlm_flax.py

run_mlm_flax.py

@@ -751,38 +751,39 @@ def main():
     state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer)
 
     def train_step(state, batch, dropout_rng):
-            dropout_rng, new_dropout_rng = jax.random.split(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]
+        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
+            # 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
     
-            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
+            # 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)
     
-            # 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)
+        # 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)}
+        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
     
-            return new_state, metrics, new_dropout_rng
+        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,))