Support moe_dropless_fallback=layer on the ragged all-to-all MoE path

src/maxtext/configs/base.yml

@@ -298,6 +298,7 @@ ragged_buffer_factor: -1.0 # a factor to determine the size of the ragged buffer
 eval_ragged_buffer_factor: -1.0 # evaluation override for ragged_buffer_factor.
 # When the ragged buffer would drop tokens: None = drop them; step = replay the whole train step with a
 # dropless buffer; layer = per MoE layer, lax.cond into a chunked dropless branch (no step replay).
+# step requires use_ring_of_experts=True; layer works with ring of experts and with the ragged all-to-all path.
 # Valid values: None, step, layer
 moe_dropless_fallback: None
 retry_dropless_first_steps: 0 # with moe_dropless_fallback=step, run steps with step - start_step < N directly with the dropless program (no normal attempt). 0 = off.

src/maxtext/configs/types.py

@@ -1122,7 +1122,8 @@ class MoEGeneral(BaseModel):
           "What to do when the ragged buffer (ragged_buffer_factor > 0) would drop tokens. None: drop them."
           " 'step': roll the step's update back in-graph and replay the whole train step with a dropless buffer."
           " 'layer': per MoE layer, calculate the overflow from the router's top-k and take a chunked dropless"
-          " branch (lax.cond) for that layer only; no step replay."
+          " branch (lax.cond) for that layer only; no step replay. 'step' requires use_ring_of_experts=True;"
+          " 'layer' also supports the ragged all-to-all path."
       ),
   )
   retry_dropless_first_steps: int = Field(
@@ -4281,7 +4282,7 @@ def validate_moe_dropless_fallback(self):
       raise ValueError(f"{prefix} requires num_experts > 1.")
     if self.ragged_buffer_factor <= 0:
       raise ValueError(f"{prefix} requires ragged_buffer_factor > 0.0.")
-    if not self.use_ring_of_experts:
+    if mode == "step" and not self.use_ring_of_experts:
       raise ValueError(f"{prefix} is currently only supported with use_ring_of_experts=True.")
     if not self.use_ragged_sort:
       raise ValueError(f"{prefix} requires use_ragged_sort=True.")

src/maxtext/layers/moe.py

@@ -2913,11 +2913,15 @@ def ra2a_and_route(
         rngs,
         input_ids=None,
         forced_routed_experts=None,
+        force_dropless=False,
         per_expert_scale=None,
     ):
       local_sorted_indices = None
       all_shards_group_sizes = None
       reshaped_group_sizes = None
+      # force_dropless (the in-layer dropless fallback) sizes the receive buffer for the worst case
+      # instead of ragged_buffer_factor * balanced, so no tokens can be dropped.
+      ragged_buffer_factor = -1.0 if force_dropless else self.get_ragged_buffer_factor()
       (
           x,
           sorted_selected_experts,
@@ -2959,13 +2963,13 @@ def ra2a_and_route(
               num_ep,
               self.config.num_experts,
               self.config.num_experts_per_tok,
-              self.get_ragged_buffer_factor(),
+              ragged_buffer_factor,
           )
           input_offsets, send_sizes, output_offsets, recv_sizes = RoutedMoE.get_all_to_all_params(
               all_shards_group_sizes,
               expert_shard_id,
               num_ep,
-              ragged_buffer_factor=self.get_ragged_buffer_factor(),
+              ragged_buffer_factor=ragged_buffer_factor,
               buffer_size=buffer_size,
           )
 
@@ -2988,7 +2992,7 @@ def ra2a_and_route(
               shard_index=expert_shard_id,
               use_custom_sort_vjp=self.config.use_custom_sort_vjp,
               use_ragged_sort=self.config.use_ragged_sort,
-              ragged_buffer_factor=self.get_ragged_buffer_factor(),
+              ragged_buffer_factor=ragged_buffer_factor,
               use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           )
         else:
@@ -3001,7 +3005,7 @@ def ra2a_and_route(
               global_sorted_experts=selected_experts,
               use_custom_sort_vjp=self.config.use_custom_sort_vjp,
               use_ragged_sort=self.config.use_ragged_sort,
-              ragged_buffer_factor=self.get_ragged_buffer_factor(),
+              ragged_buffer_factor=ragged_buffer_factor,
               use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           )
 
@@ -3061,6 +3065,7 @@ def route(
             rngs,
             input_ids=input_ids,
             forced_routed_experts=forced_routed_experts,
+            force_dropless=force_dropless,
             per_expert_scale=per_expert_scale,
         )
 
@@ -3862,7 +3867,9 @@ def _per_device_flag(has_overflow):
         return jnp.reshape(jnp.asarray(has_overflow, dtype=jnp.bool_), (1,))
 
       fallback_chunks = None
-      if self.config.moe_dropless_fallback == "layer":
+      # The all-to-all path only has a capped receive buffer (and so can only drop tokens) when the
+      # batch is sharded by expert; otherwise every EP shard already holds all tokens.
+      if self.config.moe_dropless_fallback == "layer" and (self.config.use_ring_of_experts or is_batch_sharded_by_expert):
         fallback_chunks = self.get_dropless_fallback_num_chunks(x.shape[1], n_chunks)
       if fallback_chunks is None:
         out, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio = _route_and_compute(force_dropless)
@@ -3888,10 +3895,58 @@ def _capped_branch():
         out, lb_loss, bias_updates, has_overflow, required_rbf, max_load_ratio = _route_and_compute(False)
         return out, lb_loss, bias_updates, jnp.asarray(has_overflow, dtype=jnp.bool_), required_rbf, max_load_ratio
 
+      def _a2a_dropless_chunks(n_chunks):
+        """Dropless ragged all-to-all path over n_chunks sequence chunks, run as a lax.scan.
+
+        Each chunk sizes its receive buffer for the worst case (force_dropless), which
+        get_dropless_fallback_num_chunks keeps no larger than the capped path's buffer. The scan body
+        is checkpointed, so the backward recomputes one chunk at a time and accumulates the expert
+        weight gradients in the scan carry instead of materializing one per chunk.
+        """
+        batch, seq_len = x.shape[0], x.shape[1]
+        chunk = seq_len // n_chunks
+
+        def split(a):  # [batch, seq, ...] -> [n_chunks, batch, chunk, ...]
+          if a is None:
+            return None
+          return jnp.moveaxis(a.reshape((batch, n_chunks, chunk) + a.shape[2:]), 1, 0)
+
+        def body(carry, xs):
+          x_c, logits_c, pre_bias_logits_c, input_ids_c, forced_c = xs
+          out_c, lb_c, bu_c, _, req_c, ratio_c, _ = _moe_body(
+              x_c,
+              logits_c,
+              pre_bias_logits_c,
+              w0,
+              w1,
+              wo,
+              w0_bias,
+              w1_bias,
+              wo_bias,
+              input_ids_c,
+              rngs,
+              forced_c,
+              force_dropless=True,
+              per_expert_scale=per_expert_scale,
+          )
+          return carry, (out_c, lb_c, bu_c, req_c, ratio_c)
+
+        xs = tuple(split(a) for a in (x, logits, pre_bias_logits, sharded_input_ids, forced_routed_experts))
+        body = jax.checkpoint(body, policy=jax.checkpoint_policies.nothing_saveable)
+        _, (outs, lb_losses, bias_updates, required_rbfs, max_load_ratios) = jax.lax.scan(body, None, xs)
+        output = jnp.moveaxis(outs, 0, 1).reshape((batch, seq_len) + outs.shape[3:])
+        lb_loss = None if lb_losses is None else jnp.mean(lb_losses, axis=0)
+        bias_updates = None if bias_updates is None else finalize_bias_updates(jnp.sum(bias_updates, axis=0))
+        required_rbf = None if required_rbfs is None else jnp.max(required_rbfs)
+        max_load_ratio = None if max_load_ratios is None else jnp.max(max_load_ratios)
+        return output, lb_loss, bias_updates, jnp.bool_(False), required_rbf, max_load_ratio
+
       # Save nothing from the rare dropless branch: lax.cond gives both branches the union of
       # their residuals, so the (larger) dropless residuals would otherwise be held every step.
       @functools.partial(jax.checkpoint, policy=jax.checkpoint_policies.nothing_saveable)
       def _dropless_branch():
+        if not self.config.use_ring_of_experts:
+          return _a2a_dropless_chunks(fallback_chunks)
         out, lb_loss, bias_updates, _, required_rbf, max_load_ratio = _route_and_compute(
             True, n_chunks=fallback_chunks, barrier_enabled=True
         )

tests/unit/configs_value_test.py

@@ -294,6 +294,13 @@ def test_moe_dropless_fallback_modes(self):
     with self.assertRaises(pydantic.ValidationError):
       types.MaxTextConfig(**common_config, moe_dropless_fallback="both")
 
+    # The ragged all-to-all path (no ring of experts) supports the in-layer fallback, not the step replay.
+    a2a_config = {**common_config, "use_ring_of_experts": False}
+    self.assertEqual(types.MaxTextConfig(**a2a_config, moe_dropless_fallback="layer").moe_dropless_fallback, "layer")
+    with self.assertRaises(pydantic.ValidationError) as context:
+      types.MaxTextConfig(**a2a_config, moe_dropless_fallback="step")
+    self.assertIn("only supported with use_ring_of_experts=True", str(context.exception))
+
   def test_tpu_tokamax_ring_config_validation_accepts_initial_config(self):
     argv = [
         "",

tests/unit/moe_test.py

@@ -2058,11 +2058,14 @@ def _build_retry_test_model(
       ragged_buffer_factor,
       moe_dropless_fallback: str | None = None,
       force_dropless: bool = False,
+      use_ring_of_experts: bool = True,
   ):
     """Builds a mixtral-8x7b RoutedMoE with the given ragged buffer/retry settings."""
     cfg = pyconfig.initialize(
         [None, get_test_config_path()],
-        run_name=(f"moe_retry_test_{ragged_buffer_factor}_{moe_dropless_fallback}_{force_dropless}"),
+        run_name=(
+            f"moe_retry_test_{ragged_buffer_factor}_{moe_dropless_fallback}_{force_dropless}_{use_ring_of_experts}"
+        ),
         enable_checkpointing=False,
         model_name="mixtral-8x7b",
         override_model_config=True,
@@ -2074,7 +2077,7 @@ def _build_retry_test_model(
         sparse_matmul=True,
         per_device_batch_size=4,
         ici_expert_parallelism=2,
-        use_ring_of_experts=True,
+        use_ring_of_experts=use_ring_of_experts,
         max_target_length=128,
         float32_gate_logits=True,
         use_ragged_sort=True,
@@ -2216,15 +2219,18 @@ def test_step_dropless_fallback_asymmetric_shard_overflow(self):
       )
     assert_moe_close(out_replay, out_dropless, cfg_dropless.dtype)
 
+  @parameterized.named_parameters(("ring_of_experts", True), ("all_to_all", False))
   @pytest.mark.tpu_only
-  def test_layer_dropless_fallback_matches_dropless(self):
+  def test_layer_dropless_fallback_matches_dropless(self, use_ring_of_experts):
     """A tiny buffer + in-layer fallback matches dropless output and gradients in one pass (no replay)."""
     mesh = self._build_retry_test_mesh()
     rng = jax.random.PRNGKey(2345)
     rng_model, rng_hidden_states = jax.random.split(rng)
     device_count = jax.device_count()
 
-    cfg_dropless, model_dropless = self._build_retry_test_model(mesh, ragged_buffer_factor=-1.0)
+    cfg_dropless, model_dropless = self._build_retry_test_model(
+        mesh, ragged_buffer_factor=-1.0, use_ring_of_experts=use_ring_of_experts
+    )
     hidden_states = jax.random.uniform(
         rng_hidden_states,
         (
@@ -2238,15 +2244,19 @@ def test_layer_dropless_fallback_matches_dropless(self):
       variables = model_dropless.init({"params": rng_model, "dropout": rng_model}, hidden_states)
       out_dropless, _, grads_dropless = self._out_loss_and_grad(model_dropless, variables["params"], hidden_states)
 
-    _, model_no_fallback = self._build_retry_test_model(mesh, ragged_buffer_factor=0.1)
+    _, model_no_fallback = self._build_retry_test_model(
+        mesh, ragged_buffer_factor=0.1, use_ring_of_experts=use_ring_of_experts
+    )
     with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_dropless.logical_axis_rules):
       out_no_fallback, _, _ = self._out_loss_and_grad(model_no_fallback, variables["params"], hidden_states)
     self.assertFalse(
         jnp.allclose(out_no_fallback.astype(jnp.float32), out_dropless.astype(jnp.float32), rtol=1e-2, atol=1e-2),
         msg="ragged_buffer_factor=0.1 unexpectedly matches dropless -- buffer isn't forcing an overflow.",
     )
 
-    _, model_fallback = self._build_retry_test_model(mesh, ragged_buffer_factor=0.1, moe_dropless_fallback="layer")
+    _, model_fallback = self._build_retry_test_model(
+        mesh, ragged_buffer_factor=0.1, moe_dropless_fallback="layer", use_ring_of_experts=use_ring_of_experts
+    )
     with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_dropless.logical_axis_rules):
       out_fallback, _, grads_fallback = self._out_loss_and_grad(model_fallback, variables["params"], hidden_states)
       _, mutated = model_fallback.apply({"params": variables["params"]}, hidden_states, mutable=["intermediates"])
@@ -2262,15 +2272,18 @@ def test_layer_dropless_fallback_matches_dropless(self):
     ):
       assert_moe_close(g_fallback, g_dropless, cfg_dropless.dtype)
 
+  @parameterized.named_parameters(("ring_of_experts", True), ("all_to_all", False))
   @pytest.mark.tpu_only
-  def test_layer_dropless_fallback_asymmetric_shard_overflow(self):
+  def test_layer_dropless_fallback_asymmetric_shard_overflow(self, use_ring_of_experts):
     """The predicate must fire (on every device) when only one EP shard overflows, and match dropless."""
     mesh = self._build_retry_test_mesh()
     rng = jax.random.PRNGKey(2345)
     rng_model, rng_hidden_states = jax.random.split(rng)
     device_count = jax.device_count()
 
-    cfg_dropless, model_dropless = self._build_retry_test_model(mesh, ragged_buffer_factor=-1.0)
+    cfg_dropless, model_dropless = self._build_retry_test_model(
+        mesh, ragged_buffer_factor=-1.0, use_ring_of_experts=use_ring_of_experts
+    )
     hidden_states = jax.random.uniform(
         rng_hidden_states,
         (
@@ -2293,7 +2306,9 @@ def test_layer_dropless_fallback_asymmetric_shard_overflow(self):
           {"params": variables["params"]}, hidden_states, forced_routed_experts=forced_routed_experts
       )
 
-    _, model_fallback = self._build_retry_test_model(mesh, ragged_buffer_factor=0.1, moe_dropless_fallback="layer")
+    _, model_fallback = self._build_retry_test_model(
+        mesh, ragged_buffer_factor=0.1, moe_dropless_fallback="layer", use_ring_of_experts=use_ring_of_experts
+    )
     with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_dropless.logical_axis_rules):
       (out_fallback, _, _), mutated = model_fallback.apply(
           {"params": variables["params"]},