update

Author: suexu1025

src/maxtext/configs/models/deepseek3-671b-batchsplit.yml

@@ -75,12 +75,14 @@ logical_axis_rules: [
     ['q_lora', ['fsdp']],
     ['kv_lora', ['fsdp']],
     ['layers', 'stage'],
-    ['q_lora_up_proj', ['fsdp_transpose', 'expert']],
-    ['kv_lora_up_proj', ['fsdp_transpose', 'expert']],
-    ['q_heads', ['fsdp_transpose', 'expert']],
-    ['kv_heads', ['fsdp_transpose', 'expert']],
-    ['heads', ['fsdp_transpose', 'expert']],
-    ['mlp', ['fsdp_transpose', 'expert']],
+    ['q_lora_up_proj', ['fsdp_transpose']],
+    ['kv_lora_up_proj', ['fsdp_transpose']],
+    ['q_heads', ['fsdp_transpose']],
+    ['kv_heads', ['fsdp_transpose']],
+    ['heads', ['fsdp_transpose']],
+    ['mlp', ['fsdp_transpose']],
     ['mlp_only_fsdp_transpose', ['fsdp_transpose']],
-    ['mlp_only_tensor', ['expert']],
+    ['expert_only', ['expert']],
+    ['fsdp_transpose_only', ['fsdp_transpose']],
+    ['fsdp_transpose_and_expert', ['fsdp_transpose', 'expert']],
 ]