[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

[gekurian]

This comment isn't clear.

What does k | data_parallel_shard_degree mean?

[tianyu-l]

It stands for data_parallel_shard_degree % k == 0

[gekurian]

dp2ep here would be using the DP mesh to shard non-shared experts on the num_experts dimension? If so, could you make it clear in the comments?

[tianyu-l]

dp2ep would use "the entire cp mesh (if existing) + part of dp_shard mesh (namely dp_shard_2)" to shard non-shared experts.
Sorry for the confusion -- these PRs are not meant for landing without change. We'll definitely polish the descriptions later. Reading the parallel_dims.py might be more informative for now.