Update run_mlm_flax.py

run_mlm_flax.py

@@ -687,8 +687,8 @@ def main():
     per_device_eval_batch_size = training_args.per_device_eval_batch_size
     eval_batch_size = per_device_eval_batch_size * local_device_count
 
-    num_train_steps = len(tokenized_datasets["train"]) // train_batch_size * num_epochs
-
+    num_train_steps = (len(tokenized_datasets["train"]) // (train_batch_size * jax.process_count())) * num_epochs
+    
     # Create learning rate schedule
     warmup_fn = optax.linear_schedule(
         init_value=0.0, end_value=training_args.learning_rate, transition_steps=training_args.warmup_steps
@@ -814,9 +814,8 @@ def main():
 
         # Generate an epoch by shuffling sampling indices from the train dataset
         num_train_samples = len(tokenized_datasets["train"])
-        # Avoid using jax.numpy here in case of TPU training
-        #train_samples_idx = np.random.permutation(np.arange(num_train_samples))
-        
+
+        train_samples_idx = np.arange(num_train_samples)
         train_samples_idx = np.random.permutation(train_samples_idx)
         # Split the training indices across processes        train_samples_idx = np.array_split(train_samples_idx, jax.process_count())[jax.process_index()]
         train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size, drop_last=True)