Introduce Megatron-style parallel state management

[eternalNight]

Summary

This PR, authored by @hahaha3210 and @Daydreamer-Li, introduces ParallelState, a class that manages process groups for an arbitrary combination of parallel strategies including TP, EP, PP and DP.

As is discussed in #7680, the primary approach is borrowing the process group creation logic from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py but encapsulating the states (i.e., process groups, ranks and world sizes) into a class. This design enables the coexistence of multiple, independent parallelism configurations within a single process, which is particularly valuable in scenarios involving multiple models, such as in reinforcement learning (RL) workflows. Objects of ParallelState can be created prior to calls to deepspeed.initialize so that process groups are available to custom modules, such as UlyssesSPAttentionHF, at an early stage.

Compatibility of ParallelState and current process group management facilities (including deepspeed.runtime.sequence_parallel.parallel_state_sp and deepspeed.utils.groups) is tested by test_mpu.py.

Usage

Basic usage (single global instance passed to deepspeed.initialize):

    from deepspeed import DeepSpeedConfig
    ds_config = DeepSpeedConfig("config.json")

    parallel_state = ps.initialize_parallel_state_from_config(config_dict)
    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        model_parameters=model.parameters(),
        config=config_dict,
        mpu=parallel_state,
    )

The ParallelState class is compatible with Megatron's parallel_state module and thus can also be passed as mpu to existing implementations of parallelism techniques.

Opens

• Support for Ulysses SP is yet to be added. Dimensions for Ulysses SP can now be set by passing sequence_parallel_size to initialize_parallel_state_from_config.
• Support creating a ParallelState from a config object rather than specifying different parallel dimensions explicitly. initialize_parallel_state_from_config now accepts a config argument which can be either a dict or a config object.
• Are wrappers in parallel_state_deepspeed.py necessary? If so, is there a better way to implement more concisely its APIs sharing similar code patterns? parallel_state_deepspeed.py is preserved and renamed as parallel_state_wrappers.py to provide helper methods (create from config, wrap with context, etc.) around ParallelState for easier use.
• Are GLOO process groups necessary for DeepSpeed? If not, we can strip them from the draft. GLOO process groups are not created by default as deepspeed.comm.new_group does not accept a backend argument yet. The related code is kept so that we can easily bring it back when necessary.
• Tweaking NCCL options require ProcessGroupNCCL.options from torch.distributed, and that is not provided by deepspeed.comm today. Should we introduce that to deepspeed.comm, or make the format-checking script allowing that specific use of torch.distributed? Leave it as future work.

Future work

1. Configuration of process group options (only ProcessGroupNCCL.Options as of pytorch 2.9).
2. Add config items for parallelism dimensions not covered in json today.
3. Lazy creation of process groups. The Metatron-style manager creates all process groups for all potential combinations of parallelism in one shot. That won't impact much on GPU memory usage as recent NCCL allocates GPU memory on the first collective operation, not on communicator creation. But still it forks a significant amount of threads (~2000 on 8 GPUs) and may hit container PID limits. Possible optimizations include (1) not creating pgs for parallelism of size 1 at all, assuming they'll never be used; (2) creating pgs lazily (i.e., on first getter), assuming all ranks will call the same getter in the same order; (3) merge pgs among the same ranks, at the cost of synchronizing ops on them. That will be looked into in a follow-up PR.
4. While unit tests in this PR works properly, there're still code referring to specific mpu/parallel_state modules explicitly and thus not friendly to explicitly-passed mpu.

[sfc-gh-truwase]

@stas00 @tohtana @delock FYI

[delock]

I like the idea of puting parallel dimension in a single place rather than relying on user reading deeply into document to figure out how to turn on each parallelism dimension. I also agree with the open that the class can be created from a config object. Does it make sense to have the config in config.json file, or a seperate config file is more flexible?

[eternalNight]

I like the idea of puting parallel dimension in a single place rather than relying on user reading deeply into document to figure out how to turn on each parallelism dimension. I also agree with the open that the class can be created from a config object. Does it make sense to have the config in config.json file, or a seperate config file is more flexible?

My rough idea is to reuse the current config.json which already provides dimensions of various parallel techniques. Moving parallelism-related configs to a separate file is such a huge change that can impose a big obstable to users who try to upgrade.

[tohtana]

Hi @eternalNight, this is amazing!
Do you have a plan to add a document and unit tests? They will help users learn how to use this new style of parallelization configs.

[eternalNight]

Hi @eternalNight, this is amazing! Do you have a plan to add a document and unit tests? They will help users learn how to use this new style of parallelization configs.

Good point & will do.

[sfc-gh-truwase]

@eternalNight, I am curious how this is going? Thanks

[eternalNight]

@eternalNight, I am curious how this is going? Thanks

We're just back from the lunar new year holidays and are polishing the documentations. Will update another version this week probably.

[delock]

@eternalNight Can you sign-off to fix DCO error? Thanks!

[eternalNight]

@eternalNight Can you sign-off to fix DCO error? Thanks!

All commits are signed on this branch. The DCO errors are due to the mismatch between author info in the commit metadata (which are less official due to local configs) and the signed-off-by (which uses the official names and mail addresses).

[eternalNight]

@sfc-gh-truwase @delock @tohtana @stas00 Here's the final refreshment and rebase of the ParallelState class and helpers around it, with unit tests and documentation on usage. Please kindly review.

Future works are listed in the PR description.