[REQUEST] Switch to and unify DeviceMesh management for DP/TP/PP/EP/SP

[eternalNight]

Problem Statement

Today we have multiple sets of ProcessGroup management in the codebase for different parallel scenarios, namely:

• PipelineModule uses PipelineParallelGrid for pipeline parallelism.
• AutoTP uses one set of APIs (whose names start with letters) from utils.groups.
• MoE uses another set of APIs (whose names start with _) from utils.groups.
• Ulysses SP uses a DeviceMesh created on deepspeed.initialize.

All of those implementations deal with the process group creation (either using DeviceMesh or customized logic) under complex combinations. It significantly increases maintenance burden as well as difficulty in introducing new parallelism techniques.

Besides, that divergence also causes user confusions. deepspeed.initialize() accepts two parameters, i.e., mpu and mesh_param. mpu is for DP/TP (without PP/EP/SP), mesh_param for SP, and for PP/EP parallelism topology should be configured on model initialization. That does not provide a good experience for users who want to try out different configurations for best efficiency.

Proposal

We would like to refactor those ProcessGroup management facilities so that:

• A unified DeviceMesh based module serves all parallelism strategies and their combinations.
• Let DeviceMesh instances create and manage ProcessGroups when possible.
• Simplify the mesh topology configuration interface of deepseed.initialize.
• (Advanced) Allow different models in the same world to use different parallelism strategies. This is mainly for RL post-training, but need more careful feasibility investigation due to the global map design of ProcessGroup. See [RFC] Megatron-LM and MCore maintaining issues for veRL volcengine/verl#897 MCore pain point 1 for more information.
• (Advanced) Make it easy to extend ZeRO to support HSDP-style replicate + shared parallelism for asymmetric GPU clusters. The extension itself will be tracked in a separate issue.

While supporting multiple dimensions, DeviceMesh does not fit all parallelism techniques and thus need extension. Essentially, DeviceMesh, for each dimension, creates a ProcessGroup among ranks that share the same coordinate except that dimension. It fits the need of DP, TP and SP, but PP and EP have additional requirements: P2P groups among adjacent stages for PP and global data parallel groups for EP. We would like to subclass DeviceMesh in DeepSpeed (e.g. DeepSpeedDeviceMesh?) to collect parallelism configurations from both the model and configurations and create those additional groups. Each model has its own DeepSpeedDeviceMesh instance at DeepSpeedEngine.mesh_device and fetches ProcessGroups from there.

This is an early-stage idea yet. Any comment or suggestion will be welcomed.

[eternalNight]

@stas00 Following your suggestion, I just created this issue to focus on discussions on the DeviceMesh topic. Please review and feel free to comment.

[stas00]

Thank you for starting the discussion, Junjie. This is indeed something we need to solve.

I'm not quite sure DeviceMesh is the way to standardize on since:

1. you have no intuitive control over the rank placement - e.g. TP ranks have to be on the same node. I have asked about it the pytorch dev and their answer was that you have to create the mesh so that TP ranks are listed last, so it's sort of not explicit and or so easy to mess up.
2. Moreover we have situations where groups need to be created before deepspeed.initialize - like UlyssesSP - and DeviceMesh forces you to create all groups at once - it'll fail if you try init just one group not utilizing them all.

On the other hand Megatron-LM's latest incarnation has a full-blown 2K lines parallel-state module (mpu) which gives you all the flexibility you need. https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py

I think DeviceMesh is a super-nice abstraction for a lay user who needs a quick setup. But for such a complex thing as we have I'd investigate it more. Perhaps using separate components of DeviceMesh could be used, which don't force you to init all ranks at once, but then the neatness is lost. I'm not sure.

I think we need full control and things being very explicit to avoid mistakes that would lead to performance deterioration.