[MoE][PoC] Expert Parallel: dp2ep

[tianyu-l]

Stack from ghstack (oldest at bottom):

• -> [MoE][PoC] Expert Parallel: dp2ep #732
• [MoE][PoC] Expert Parallel: tp and tp2ep #731
• [MoE][PoC] model code #730

Temporary changes to unblock exploration

• [pytorch] comment out the check at https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/parallel/api.py#L66
• [torchtitan] turn optimizers foreach and clip_grad_norm_ off, as not all parameters are DTensors on the same meshes (e.g. (1) MoE non-shared experts and other params are on different FSDP meshes, and (2) moe.router.gate is a replicate torch.Tensor)

Also need to

• turn the block-level compile to full_graph=False because there will be an additional FSDP inside a TransformerBlock at the non shared experts level.

Things won't work

• For EP + TP, DCP resharding likely would fail due to the fact that experts would "forget" they are sharded because this meta info is not tracked as part of the 1-D DTensor. This can be solved by storing a 2-D DTensor (ep + tp), but requires several code changes including strided sharding from FSDP given a 2D DTensor.

Not including

• shared expert overlapping