add custom mesh and logical rule support

Author: NuojCheng

src/maxtext/configs/base.yml

@@ -352,13 +352,13 @@ moba_topk: 8
 
 # DeepSeek Sparse Attention (DSA)
 # deepseek3.2 introduces indexer in MLA
-use_indexer: False
-indexer_head_dim: 128
-indexer_n_heads: 64
-indexer_topk: 2048
-# Determines the training strategy for the indexer:
-# - False (Dense Warm-up): Computes indexer loss over all tokens. Used with `trainable_parameters_mask` to freeze other model parameters.
-# - True (Sparse Training): Computes indexer loss over top-k tokens only and detaches the indexer input for independent optimization.
+use_sparse_indexer: False
+index_head_dim: 128
+index_n_heads: 64
+index_topk: 2048
+# Determines the token selection strategy for indexer loss:
+# - False: Uses all tokens (Dense Warm-up).
+# - True: Uses only top-k tokens (Sparse Training).
 # Note: This is only active when `indexer_loss_scaling_factor` > 0.
 indexer_sparse_training: False
 # Multiplier for the indexer KL divergence loss
@@ -426,6 +426,7 @@ internal_compile_num_devices: -1 # You must specify the number of devices when u
 
 # Parallelism
 shard_mode: "auto" # can be either auto or explicit
+custom_mesh_and_rule: "" # replace default mesh and logical rule by specifying yml name under config/mesh_and_rule/.
 mesh_axes:
   [
     "diloco",

src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml

@@ -28,9 +28,7 @@ mesh_axes: ['data', 'stage', 'fsdp', 'tensor', 'expert']
 data_sharding: [['data', 'stage', 'fsdp', 'tensor', 'expert']]
 logical_axis_rules: [
                       ['activation_batch', ['data', 'fsdp', 'expert']],
-                      ['activation_batch_moe', ['data', 'fsdp', 'expert']],
                       ['activation_batch_no_exp', ['data', 'fsdp']],
-                      ['activation_batch_no_exp_moe', ['data', 'fsdp']],
                       ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']],
                       ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'expert']],
                       ['activation_heads', ['tensor']],
@@ -40,7 +38,6 @@ logical_axis_rules: [
                       ['activation_q_length', ['expert']],
                       ['activation_attn_embed', ['tensor']],
                       ['activation_embed', ['tensor']],
-                      ['activation_embed_moe', ['tensor']],
                       ['activation_mlp', ['tensor']],
                       ['activation_kv', ['tensor']],
                       ['activation_prefill_kv_batch', ['data', 'fsdp', 'expert']],
@@ -58,10 +55,7 @@ logical_axis_rules: [
                       ['q_heads', ['tensor']],
                       ['kv_heads', ['tensor']],
                       ['embed', ['fsdp', 'expert']],
-                      ['embed_moe', ['fsdp', 'expert']],
                       ['embed_no_exp', ['fsdp']],
-                      ['embed_no_exp_moe', ['fsdp']],
-                      ['embed_moe', ['fsdp']],
                       ['q_lora', ['fsdp']],
                       ['kv_lora', ['fsdp']],
                       ['norm', ['tensor']],

src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml

@@ -19,8 +19,6 @@ data_sharding: [['fsdp']]
 logical_axis_rules: [
                       ['activation_batch', ['fsdp']],
                       ['activation_batch_no_exp', ['fsdp']],
-                      ['activation_batch_moe', ['fsdp']],
-                      ['activation_batch_no_exp_moe', ['fsdp']],
                       ['activation_embed_and_logits_batch', ['fsdp']],
                       ['activation_embed_and_logits_batch_sequence', ['fsdp']],
                       ['activation_prefill_kv_batch', ['fsdp']],
@@ -29,8 +27,6 @@ logical_axis_rules: [
                       ['decode_batch', ['fsdp']],
                       ['embed', ['fsdp']],
                       ['embed_no_exp', ['fsdp']],
-                      ['embed_moe', ['fsdp']],
-                      ['embed_no_exp_moe', ['fsdp']],
                       ['q_lora', ['fsdp']],
                       ['kv_lora', ['fsdp']],
                       ['exp_with_fsdp', 'fsdp'],