[QST] [DTensor] MoE + EP support

[jeromeku]

@wanchaol @tianyu-l

In the dev forum post on DTensor there was mention of upcoming support for MoE and EP, with a link to a design doc.

Wondering if there are any updates on this? Would it be possible to access the design doc, which is gated?

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k

[tianyu-l]

Hey I recently landed pytorch/torchtitan#1324 in torchtitan, which pretty much reflects what's covered in the doc. The doc itself is a bit outdated.

[jeromeku]

@tianyu-l

Thanks for the response!

Is it possible to implement MoE parallel folding as described here and implemented in Megatron-LM with the current state of DTensor and FSDP2?

More specifically, use different parallelisms for the attention and experts layers:

• Suppose we have 2 nodes, 8 gpus per node, so world_size = 16. Setting tensor parallel size: 8, expert_model_parallel_size: 4, expert_tensor_parallel_size: 2, we would have:

  • In the attention layers:

    • tp groups: [[0-7], [8-15]]
    • dp groups: [[0,8], [1,9],...]

  • In the expert layers:

    • ep model parallel groups: [[0, 2, 4, 6], [1, 3, 5, 7], ...]
    • ep tensor parallel groups: [[0, 1], [2, 3], ...]
    • dp groups: same as attention layer dp

This assumes tp-cp-ep-dp-pp ordering.

Additionally, one should also be able to add pipeline parallel to both parallelism groups and context parallel to the attention group.

[tianyu-l]

This is (partial) tp2ep implementation of EP where EP borrows degree from TP. What I have in pytorch/torchtitan#1324 is (partial) dp2ep where EP borrows degree from DP/CP. The implementation won't be that different.

I plan to implement a general solution to borrow from DP/CP and TP in flexible way (with tp-cp/ep-dp-pp ordering) which covers both cases and more, when I find time to.