Update 2dfsdp logical rule

[NuojCheng]

Description

Update 2dfsdp logical rule supporting FSDP+FSDP_T

Tests

Train compile test: https://paste.googleplex.com/4570213325602816

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[github-actions]

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[RissyRan]

Thanks for the fix!

[github-actions]

## 📋 Review Summary

The Pull Request correctly updates the logical axis rules for the DeepSeek3 2D FSDP configuration. It enables more efficient memory distribution for the large vocabulary embedding layer and aligns the MoE logical axes with the core model implementation.

🔍 General Feedback

• Optimization: Enabling 2D sharding (fsdp and fsdp_transpose) for embed_vocab is a solid improvement for handling the large embedding weights in a 2D mesh configuration.
• Consistency: Adding the mlp_moe rule ensures consistent sharding behavior across MoE components, matching the expected axis names in the model code.
• Cleanup: Removing unused axis names (embed_no_exp, mlp_only_tensor, etc.) reduces clutter and potential confusion in the configuration file.