Merge pull request #4367 from Shuwen-Fang:deprecate_tp_transpose

PiperOrigin-RevId: 943625535

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -492,7 +492,7 @@ compile_xla_flags: "" # Compiler options e.g. compile_xla_flags="--xla_tpu_num_s
 # 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', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']
+mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']
 logical_axis_rules: [
                       ['circular_repeats', []],
                       # ==========================================
@@ -501,40 +501,36 @@ logical_axis_rules: [
                       # Vocab Activations
                       ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'expert']],
-                      ['activation_vocab', ['tensor', 'tensor_transpose', 'tensor_sequence']],
-                      ['activation_vocab', ['tensor', 'tensor_transpose']],
+                      ['activation_vocab', ['tensor', 'tensor_sequence']],
+                      ['activation_vocab', ['tensor']],
                       ['activation_vocab', 'tensor_sequence'],
                       # Vocab Weights
-                      ['vocab', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
+                      ['vocab', ['tensor', 'tensor_sequence', 'autoregressive']],
                       ['embed_vocab', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                       # ==========================================
                       # Attention
                       # ==========================================
                       # Attention Activations
                       ['activation_batch_attn', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
-                      ['activation_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
-                      ['activation_kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence']],
+                      ['activation_heads', ['tensor', 'tensor_sequence', 'autoregressive']],
+                      ['activation_kv_heads', ['tensor', 'tensor_sequence']],
                       ['activation_length_attn', ['context']],
                       ['activation_q_length', ['context']],
                       ['activation_kv_length', []],
-                      ['activation_embed_attn', ['tensor', 'tensor_transpose']],
-                      ['activation_kv', ['tensor', 'tensor_transpose', 'tensor_sequence']],
+                      ['activation_embed_attn', ['tensor']],
+                      ['activation_kv', ['tensor', 'tensor_sequence']],
                       ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
-                      ['activation_kv_head_dim', ['tensor', 'tensor_transpose', 'tensor_sequence']],
+                      ['activation_kv_head_dim', ['tensor', 'tensor_sequence']],
                       # Attention Weights
-                      ['heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
-                      ['q_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
-                      ['kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
+                      ['heads', ['tensor', 'tensor_sequence', 'autoregressive']],
+                      ['q_heads', ['tensor', 'tensor_sequence', 'autoregressive']],
+                      ['kv_heads', ['tensor', 'tensor_sequence', 'autoregressive']],
                       ['qkv', []],
                       ['kv', []],
                       ['kv_head_dim', []],
-                      ['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'tensor_transpose', 'expert']],
-                      ['q_lora', ['fsdp', 'context', 'tensor_transpose', 'expert']],
                       ['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                       ['q_lora', ['fsdp', 'context', 'expert']],
                       ["q_lora_up_proj", []],
-                      ['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'tensor_transpose', 'expert']],
-                      ['kv_lora', ['fsdp', 'context', 'tensor_transpose', 'expert']],
                       ['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                       ['kv_lora', ['fsdp', 'context', 'expert']],
                       ["kv_lora_up_proj", []],
@@ -545,37 +541,33 @@ logical_axis_rules: [
                       ['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_length_moe', ['context']],
                       ['activation_norm_length_moe', ['tensor_sequence', 'context']],
-                      ['activation_embed_moe', ['tensor', 'tensor_transpose']],
-                      ['activation_mlp_moe', ['tensor', 'tensor_transpose', 'tensor_sequence']],
+                      ['activation_embed_moe', ['tensor']],
+                      ['activation_mlp_moe', ['tensor', 'tensor_sequence']],
                       ['activation_exp', ['expert']],
                       # MoE Weights
                       ['exp', 'expert'],
                       ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
-                      ['embed_moe', ['fsdp', 'fsdp_transpose', 'tensor_transpose', 'context']],
-                      ['embed_moe', ['fsdp', 'tensor_transpose', 'context']],
                       ['embed_moe', ['fsdp', 'fsdp_transpose', 'context']],
                       ['embed_moe', ['fsdp', 'context']],
                       # ==========================================
                       # Standard MLP / Dense Layers / Model Structure
                       # ==========================================
                       # Dense Activations
-                      ['activation_mlp', ['tensor', 'tensor_transpose', 'tensor_sequence']],
+                      ['activation_mlp', ['tensor', 'tensor_sequence']],
                       # Note activation batch and length also get used in vocab
                       ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_length', ['context']],
                       ['activation_norm_length', ['tensor_sequence', 'context']],
-                      ['activation_embed', ['tensor', 'tensor_transpose']],
+                      ['activation_embed', ['tensor']],
                       ['activation_stage', 'stage'],
                       # General Weights
                       ['mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                       # GDN (linear-attention) projections shard like 'mlp' during training; the
                       # vLLM serving config overrides this to match tpu-inference's ATTN_HEAD order.
                       ['gdn_head', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
-                      ['embed', ['fsdp', 'fsdp_transpose', 'tensor_transpose', 'context', 'expert']],
-                      ['embed', ['fsdp', 'tensor_transpose', 'context', 'expert']],
                       ['embed', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                       ['embed', ['fsdp', 'context', 'expert']],
-                      ['norm', ['tensor', 'tensor_transpose']],
+                      ['norm', ['tensor']],
                       ['layers', 'stage'],
                       ['diloco', 'diloco'],
                       ['engram_dim', ['tensor']],
@@ -590,7 +582,6 @@ logical_axis_rules: [
                       ['activation_prefill_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['decode_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['decode_length', []],
-                      ['cache_heads', ['autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence']],
                       ['cache_heads', ['autoregressive', 'tensor', 'tensor_sequence']],
                       ['paged_kv_heads', ['tensor']],
                       ['cache_batch_prefill', []],
@@ -605,11 +596,10 @@ logical_axis_rules: [
                       # Deprecated / Scheduled for Removal
                       # ==========================================
                       ['mlp_no_fsdp', ['tensor', 'tensor_sequence', 'autoregressive']],
-                      ['embed_tensor_transpose', ['tensor_transpose']],
                       ['exp_with_fsdp', 'fsdp'],
                   ]
 # Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details
-data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']]
+data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']]
 input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length']
 # Determines which physical axis plays the role of context parallelism for input data processing and load balancing
 # only supports "context" or "expert" (when custom_mesh_and_rule=ep-as-cp)
@@ -630,7 +620,6 @@ dcn_sequence_parallelism: 1  # never recommended
 dcn_context_parallelism: 1
 dcn_context_autoregressive_parallelism: 1
 dcn_tensor_parallelism: 1 # never recommended
-dcn_tensor_transpose_parallelism: 1
 dcn_tensor_sequence_parallelism: 1 # never recommended
 dcn_pipeline_parallelism: 1
 dcn_expert_parallelism: 1
@@ -643,7 +632,6 @@ ici_sequence_parallelism: 1
 ici_context_parallelism: 1
 ici_context_autoregressive_parallelism: 1
 ici_tensor_parallelism: 1
-ici_tensor_transpose_parallelism: 1
 ici_tensor_sequence_parallelism: 1
 ici_autoregressive_parallelism: 1
 ici_pipeline_parallelism: 1

src/maxtext/configs/decoupled_base_test.yml

@@ -42,7 +42,6 @@ ici_pipeline_parallelism: 1
 ici_expert_parallelism: 1
 ici_sequence_parallelism: 1
 ici_context_parallelism: 1
-ici_tensor_transpose_parallelism: 1
 ici_tensor_sequence_parallelism: 1
 ici_autoregressive_parallelism: 1
 ici_fsdp_parallelism: -1
@@ -57,7 +56,6 @@ dcn_pipeline_parallelism: 1
 dcn_expert_parallelism: 1
 dcn_sequence_parallelism: 1
 dcn_context_parallelism: 1
-dcn_tensor_transpose_parallelism: 1
 dcn_tensor_sequence_parallelism: 1
 dcn_autoregressive_parallelism: 1
 dcn_fsdp_parallelism: 1

src/maxtext/configs/pyconfig_deprecated.py

@@ -918,7 +918,6 @@ def create_parallelisms_list(raw_keys):
       raw_keys["ici_context_parallelism"],
       raw_keys["ici_context_autoregressive_parallelism"],
       raw_keys["ici_tensor_parallelism"],
-      raw_keys["ici_tensor_transpose_parallelism"],
       raw_keys["ici_tensor_sequence_parallelism"],
       raw_keys["ici_expert_parallelism"],
       raw_keys["ici_autoregressive_parallelism"],
@@ -932,7 +931,6 @@ def create_parallelisms_list(raw_keys):
       raw_keys["dcn_context_parallelism"],
       raw_keys["dcn_context_autoregressive_parallelism"],
       raw_keys["dcn_tensor_parallelism"],
-      raw_keys["dcn_tensor_transpose_parallelism"],
       raw_keys["dcn_tensor_sequence_parallelism"],
       raw_keys["dcn_expert_parallelism"],
       raw_keys["dcn_autoregressive_parallelism"],
@@ -1029,7 +1027,6 @@ def pipeline_first_axis(raw_keys):
           raw_keys["ici_context_parallelism"],
           raw_keys["ici_context_autoregressive_parallelism"],
           raw_keys["ici_tensor_parallelism"],
-          raw_keys["ici_tensor_transpose_parallelism"],
           raw_keys["ici_tensor_sequence_parallelism"],
           raw_keys["ici_expert_parallelism"],
           raw_keys["ici_autoregressive_parallelism"],
@@ -1043,7 +1040,6 @@ def pipeline_first_axis(raw_keys):
           raw_keys["dcn_context_parallelism"],
           raw_keys["dcn_context_autoregressive_parallelism"],
           raw_keys["dcn_tensor_parallelism"],
-          raw_keys["dcn_tensor_transpose_parallelism"],
           raw_keys["dcn_tensor_sequence_parallelism"],
           raw_keys["dcn_expert_parallelism"],
           raw_keys["dcn_autoregressive_parallelism"],
@@ -1057,7 +1053,6 @@ def pipeline_first_axis(raw_keys):
           "context",
           "context_autoregressive",
           "tensor",
-          "tensor_transpose",
           "tensor_sequence",
           "expert",
           "autoregressive",
@@ -1072,7 +1067,6 @@ def pipeline_first_axis(raw_keys):
               "context",
               "context_autoregressive",
               "tensor",
-              "tensor_transpose",
               "tensor_sequence",
               "expert",
               "autoregressive",
@@ -1217,9 +1211,7 @@ def validate_shard_expert_on_fsdp(raw_keys):
   if raw_keys["shard_exp_on_fsdp"] and raw_keys["num_experts"] % raw_keys["ici_fsdp_parallelism"] != 0:
     raise ValueError("shard_exp_on_fsdp requires num_experts is divisiable by ici_fsdp_parallelism.")
   if raw_keys["shard_exp_on_fsdp"] and (using_tensor_parallelism(raw_keys) or using_expert_parallelism(raw_keys)):
-    raise ValueError(
-        "shard_exp_on_fsdp requires ici_expert_parallelism = 1 and ici_tensor_parallelism/ici_tensor_transpose_parallelism = 1."
-    )
+    raise ValueError("shard_exp_on_fsdp requires ici_expert_parallelism = 1 and ici_tensor_parallelism = 1.")
 
 
 def validate_ragged_dot(raw_keys):

src/maxtext/configs/types.py

@@ -963,7 +963,6 @@ class HardwareAndMesh(BaseModel):
           "context",
           "context_autoregressive",
           "tensor",
-          "tensor_transpose",
           "tensor_sequence",
           "expert",
           "autoregressive",
@@ -1036,7 +1035,6 @@ class DcnParallelism(BaseModel):
   dcn_context_parallelism: int = Field(1, description="DCN axis for context parallelism.")
   dcn_context_autoregressive_parallelism: int = Field(1, description="DCN axis for context autoregressive parallelism.")
   dcn_tensor_parallelism: int = Field(1, description="DCN axis for tensor parallelism (not recommended).")
-  dcn_tensor_transpose_parallelism: int = Field(1, description="DCN axis for tensor transpose parallelism.")
   dcn_tensor_sequence_parallelism: int = Field(
       1, description="DCN axis for tensor sequence parallelism (not recommended)."
   )
@@ -1056,7 +1054,6 @@ class IciParallelism(BaseModel):
   ici_context_parallelism: int = Field(1, description="ICI axis for context parallelism.")
   ici_context_autoregressive_parallelism: int = Field(1, description="ICI axis for context autoregressive parallelism.")
   ici_tensor_parallelism: int = Field(1, description="ICI axis for tensor parallelism.")
-  ici_tensor_transpose_parallelism: int = Field(1, description="ICI axis for tensor transpose parallelism.")
   ici_tensor_sequence_parallelism: int = Field(1, description="ICI axis for tensor sequence parallelism.")
   ici_autoregressive_parallelism: int = Field(1, description="ICI axis for autoregressive parallelism.")
   ici_pipeline_parallelism: int = Field(1, description="ICI axis for pipeline parallelism.")
@@ -3407,7 +3404,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         "context": self.ici_context_parallelism,
         "context_autoregressive": self.ici_context_autoregressive_parallelism,
         "tensor": self.ici_tensor_parallelism,
-        "tensor_transpose": self.ici_tensor_transpose_parallelism,
         "tensor_sequence": self.ici_tensor_sequence_parallelism,
         "model": self.ici_tensor_parallelism,
         "expert": self.ici_expert_parallelism,
@@ -3427,7 +3423,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         "context": self.dcn_context_parallelism,
         "context_autoregressive": self.dcn_context_autoregressive_parallelism,
         "tensor": self.dcn_tensor_parallelism,
-        "tensor_transpose": self.dcn_tensor_transpose_parallelism,
         "tensor_sequence": self.dcn_tensor_sequence_parallelism,
         "model": self.dcn_tensor_parallelism,
         "expert": self.dcn_expert_parallelism,