Moe dynamic redundancy

src/maxtext/configs/base.yml

@@ -225,6 +225,9 @@ load_balance_loss_weight: 0.0 # weight for the load balance loss
 use_random_routing: false # whether to use random routing for debug/test purpose
 use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul
 use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
+enable_moe_dynamic_redundant_experts: false # whether to use MoonEP-style dynamic redundant expert duplication and rebalancing
+moe_redundant_slots_per_rank: 1 # number of prefetch slots (B) per EP rank for dynamic expert duplication
+moe_rebalance_threshold_ratio: 1.15 # load threshold ratio (rank_load / avg_load) to trigger expert duplication
 num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension
 # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM
 # stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe).
@@ -269,8 +272,9 @@ merge_gating_gmm: false
 
 norm_topk_prob: false # boolean to enable the top-k probability normalization. qwen3-specific normalization of router weights.
 
-# when moe weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls
 moe_fsdp_use_two_stage_all_gather: false
+# Enable token activation deduplication in MoE ragged all-to-all communication
+enable_moe_token_activation_dedup: false
 # Shard the expert dimension of the MLP weights on the FSDP axis.
 # This configuration is recommended only when num_experts is a multiple of fsdp_parallelism
 shard_exp_on_fsdp: false
@@ -739,6 +743,7 @@ sft_train_on_completion_only: false
 # dataset_type must be synthetic, hf, grain, tfds
 # details in: https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline.md
 dataset_type: tfds
+synthetic_data_distribution: 'uniform'
 # for TFDS input pipeline (dataset_type=tfds)
 dataset_path: "" # your path given as argument in download_dataset.sh, e.g. "gs://my-maxtext-dataset/"
 dataset_name: 'c4/en:3.1.0'

src/maxtext/configs/types.py

@@ -860,6 +860,18 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to use Ring of Experts for sparse matmul expert parallelism.",
   )
+  enable_moe_dynamic_redundant_experts: bool = Field(
+      False,
+      description="Whether to use MoonEP-style dynamic redundant expert duplication and rebalancing.",
+  )
+  moe_redundant_slots_per_rank: int = Field(
+      2,
+      description="Number of prefetch slots (B) per EP rank for dynamic expert duplication.",
+  )
+  moe_rebalance_threshold_ratio: float = Field(
+      1.15,
+      description="Load threshold ratio (rank_load / avg_load) to trigger expert duplication.",
+  )
   moe_dispatch_no_expert_sharding: bool = Field(
       False,
       description=(
@@ -924,6 +936,10 @@ class MoEGeneral(BaseModel):
       False,
       description="Use `fsdp` and `fsdp_transpose` axes for 2D FSDP sharding.",
   )
+  enable_moe_token_activation_dedup: bool = Field(
+      False,
+      description="Enable token activation deduplication in MoE ragged all-to-all communication.",
+  )
   norm_topk_prob: bool = Field(
       False,
       description="Enable top-k probability normalization for router weights (Qwen3-specific).",
@@ -1336,6 +1352,7 @@ class DatasetGeneral(BaseModel):
   """General configuration for dataset and data loading."""
 
   dataset_type: DatasetType = Field(DatasetType.TFDS, description="The type of the data loading pipeline.")
+  synthetic_data_distribution: str = Field("uniform", description="Token distribution for synthetic dataset ('uniform' or 'zipf').")
   per_device_batch_size: int | float = Field(12, description="The batch size per device.")
   eval_per_device_batch_size: int | float = Field(
       0.0,

src/maxtext/input_pipeline/synthetic_data_processing.py

@@ -41,13 +41,25 @@ def __init__(self, config, mesh):
         SyntheticDataIterator.raw_generate_synthetic_data, out_shardings=data_pspec_shardings, static_argnums=0
     )
 
-    tokens = jax.random.randint(
-        jax.random.PRNGKey(0),
-        (config.global_batch_size_to_load, config.max_target_length + 1),
-        0,
-        config.vocab_size,
-        dtype=jnp.int32,
-    )
+    if getattr(config, "synthetic_data_distribution", "uniform") == "zipf":
+      alpha = 1.25
+      ranks = np.arange(1, config.vocab_size + 1, dtype=np.float64)
+      probs = 1.0 / (ranks ** alpha)
+      probs /= probs.sum()
+      flat_size = config.global_batch_size_to_load * (config.max_target_length + 1)
+      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)
+      tokens = jnp.array(np_tokens)
+    else:
+      tokens = jax.random.randint(
+          jax.random.PRNGKey(0),
+          (config.global_batch_size_to_load, config.max_target_length + 1),
+          0,
+          config.vocab_size,
+          dtype=jnp.int32,
+      )
 
     sequence_positions = jnp.arange(0, config.max_target_length + 1, dtype=jnp.int32).reshape(1, -1)
     batch_positions = jnp.broadcast_to(