updated.

Author: tanyuqian

redco/deployers/deployer.py

@@ -14,6 +14,7 @@
 
 import os
 import jax
+from flax.training.common_utils import shard_prng_key
 import orbax.checkpoint as ocp
 
 from .data_utils import get_host_examples, get_data_batches
@@ -274,6 +275,14 @@ def gen_rng(self):
         self._rng, new_rng = jax.random.split(self._rng)
         return new_rng
 
+    def gen_model_step_rng(self):
+        rng = self.gen_rng()
+        if self.mesh is None:
+            rng = jax.random.split(
+                rng, num=jax.process_count())[jax.process_index()]
+            rng = shard_prng_key(rng)
+        return rng
+
     def log_info(self, info, title=None, step=None):
         """Logs a messages"""
         log_info(

redco/predictors/predictor.py

@@ -136,11 +136,7 @@ def predict(self,
                 self.setup_running_step(
                     dummy_batch=batch, params_shape_or_params=params)
 
-            rng = self._deployer.gen_rng()
-            if self.mesh is None:
-                rng = jax.random.split(
-                    rng, num=jax.process_count())[jax.process_index()]
-                rng = shard_prng_key(rng)
+            rng = self._deployer.gen_model_step_rng()
             batch_preds_with_idxes = self._deployer.run_model_step(
                 step_fn=self._p_pred_step, input_args=(rng, params, batch))
             batch_preds = process_batch_preds(

redco/trainers/trainer.py

@@ -22,7 +22,6 @@
 from jax.sharding import PartitionSpec as P
 from flax.jax_utils import replicate, unreplicate
 from flax.training import train_state
-from flax.training.common_utils import shard_prng_key
 from flax.core.frozen_dict import freeze
 from orbax.checkpoint.utils import \
     fully_replicated_host_local_array_to_global_array
@@ -227,11 +226,7 @@ def train(self, examples, per_device_batch_size, desc=None):
             if self._p_train_step is None:
                 self.setup_running_step(dummy_batch=batch)
 
-            rng = self._deployer.gen_rng()
-            if self.mesh is None:
-                rng = jax.random.split(
-                    rng, num=jax.process_count())[jax.process_index()]
-                rng = shard_prng_key(rng)
+            rng = self._deployer.gen_model_step_rng()
             self._state, metrics = self._deployer.run_model_step(
                 step_fn=self._p_train_step,
                 input_args=(rng, self._state, batch))
@@ -265,9 +260,9 @@ def eval_loss(self, examples, per_device_batch_size, desc=None):
             if self._p_eval_step is None:
                 self.setup_running_step(dummy_batch=batch)
 
+            rng = self._deployer.gen_model_step_rng()
             metrics = self._deployer.run_model_step(
-                step_fn=self._p_eval_step,
-                input_args=(jax.random.PRNGKey(0), self._state, batch))
+                step_fn=self._p_eval_step, input_args=(rng, self._state, batch))
             if self.mesh is None:
                 metrics = unreplicate(metrics)