deprecate expert_shard_attention_option flag

Author: NuojCheng

docs/reference/core_concepts/moe_configuration.md

@@ -96,11 +96,6 @@ Dropping:
 
 ## 2. Sharding
 
-`expert_shard_attention_option`: Determines how the "expert" axis is interpreted when sharding attention layers. Options include:
-
-- `fsdp`: Treats the expert axis as a FSDP axis.
-- `context`: Treats the expert axis as a context parallelism axis, useful for long context.
-
 `use_ring_of_experts` (experimental): This feature requires expert parallelism. If enabled, it replaces the standard two All-to-All communications with All-Gather in dispatch and Reduce-Scatter in collect. By gathering inputs across all shards, it allows for local routing and Top-K calculations, followed by result aggregation via Reduce-Scatter. This approach is particularly effective for models with a large Top-K, as it gathers activations before they are replicated k times to reduce communication.
 
 `moe_fsdp_use_two_stage_all_gather`: If enabled, split the All-Gather operation for MoE weights into two separate stages when using FSDP/FSDP-transpose sharding. This is preferred when 3D All-Gather support is unavailable.

src/maxtext/common/common_types.py

@@ -66,10 +66,6 @@
 MODEL_MODE_PREFILL = "prefill"
 MODEL_MODE_TRAIN = "train"
 
-# expert_shard_attention_option
-EP_AS_CONTEXT = "context"
-EP_AS_FSDP = "fsdp"
-
 DECODING_ACTIVE_SEQUENCE_INDICATOR = 1
 
 # A large negative mask value is used for masking to ensure that the

src/maxtext/configs/base.yml

@@ -237,11 +237,6 @@ merge_gating_gmm: False
 
 norm_topk_prob: false # boolean to enable the top-k probability normalization. qwen3-specific normalization of router weights.
 
-# how the expert axis is used to shard attention weights and activations
-# "fsdp" (ep acts as fsdp parallelism)
-# "context" (ep acts as context parallelism, training only)
-expert_shard_attention_option: "fsdp"
-
 # 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
 # Shard the expert dimension of the MLP weights on the FSDP axis.
