[d0t2]

Description

Support moe_dropless_fallback=layer on the ragged all-to-all MoE path (use_ring_of_experts=False). Before, the in-layer dropless fallback was ring-of-experts only and config validation rejected it otherwise.

This PR is required by #5785.

Why. With ragged_buffer_factor > 0, each EP shard's all-to-all receive buffer is capped at factor x balanced, and tokens beyond it are silently dropped. Without a cap (-1), the buffer is sized for the worst case: EP x local tokens. At EP16 that doesn't fit: 121 GiB of HBM for this recipe (available: 95.7 GiB).

How.

• Before the all-to-all, _check_ragged_overflow (the same check the ring path uses) computes from the router's top-k whether any EP shard would receive more than its buffer. lax.cond then picks a branch:
  • No overflow: the capped path, unchanged.
  • Overflow: _a2a_dropless_chunks runs the layer over get_dropless_fallback_num_chunks sequence chunks. Each chunk has a worst-case (force_dropless) receive buffer, which is no larger than the capped one, so peak HBM doesn't grow.
• The chunks run in a lax.scan with a jax.checkpoint(nothing_saveable) body. A Python-unrolled chunk loop like the ring path uses OOMed (116 GiB): the backward kept every chunk's f32 expert-weight gradient alive at once. With the scan, the backward recomputes one chunk at a time and accumulates the weight gradients in the carry (87.4 GiB).
• ra2a_and_route now honours force_dropless (worst-case buffer). route() passes it through.
• The fallback is only enabled when the batch is sharded by expert. Otherwise each EP shard already holds all tokens and nothing can be dropped.
• moe_dropless_fallback=step still requires use_ring_of_experts=True. Ring-of-experts code paths are unchanged.

Shortcomings.

• When a layer does overflow, its MoE runs about n_chunks× slower for that step (8× for EP16 / factor 2).
• log_required_ragged_buffer_factor is still ring-only, so the all-to-all path only reports whether the fallback ran (learning/moe_dropless_fallback_layers), not the required factor.

Tests

Unit tests

• test_layer_dropless_fallback_matches_dropless and test_layer_dropless_fallback_asymmetric_shard_overflow are now parameterized over ring_of_experts / all_to_all.
  • ragged_buffer_factor=0.1 forces overflow; the test checks that the capped model diverges from dropless.
  • With layer, outputs and gradients match the dropless model, and the fallback flag fires on every device even when only one EP shard overflows.
• test_force_dropless_matches_dropless still passes.
• test_moe_dropless_fallback_modes now also checks that layer is accepted and step rejected without ring-of-experts.

python3 -m pytest tests/unit/moe_test.py -k "layer_dropless_fallback or force_dropless_matches"   # 5 passed (TPU v4-8)
JAX_PLATFORMS=cpu python3 -m pytest tests/unit/configs_value_test.py -k moe_dropless_fallback \
  tests/unit/defer_small_all_reduces_test.py tests/unit/pyconfig_test.py                         # 11 passed

End-to-end

Llama4-Maverick on v5p-256 at FSDP8 x EP16 (full 48 layers, per_device_batch_size=4, max_target_length=8192, synthetic data, average of steps 16-19). Full recipe and results in #5785.

Config	s/step	TFLOP/s/dev	loss@19	XProf
ragged_buffer_factor=2.0, moe_dropless_fallback=layer	14.433	247.4	9.045	trace

The fallback never triggered in this run (synthetic data, no "overflow detected" log lines), so it measures the cost of the overflow check, which is negligible.

AOT HBM (train_compile, compile_topology=v5p-256, libtpu 0.0.46):

Config	HBM
FSDP8 x EP16, dropless (worst-case buffer)	121.06 GiB (OOM)
FSDP8 x EP16, ragged_buffer_factor=2.0	85.72 GiB
FSDP8 x EP16, ragged_buffer_factor=2.0, moe_dropless_fallback=layer (scan)	87.36 GiB
same, Python-unrolled chunk loop (not used)	116 GiB (OOM)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!