add support for v3-32

run_mlm_flax_stream.py

@@ -551,6 +551,10 @@ if __name__ == "__main__":
     # define number steps per stream epoch
     num_train_steps = data_args.num_train_steps
 
+    num_of_hosts = jax.process_count()
+    current_host_idx = jax.process_index()
+
+
     # 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
@@ -714,8 +718,13 @@ if __name__ == "__main__":
         # process input samples
         model_inputs = data_collator(samples, pad_to_multiple_of=16)
 
+        local_host_model_inputs = {
+                key: np.split(model_inputs.data[key], num_of_hosts, axis=0)[current_host_idx]
+                for key, value in model_inputs.data.items()
+            }
+
         # Model forward
-        model_inputs = shard(model_inputs.data)
+        model_inputs = shard(local_host_model_inputs)
         state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
 
         train_metrics.append(train_metric)