Moe dynamic redundancy

[snehalv2002]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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.

[gemini-code-assist]

Code Review

This pull request introduces MoonEP-style dynamic redundant expert duplication and token activation deduplication in MoE ragged all-to-all communication, along with Zipf distribution support for synthetic data generation. While these are valuable performance optimizations, several critical issues must be addressed. Specifically, the dynamic redundant expert duplication is currently functionally incomplete as the routing logic does not actually divert tokens to the duplicated slots. Additionally, the backward pass of the dynamic weight manager introduces a massive memory bottleneck by allocating a global sparse tensor and performing a psum over it, which will likely cause OOM errors on large models. Other issues include silent token dropping under routing imbalance in the deduplication path, redundant communication when weights are already gathered, potential double-duplication in the planning algorithm, and thread-safety concerns from modifying the global NumPy random state.

[gemini-code-assist]

The implementation of MoonEP-style dynamic redundant expert duplication is functionally incomplete because the router does not actually divert or reroute any tokens to the duplicated expert slots.

Specifically:

1. The reroute_fractions returned by plan_expert_duplication is ignored (assigned to _ at line 2473).
2. Neither slot_assignments nor reroute_fractions is passed to the routing/permute functions (route, moe_emb_chunking, permute).
3. In get_gmm_for_local_experts (line 2253), the redundant slots are padded with 0 in group_sizes (group_sizes = jnp.pad(group_sizes, (0, num_slots_B), constant_values=0)), meaning the GMM will never process any tokens for them.

Without modifying the routing logic to actually divert a fraction of tokens (based on reroute_fractions) from the overloaded experts to the redundant slots on the underloaded ranks, the prefetched weights in the redundant slots will remain completely unused, wasting memory and communication bandwidth.

[gemini-code-assist]

In _manage_dynamic_expert_weights_bwd, allocating a global sparse tensor of size [E, H, H_ffn] (where E = num_ep * local_e is the global number of experts) and performing a psum over it introduces a massive memory and communication bottleneck.

For large models like DeepSeek-V3, E is 256, and H * H_ffn is extremely large (e.g., 7168 * 2048 = 14.6M elements). Allocating g_global_sparse of size [256, 7168, 2048] requires ~7.5 GB of memory per layer in the backward pass, and performing jax.lax.psum over this giant, mostly-zero tensor across all EP ranks will be extremely slow and will likely cause Out-Of-Memory (OOM) errors.

Recommendation:
Instead of allocating a global sparse tensor and using psum, consider all-gathering the active slot gradients and indices (which are of size num_slots_B per rank, where num_slots_B is very small, e.g., 2) and then locally scattering/accumulating them, or using a more targeted communication pattern.

[gemini-code-assist]

In ra2a_dedup_and_route, expand_indices is truncated to total_local_tokens (defined as bsz_times_seq_len * self.num_experts_per_tok). Under routing imbalance, a single EP shard can receive more than total_local_tokens unique tokens. Truncating expand_indices to total_local_tokens will silently drop these excess tokens, leading to incorrect outputs.

Recommendation:
Use a safe upper bound (like buffer_size or a configurable capacity factor) instead of total_local_tokens to avoid silent token dropping.

[gemini-code-assist]

When the expert weights are already fully gathered on each device (i.e., w0.shape[0] == self.config.num_experts), slicing them to w0_home and then performing all_gather inside _manage_dynamic_expert_weights_fwd introduces redundant communication. Similarly, in the backward pass, performing psum is redundant because the weights are already replicated/gathered.

Recommendation:
Bypass the collective communication (all_gather and psum) when the weights are already fully gathered on each device.

[gemini-code-assist]

In plan_expert_duplication, if all experts of the overloaded rank r_max are already duplicated, r_max_expert_mask & (~dup_mask) will be all False, making candidate_expert_scores all -1. In this case, jnp.argmax will return 0 (or r_max_expert_start), which is an already duplicated expert. This will cause the already duplicated expert to be duplicated again to another slot, potentially overwriting its divert_fraction and wasting slots.

Recommendation:
Add a check to ensure can_balance is only True if there is at least one un-duplicated expert available on r_max.

Suggested change

	# Find hottest un-duplicated expert belonging to r_max
	r_max_expert_start = r_max * experts_per_rank
	expert_indices = jnp.arange(num_experts)
	r_max_expert_mask = (expert_indices >= r_max_expert_start) & (expert_indices < r_max_expert_start + experts_per_rank)
	candidate_expert_scores = jnp.where(
	r_max_expert_mask & (~dup_mask),
	total_tokens_per_expert,
	-1,
	)
	# Find hottest un-duplicated expert belonging to r_max
	r_max_expert_start = r_max * experts_per_rank
	expert_indices = jnp.arange(num_experts)
	r_max_expert_mask = (expert_indices >= r_max_expert_start) & (expert_indices < r_max_expert_start + experts_per_rank)
	has_candidate = jnp.any(r_max_expert_mask & (~dup_mask))
	can_balance = jnp.logical_and(can_balance, has_candidate)
	candidate_expert_scores = jnp.where(
	r_max_expert_mask & (~dup_mask),
	total_tokens_per_expert,
	-1,
	)

[gemini-code-assist]

Using np.random.seed(42) and np.random.choice modifies the global random state of NumPy, which is not thread-safe and can lead to non-reproducible behavior if other parts of the application or tests also use NumPy's global random generator.

Recommendation:
Use a local np.random.Generator or np.random.RandomState to ensure thread-safety and avoid side effects on the global random state.

Suggested change

	np.random.seed(42)
	np_tokens = np.random.choice(config.vocab_size, size=flat_size, p=probs).reshape(
	config.global_batch_size_to_load, config.max_target_length + 1
	).astype(np.int32)
	rng = np.random.default_rng(42)
	np_tokens = rng.choice(config.vocab_size, size=flat_size, p=probs).reshape(
	config.global_batch_size_to_load, config.max_target_length + 1
	).astype(np.int32)