Merge pull request #3805 from AI-Hypercomputer:chengnuojin-fix-2dfsdp

PiperOrigin-RevId: 910198099

Author: Google-ML-Automation

src/maxtext/configs/models/deepseek3-671b-2dfsdp.yml

@@ -71,9 +71,7 @@ logical_axis_rules: [
     ['activation_stage', 'stage'],
     ['embed', ['fsdp']],
     ['embed_moe', ['fsdp']],
-    ['embed_vocab', ['fsdp']],
-    ['embed_no_exp', ['fsdp']],
-    ['embed_no_exp_moe', ['fsdp']],
+    ['embed_vocab', ['fsdp', 'fsdp_transpose']],
     ['q_lora', ['fsdp']],
     ['kv_lora', ['fsdp']],
     ['layers', 'stage'],
@@ -83,7 +81,6 @@ logical_axis_rules: [
     ['kv_heads', ['fsdp_transpose', 'expert']],
     ['heads', ['fsdp_transpose', 'expert']],
     ['mlp', ['fsdp_transpose', 'expert']],
-    ['mlp_only_fsdp_transpose', ['fsdp_transpose']],
-    ['mlp_only_tensor', ['expert']],
+    ['mlp_moe', ['fsdp_transpose', 'expert']],
     ['diloco', 'diloco'],
 ]