Merge pull request #4492 from AI-Hypercomputer:single_random_seed

PiperOrigin-RevId: 949273858

Author: Google-ML-Automation

src/maxtext/layers/moe.py

@@ -691,7 +691,7 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None):
       if rngs is None:
         raise ValueError("The random key cannot be None for random routing.")
       # Reuse the 'params' RNG stream to ensure random routing
-      rng = rngs.params()
+      rng = rngs.params() if hasattr(rngs, "params") and callable(getattr(rngs, "params")) else rngs
       top_k_weights, top_k_indices = random_routing(rng, gate_logits, self.num_experts_per_tok)
       return top_k_weights, top_k_indices
 
@@ -1949,16 +1949,9 @@ def moe_emb_chunking(
       expert_shard_id = jax.lax.axis_index(self._expert_parallelism_name) if num_ep > 1 else 0
 
       chunk_dim = embed_dim // self.config.num_moe_emb_chunks
-      chunk_rngs = (
-          jax.random.split(rngs.params(), self.config.num_moe_emb_chunks)
-          if self.config.use_random_routing
-          and rngs is not None
-          and hasattr(rngs, "params")
-          and callable(getattr(rngs, "params"))
-          else [rngs] * self.config.num_moe_emb_chunks
+      chunk_rngs_key = (
+          rngs.params() if rngs is not None and hasattr(rngs, "params") and callable(getattr(rngs, "params")) else None
       )
-      if rngs is not None:
-        chunk_rngs = jax.random.split(rngs.params(), self.config.num_moe_emb_chunks)
 
       first_x_unrouted = jax.lax.dynamic_slice_in_dim(x, 0, chunk_dim, axis=2)
       cur_x_chunk, routing, route_metadata = roe_ag_and_route(
@@ -1967,7 +1960,7 @@ def moe_emb_chunking(
           pre_bias_logits,
           num_ep,
           expert_shard_id,
-          nnx.Rngs(params=chunk_rngs[0]) if chunk_rngs is not None else None,
+          chunk_rngs_key,
           input_ids=sharded_input_ids,
       )
 
@@ -2001,7 +1994,7 @@ def scan_fn(carry, _):
             pre_bias_logits,
             num_ep,
             expert_shard_id,
-            nnx.Rngs(params=chunk_rngs[chunk_idx + 1]) if chunk_rngs is not None else None,
+            chunk_rngs_key,
             input_ids=sharded_input_ids,
         )
         return (next_x, next_ps0, next_ps1, chunk_idx + 1), None

tests/unit/moe_test.py

@@ -537,6 +537,90 @@ def test_dense(self):
     actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
     self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))
 
+  @pytest.mark.tpu_only
+  def test_moe_emb_chunking_random_routing(self):
+    cfg_chunked = pyconfig.initialize(
+        [None, get_test_config_path()],
+        run_name="moe_block_chunking_rr_test",
+        enable_checkpointing=False,
+        model_name="mixtral-8x7b",
+        dtype="bfloat16",
+        weight_dtype="bfloat16",
+        megablox=False,
+        sparse_matmul=True,
+        use_tokamax_gmm=True,
+        use_gmm_v2=True,
+        num_moe_emb_chunks=4,
+        use_ring_of_experts=True,
+        use_random_routing=True,
+        mlp_bias=True,
+        per_device_batch_size=1,
+        max_target_length=128,
+    )
+
+    cfg_non_chunked = pyconfig.initialize(
+        [None, get_test_config_path()],
+        run_name="moe_block_non_chunking_rr_test",
+        enable_checkpointing=False,
+        model_name="mixtral-8x7b",
+        dtype="bfloat16",
+        weight_dtype="bfloat16",
+        megablox=False,
+        sparse_matmul=True,
+        use_tokamax_gmm=True,
+        use_gmm_v2=True,
+        num_moe_emb_chunks=0,
+        use_ring_of_experts=True,
+        use_random_routing=True,
+        mlp_bias=True,
+        per_device_batch_size=1,
+        max_target_length=128,
+    )
+
+    rng = jax.random.PRNGKey(1234)
+    rng_model, rng_hidden_states = jax.random.split(rng)
+    device_count = jax.device_count()
+    hidden_states = jax.random.uniform(
+        rng_hidden_states,
+        (int(cfg_chunked.per_device_batch_size) * device_count, cfg_chunked.max_target_length, cfg_chunked.base_emb_dim),
+        dtype=cfg_chunked.dtype,
+    )
+
+    devices_array = maxtext_utils.create_device_mesh(cfg_chunked)
+    mesh = Mesh(devices_array, cfg_chunked.mesh_axes)
+
+    moe_chunked = moe.RoutedMoE(
+        config=cfg_chunked,
+        num_experts=cfg_chunked.num_experts,
+        num_experts_per_tok=cfg_chunked.num_experts_per_tok,
+        mesh=mesh,
+        kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+        kernel_axes=("embed", "mlp"),
+        dtype=cfg_chunked.dtype,
+        rngs=nnx.Rngs(params=rng_model),
+    )
+
+    moe_non_chunked = moe.RoutedMoE(
+        config=cfg_non_chunked,
+        num_experts=cfg_non_chunked.num_experts,
+        num_experts_per_tok=cfg_non_chunked.num_experts_per_tok,
+        mesh=mesh,
+        kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+        kernel_axes=("embed", "mlp"),
+        dtype=cfg_non_chunked.dtype,
+        rngs=nnx.Rngs(params=rng_model),
+    )
+
+    moe_non_chunked.gate.kernel.value = moe_chunked.gate.kernel.value
+    moe_non_chunked.wi_0.value = moe_chunked.wi_0.value
+    moe_non_chunked.wi_1.value = moe_chunked.wi_1.value
+    moe_non_chunked.wo.value = moe_chunked.wo.value
+
+    chunked_out, _, _ = moe_chunked(hidden_states)
+    non_chunked_out, _, _ = moe_non_chunked(hidden_states)
+
+    self.assertTrue(jax.numpy.allclose(chunked_out, non_chunked_out, rtol=1e-01, atol=1e-01, equal_nan=False))
+
   @pytest.mark.tpu_only
   def test_moe_emb_chunking_gmm_v2(self):
     cfg = pyconfig.initialize(