Initialize random seed for distributed models.

Keep track of a `global_random_seed`, and ensure it is set
when initializing `keras.distribution.initialize(...)`.

In multi-host processes in JAX, all processes require consistent
random number generation.  Otherwise, initializers on different
hosts would produce inconsistent values, resulting in both
compilation and runtime failures.

Author: cantonios

keras/src/backend/jax/distribution_lib.py

@@ -3,7 +3,10 @@
 import jax
 import numpy as np
 
+from keras.src.backend.common import global_state
+from keras.src.random import seed_generator
 from keras.src.utils import jax_utils
+from keras.src.utils import rng_utils
 
 
 def list_devices(device_type=None):
@@ -185,28 +188,82 @@ def distribute_data_input(per_process_batch, layout, batch_dim_name):
     return global_batch_array
 
 
-def initialize(job_addresses, num_processes, process_id):
-    if job_addresses and "," in job_addresses:
-        # When user provide all the job addresses, we will split and get the
-        # first one, which is the coordinator.
-        job_addresses = job_addresses.split(",")
-        # Do a sanity check to make sure the number of addresses also match
-        # the num_processes.
-        if num_processes is not None and num_processes != len(job_addresses):
-            raise ValueError(
-                f"The provided job_addresses {job_addresses} has "
-                f"{len(job_addresses)} jobs, but num_processes is "
-                f"{num_processes}"
-            )
-        coordinator_address = job_addresses[0]
-    else:
-        coordinator_address = job_addresses
+def initialize_rng():
+    """Initializes the global random number generator across processes.
 
-    jax.distributed.initialize(
-        coordinator_address=coordinator_address,
-        num_processes=num_processes,
-        process_id=process_id,
+    This is required for consistent initialization in multi-host settings.
+    """
+    global_seed = rng_utils.get_random_seed()
+    # Only set a random seed if not already set
+    # via keras.config.set_random_seed()
+    if global_seed is None:
+        # Generate a random seed on each CPU host and psum them to get a single
+        # consistent seed across all processes.
+        cpu_devices = jax.devices("cpu")
+        num_local_cpu_devices = jax.local_device_count("cpu")
+        # Seed must be in range [0, 2^32 - 1], so to ensure proper range and
+        # avoid signed integer overflow, we use uint32.
+        local_seed = jax.numpy.asarray(
+            [seed_generator.make_default_seed()] * num_local_cpu_devices,
+            dtype=jax.numpy.uint32,
+        )
+        # Sum across processes and pull out the first item.
+        global_seed = jax.pmap(
+            lambda x: jax.lax.psum(x, "all"),
+            axis_name="all",
+            devices=cpu_devices,
+        )(local_seed).item(0)
+        # Set the global seed.
+        rng_utils.set_random_seed(global_seed)
+
+    # Check if the global seed generator is set and ensure it has an initialized
+    # seed.  Otherwise, reset the seed to the global seed.
+    global_seed_generator = global_state.get_global_attribute(
+        "global_seed_generator"
     )
+    if global_seed_generator is not None:
+        seed = global_seed_generator.get_config()["seed"]
+        if seed is None:
+            global_state.set_global_attribute(
+                "global_seed_generator",
+                seed_generator.SeedGenerator(
+                    seed=global_seed,
+                    name=global_seed_generator.name,
+                    backend=global_seed_generator.backend,
+                ),
+            )
+
+
+def initialize(job_addresses, num_processes, process_id):
+    # Only call JAX initialize if not already initialized.
+    # Some JAX processes already set this up (e.g. multiprocess tests).
+    if not jax.distributed.is_initialized():
+        if job_addresses and "," in job_addresses:
+            # When user provide all the job addresses, we will split and get the
+            # first one, which is the coordinator.
+            job_addresses = job_addresses.split(",")
+            # Do a sanity check to make sure the number of addresses also match
+            # the num_processes.
+            if num_processes is not None and num_processes != len(
+                job_addresses
+            ):
+                raise ValueError(
+                    f"The provided job_addresses {job_addresses} has "
+                    f"{len(job_addresses)} jobs, but num_processes is "
+                    f"{num_processes}"
+                )
+            coordinator_address = job_addresses[0]
+        else:
+            coordinator_address = job_addresses
+
+        jax.distributed.initialize(
+            coordinator_address=coordinator_address,
+            num_processes=num_processes,
+            process_id=process_id,
+        )
+
+    # Ensure the random number generator is initialized across processes.
+    initialize_rng()
 
 
 def num_processes():

keras/src/utils/rng_utils.py

@@ -4,8 +4,11 @@
 
 from keras.src import backend
 from keras.src.api_export import keras_export
+from keras.src.backend.common import global_state
 from keras.src.utils.module_utils import tensorflow as tf
 
+GLOBAL_RANDOM_SEED = "global_random_seed"
+
 
 @keras_export("keras.utils.set_random_seed")
 def set_random_seed(seed):
@@ -46,6 +49,9 @@ def set_random_seed(seed):
             "Expected `seed` argument to be an integer. "
             f"Received: seed={seed} (of type {type(seed)})"
         )
+
+    # Store seed in global state so we can query it if set.
+    global_state.set_global_attribute(GLOBAL_RANDOM_SEED, seed)
     random.seed(seed)
     np.random.seed(seed)
     if tf.available:
@@ -54,3 +60,12 @@ def set_random_seed(seed):
         import torch
 
         torch.manual_seed(seed)
+
+
+def get_random_seed():
+    """Returns the explicit integer random seed if set.
+
+    If the seed has been explicitly set via `set_random_seed`, then
+    returns the seed.  Otherwise, returns `None`.
+    """
+    return global_state.get_global_attribute(GLOBAL_RANDOM_SEED)