@@ -461,9 +456,9 @@ logical_axis_rules: [
                       ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
                       ['activation_vocab', ['tensor', 'tensor_transpose', 'tensor_sequence']],
-                      # ['activation_vocab', ['tensor', 'tensor_transpose']],
-                      # ['activation_vocab', 'tensor_sequence'],
-                      # ['activation_vocab', ['sequence', 'context']],
+                      ['activation_vocab', ['tensor', 'tensor_transpose']],
+                      ['activation_vocab', 'tensor_sequence'],
+                      ['activation_vocab', ['sequence', 'context']],
                       # Vocab Weights
                       ['vocab', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                       ['embed_vocab', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
@@ -489,14 +484,14 @@ logical_axis_rules: [
                       ['kv', []],
                       ['kv_head_dim', []],
                       ['q_lora', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'tensor_transpose', 'expert']],
-                      # ['q_lora', ['fsdp', 'sequence', 'context', 'tensor_transpose', 'expert']],
-                      # ['q_lora', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
-                      # ['q_lora', ['fsdp', 'sequence', 'context', 'expert']],
+                      ['q_lora', ['fsdp', 'sequence', 'context', 'tensor_transpose', 'expert']],
+                      ['q_lora', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
+                      ['q_lora', ['fsdp', 'sequence', 'context', 'expert']],
                       ["q_lora_up_proj", []],
                       ['kv_lora', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'tensor_transpose', 'expert']],
-                      # ['kv_lora', ['fsdp', 'sequence', 'context', 'tensor_transpose', 'expert']],
-                      # ['kv_lora', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
-                      # ['kv_lora', ['fsdp', 'sequence', 'context', 'expert']],
+                      ['kv_lora', ['fsdp', 'sequence', 'context', 'tensor_transpose', 'expert']],
+                      ['kv_lora', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
+                      ['kv_lora', ['fsdp', 'sequence', 'context', 'expert']],
                       ["kv_lora_up_proj", []],
                       # ==========================================
                       # Mixture of Experts (MoE)
@@ -513,9 +508,9 @@ logical_axis_rules: [
                       ['exp', 'expert'],
                       ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                       ['embed_moe', ['fsdp', 'fsdp_transpose', 'sequence', 'tensor_transpose', 'context']],
-                      # ['embed_moe', ['fsdp', 'sequence', 'tensor_transpose', 'context']],
-                      # ['embed_moe', ['fsdp', 'fsdp_transpose', 'sequence', 'context']],
-                      # ['embed_moe', ['fsdp', 'sequence', 'context']],
+                      ['embed_moe', ['fsdp', 'sequence', 'tensor_transpose', 'context']],
+                      ['embed_moe', ['fsdp', 'fsdp_transpose', 'sequence', 'context']],
+                      ['embed_moe', ['fsdp', 'sequence', 'context']],
                       # ==========================================
                       # Standard MLP / Dense Layers / Model Structure
                       # ==========================================
@@ -530,9 +525,9 @@ logical_axis_rules: [
                       # General Weights
                       ['mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                       ['embed', ['fsdp', 'fsdp_transpose', 'sequence', 'tensor_transpose', 'context', 'expert']],
-                      # ['embed', ['fsdp', 'sequence', 'tensor_transpose', 'context', 'expert']],
-                      # ['embed', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
-                      # ['embed', ['fsdp', 'sequence', 'context', 'expert']],
+                      ['embed', ['fsdp', 'sequence', 'tensor_transpose', 'context', 'expert']],
+                      ['embed', ['fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
+                      ['embed', ['fsdp', 'sequence', 'context', 'expert']],
                       ['norm', ['tensor', 'tensor_transpose']],
                       ['layers', 'stage'],
                       ['diloco', 'diloco'],

src/maxtext/configs/types.py

@@ -661,10 +661,6 @@ class MoEGeneral(BaseModel):
   )
   use_random_routing: bool = Field(False, description="Whether to use random routing for debugging.")
   interleave_moe_layer_step: int = Field(1, description="Frequency of MoE layers, e.g., 2 means every 2nd layer is MoE.")
-  expert_shard_attention_option: Literal["fsdp", "context"] = Field(
-      "fsdp",
-      description="How the expert axis is used to shard attention weights and activations.",
-  )
   moe_fsdp_use_two_stage_all_gather: bool = Field(
       False,
       description="Use two separate All-Gather calls for MoE weights sharded on both FSDP and FSDP-transpose.",
@@ -2396,8 +2392,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       self.tensors_to_offload = [t for t in tensors if getattr(self, t) == "offload"]
 
     cp_size = self.ici_context_parallelism * self.dcn_context_parallelism
-    if self.expert_shard_attention_option == "context":
-      cp_size *= self.ici_expert_parallelism * self.dcn_expert_parallelism
     self.context_parallel_size = cp_size
     if self.pipeline_parallel_layers == -1:
       if self.decoder_block == DecoderBlockType.DEEPSEEK:

src/maxtext/layers/attention_mla.py

@@ -48,7 +48,6 @@
     D_KV,
     DType,
     EMBED,
-    EP_AS_CONTEXT,
     HEAD,
     Q_LORA_UP_PROJ,
     KV_BATCH,
@@ -901,9 +900,6 @@ def mla_get_key_value(self, low_rank_main, key_rope, model_mode):
     if model_mode == MODEL_MODE_PREFILL:
       key_logical_name = self.prefill_key_axis_names
       value_logical_name = self.prefill_value_axis_names
-    elif model_mode == MODEL_MODE_TRAIN and self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      key_logical_name = self.ep_key_axis_names
-      value_logical_name = self.ep_value_axis_names
     else:
       key_logical_name = self.key_axis_names
       value_logical_name = self.value_axis_names
@@ -1224,10 +1220,7 @@ def __call__(
       )
 
     out = jax.ad_checkpoint.checkpoint_name(out, "attention_out")
-    if model_mode == MODEL_MODE_TRAIN and self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      out = self._maybe_shard_with_logical(out, self.ep_out_axis_names)
-    else:
-      out = self._maybe_shard_with_logical(out, self.out_axis_names)
+    out = self._maybe_shard_with_logical(out, self.out_axis_names)
 
     out_sharding = create_sharding(self.mesh, out_logical_name)
     out = self.out_projection(out, out_sharding=out_sharding)

src/maxtext/layers/attention_op.py

@@ -55,7 +55,6 @@
     DEFAULT_MASK_VALUE,
     DType,
     D_KV,
-    EP_AS_FSDP,
     HEAD,
     KV_LENGTH,
     LENGTH,
@@ -1270,7 +1269,7 @@ def wrap_splash_kernel(single_head_mask):
 
       splash_kernel = wrap_splash_kernel(single_head_mask)
       segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH,))
-    elif self.config.use_jax_splash and self.config.expert_shard_attention_option == EP_AS_FSDP:
+    elif self.config.use_jax_splash:
       if self.config.use_max_logit_estimate > 0:
         sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate)
       segment_axis_names_splash_kernel = nn.logical_to_mesh_axes((Q_LENGTH,))

tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/logical_shardings.json

@@ -133,7 +133,7 @@
   },
   ".params/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -318,7 +318,7 @@
   ".params/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -459,7 +459,7 @@
   },
   ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -644,7 +644,7 @@
   ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -781,7 +781,7 @@
   },
   ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -966,7 +966,7 @@
   ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,

tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_4/logical_shardings.json

@@ -133,7 +133,7 @@
   },
   ".params/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -318,7 +318,7 @@
   ".params/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -459,7 +459,7 @@
   },
   ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -644,7 +644,7 @@
   ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -781,7 +781,7 @@
   },
   ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -966,7 +966,7 @@
   ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,

tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_1/logical_shardings.json

@@ -133,7 +133,7 @@
   },
   ".params/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -318,7 +318,7 @@
   ".params/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -459,7 +459,7 @@
   },
   ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -644,7 +644,7 @@
   ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -781,7 +781,7 @@
   },
   ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -966,7 +966,7 @@
   ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,

tests/utils/sharding_info/deepseek2-16b/v5p-16/slice_4/logical_shardings.json

@@ -133,7 +133,7 @@
   },
   ".params/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -318,7 +318,7 @@
   ".params/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -459,7 +459,7 @@
   },
   ".opt_state/[0]/.mu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -644,7 +644,7 @@
   ".opt_state/[0]/.mu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,
@@ -781,7 +781,7 @@
   },
   ".opt_state/[0]/.nu/['params']/['decoder']/['logits_dense']/['kernel']": {
     "partition_spec": [
-      "embed",
+      "embed_vocab",
       "vocab"
     ],
     "shape": [
@@ -966,7 +966,7 @@
   ".opt_state/[0]/.nu/['params']/['token_embedder']/['embedding']": {
     "partition_spec": [
       "vocab",
-      "embed"
+      "embed_vocab"
     ],
     "shape": [
       102400,